#!/usr/bin/env python3
"""
SubagentStop Hook: Track subagent completion and update budget.

Receives completion context via stdin (JSON with model, tokens used, etc.)
Updates budget.json with actual usage data.
"""

import json
import sys
import os
from datetime import datetime
from pathlib import Path

def get_grid_dir():
    """Get the .grid directory path."""
    project_dir = os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd())
    return Path(project_dir) / ".grid"

def load_budget():
    """Load budget configuration from .grid/budget.json."""
    budget_path = get_grid_dir() / "budget.json"
    if not budget_path.exists():
        return None

    with open(budget_path, 'r') as f:
        return json.load(f)

def save_budget(budget):
    """Save budget configuration to .grid/budget.json."""
    budget_path = get_grid_dir() / "budget.json"

    with open(budget_path, 'w') as f:
        json.dump(budget, f, indent=2)

def log_to_scratchpad(message):
    """Append a log entry to SCRATCHPAD.md."""
    scratchpad_path = get_grid_dir() / "SCRATCHPAD.md"
    timestamp = datetime.now().strftime("%H:%M:%S")

    try:
        with open(scratchpad_path, 'a') as f:
            f.write(f"\n[{timestamp}] COMPLETE: {message}")
    except Exception:
        pass

def calculate_actual_cost(model, input_tokens, output_tokens, pricing):
    """Calculate actual cost based on token usage."""
    model_lower = model.lower() if model else "sonnet"

    # Determine pricing tier
    if "opus" in model_lower:
        tier = "opus"
    elif "haiku" in model_lower:
        tier = "haiku"
    else:
        tier = "sonnet"

    tier_pricing = pricing.get(tier, {"input": 3.00, "output": 15.00})

    # Pricing is per million tokens
    input_cost = (input_tokens / 1_000_000) * tier_pricing["input"]
    output_cost = (output_tokens / 1_000_000) * tier_pricing["output"]

    return input_cost + output_cost

def main():
    # Read completion context from stdin
    try:
        stdin_data = sys.stdin.read()
        if stdin_data.strip():
            context = json.loads(stdin_data)
        else:
            context = {}
    except json.JSONDecodeError:
        context = {}

    # Load budget
    budget = load_budget()
    if not budget:
        sys.exit(0)

    # Extract completion data
    model = context.get("model", "sonnet")
    input_tokens = context.get("input_tokens", 0)
    output_tokens = context.get("output_tokens", 0)
    exit_code = context.get("exit_code", 0)
    duration_ms = context.get("duration_ms", 0)

    # Calculate actual cost
    pricing = budget.get("pricing", {
        "opus": {"input": 5.00, "output": 25.00},
        "sonnet": {"input": 3.00, "output": 15.00},
        "haiku": {"input": 1.00, "output": 5.00}
    })

    actual_cost = calculate_actual_cost(model, input_tokens, output_tokens, pricing)

    # Update history totals
    history = budget.setdefault("history", {
        "total_cost": 0,
        "total_spawns": 0,
        "total_input_tokens": 0,
        "total_output_tokens": 0,
        "sessions": []
    })

    history["total_input_tokens"] = history.get("total_input_tokens", 0) + input_tokens
    history["total_output_tokens"] = history.get("total_output_tokens", 0) + output_tokens

    # Find and update the spawn record in current session
    current_session = budget.get("current_session", {})
    spawns = current_session.get("spawns", [])

    if spawns:
        # Update the most recent spawn with actual data
        last_spawn = spawns[-1]
        last_spawn["actual_cost"] = actual_cost
        last_spawn["input_tokens"] = input_tokens
        last_spawn["output_tokens"] = output_tokens
        last_spawn["exit_code"] = exit_code
        last_spawn["duration_ms"] = duration_ms
        last_spawn["completed_at"] = datetime.now().isoformat()

        # Recalculate session cost with actuals where available
        session_cost = sum(
            s.get("actual_cost", s.get("estimated_cost", 0))
            for s in spawns
        )
        current_session["estimated_cost"] = session_cost

    budget["current_session"] = current_session
    budget["history"] = history
    save_budget(budget)

    # Log completion
    status = "OK" if exit_code == 0 else f"FAILED (exit {exit_code})"
    duration_sec = duration_ms / 1000 if duration_ms else 0

    log_to_scratchpad(
        f"{model} {status} - ${actual_cost:.4f} "
        f"({input_tokens}in/{output_tokens}out, {duration_sec:.1f}s)"
    )

    sys.exit(0)

if __name__ == "__main__":
    main()