#!/usr/bin/env python3
"""
SubagentStart Hook: Budget enforcement before spawning programs.

Receives subagent context via stdin (JSON with model, prompt, etc.)
Exits with code:
  0 - Allow spawn
  2 - Block spawn (budget exceeded)
"""

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

# Cost estimates per spawn (average expected cost in USD)
SPAWN_COST_ESTIMATES = {
    "opus": 0.15,
    "sonnet": 0.05,
    "haiku": 0.01,
    "claude-sonnet-4-20250514": 0.05,
    "claude-opus-4-20250514": 0.15,
    "claude-haiku-4-20250514": 0.01,
}

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"
    budget_path.parent.mkdir(parents=True, exist_ok=True)

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

def get_model_cost(model_name):
    """Get estimated cost for a model."""
    model_lower = model_name.lower() if model_name else "sonnet"

    for key, cost in SPAWN_COST_ESTIMATES.items():
        if key in model_lower:
            return cost

    # Default to sonnet cost if unknown
    return SPAWN_COST_ESTIMATES["sonnet"]

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}] BUDGET: {message}")
    except Exception:
        pass  # Don't fail if scratchpad unavailable

def main():
    # Read subagent 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:
        # No budget file, allow spawn
        print("No budget configuration found, allowing spawn", file=sys.stderr)
        sys.exit(0)

    # Check if budget limit is set
    budget_limit = budget.get("budget_limit")
    if budget_limit is None:
        # No limit set, allow spawn
        sys.exit(0)

    # Get current session costs
    current_session = budget.get("current_session", {})
    current_cost = current_session.get("estimated_cost", 0)

    # Get historical total
    history = budget.get("history", {})
    total_cost = history.get("total_cost", 0)

    # Estimate cost of this spawn
    model = context.get("model", "sonnet")
    spawn_cost = get_model_cost(model)

    # Calculate new totals
    new_session_cost = current_cost + spawn_cost
    new_total_cost = total_cost + spawn_cost

    # Check enforcement mode
    enforcement = budget.get("enforcement", "soft")
    warning_threshold = budget.get("warning_threshold", 0.75)
    confirmation_threshold = budget.get("confirmation_threshold", 0.90)

    # Calculate percentage of budget
    budget_percentage = new_session_cost / budget_limit if budget_limit > 0 else 0

    # Hard enforcement: block if over budget
    if enforcement == "hard" and new_session_cost > budget_limit:
        log_to_scratchpad(f"BLOCKED spawn - would exceed budget (${new_session_cost:.2f}/${budget_limit:.2f})")
        print(f"Budget exceeded: ${new_session_cost:.2f} would exceed limit of ${budget_limit:.2f}", file=sys.stderr)
        sys.exit(2)  # Exit 2 blocks the spawn

    # Warning at threshold
    if budget_percentage >= warning_threshold:
        log_to_scratchpad(f"WARNING: Budget at {budget_percentage*100:.0f}% (${new_session_cost:.2f}/${budget_limit:.2f})")

    # Track this spawn
    spawn_record = {
        "timestamp": datetime.now().isoformat(),
        "model": model,
        "estimated_cost": spawn_cost,
        "prompt_preview": context.get("prompt", "")[:100] if context.get("prompt") else ""
    }

    # Update session
    if "spawns" not in current_session:
        current_session["spawns"] = []
    current_session["spawns"].append(spawn_record)
    current_session["estimated_cost"] = new_session_cost

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

    log_to_scratchpad(f"Spawn allowed: {model} (${spawn_cost:.2f}, session total: ${new_session_cost:.2f})")

    # Allow spawn
    sys.exit(0)

if __name__ == "__main__":
    main()