"""Token budget management for LLM calls.

Tracks token usage and cost across the fix generation pipeline
to prevent runaway costs and enable intelligent model routing.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Dict

logger = logging.getLogger(__name__)

# Cost per 1K tokens (input, output) by model
MODEL_RATES: Dict[str, tuple] = {
    "claude-sonnet-4-20250514": (0.003, 0.015),
    "claude-opus-4-20250514": (0.015, 0.075),
    "claude-haiku-4-20250514": (0.00025, 0.00125),
    # Fallback for unknown models
    "default": (0.003, 0.015),
}


@dataclass
class TokenUsage:
    """Track token usage for a single LLM call."""

    model: str
    input_tokens: int
    output_tokens: int
    estimated_cost_usd: float
    duration_ms: float
    purpose: str  # e.g., "categorization", "fix_generation", "debate_round_1"


@dataclass
class TokenBudget:
    """Manage token budget across an entire fix attempt.

    Tracks cumulative usage and enforces cost limits to prevent
    runaway LLM spending on a single fix attempt.

    Usage:
        budget = TokenBudget(max_total_cost_usd=2.00)

        # Before each LLM call, check budget
        if budget.remaining_budget_usd > 0:
            result = await llm_call(...)
            budget.record_usage("categorization", model, in_tokens, out_tokens, duration)

        # Check if we should use expensive model
        if budget.should_use_expensive_model(complexity="high"):
            model = "claude-opus-4-20250514"
    """

    max_input_tokens: int = 100_000
    max_output_tokens: int = 4_000
    max_total_cost_usd: float = 2.00
    _usage_log: list = field(default_factory=list)
    _total_input_tokens: int = 0
    _total_output_tokens: int = 0
    _total_cost_usd: float = 0.0
    _start_time: float = field(default_factory=time.time)

    def estimate_cost(
        self, input_tokens: int, output_tokens: int, model: str
    ) -> float:
        """Estimate cost for a single LLM call."""
        in_rate, out_rate = MODEL_RATES.get(model, MODEL_RATES["default"])
        return (input_tokens / 1000 * in_rate) + (output_tokens / 1000 * out_rate)

    def record_usage(
        self,
        purpose: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        duration_ms: float = 0.0,
    ) -> TokenUsage:
        """Record token usage from an LLM call."""
        cost = self.estimate_cost(input_tokens, output_tokens, model)

        usage = TokenUsage(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            estimated_cost_usd=cost,
            duration_ms=duration_ms,
            purpose=purpose,
        )

        self._usage_log.append(usage)
        self._total_input_tokens += input_tokens
        self._total_output_tokens += output_tokens
        self._total_cost_usd += cost

        logger.debug(
            f"Token usage [{purpose}]: {input_tokens}in + {output_tokens}out "
            f"= ${cost:.4f} (total: ${self._total_cost_usd:.4f})"
        )

        return usage

    @property
    def remaining_budget_usd(self) -> float:
        """Remaining budget in USD."""
        return max(0.0, self.max_total_cost_usd - self._total_cost_usd)

    @property
    def total_cost_usd(self) -> float:
        """Total cost so far."""
        return self._total_cost_usd

    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return self._total_input_tokens + self._total_output_tokens

    @property
    def num_calls(self) -> int:
        """Number of LLM calls made."""
        return len(self._usage_log)

    def has_budget(self, estimated_cost: float = 0.0) -> bool:
        """Check if there's enough budget for another call."""
        if estimated_cost > 0:
            return self.remaining_budget_usd >= estimated_cost
        return self.remaining_budget_usd > 0

    def should_use_expensive_model(self, failure_complexity: str = "medium") -> bool:
        """Decide whether to use an expensive model based on complexity and remaining budget.

        Args:
            failure_complexity: "low", "medium", "high", or "critical"

        Returns:
            True if the expensive model (Opus) should be used
        """
        if failure_complexity in ("high", "critical"):
            # Use expensive model if we have enough budget
            # Opus is ~5x more expensive, so check we have 5x a Sonnet call left
            return self.remaining_budget_usd > 0.50
        return False

    def select_model(self, failure_complexity: str = "medium") -> str:
        """Select the best model given budget and complexity."""
        if self.should_use_expensive_model(failure_complexity):
            return "claude-opus-4-20250514"
        if self.remaining_budget_usd < 0.10:
            return "claude-haiku-4-20250514"  # Budget is low, use cheapest
        return "claude-sonnet-4-20250514"

    def get_summary(self) -> Dict:
        """Get a summary of token usage."""
        elapsed_ms = (time.time() - self._start_time) * 1000
        return {
            "total_calls": self.num_calls,
            "total_input_tokens": self._total_input_tokens,
            "total_output_tokens": self._total_output_tokens,
            "total_tokens": self.total_tokens,
            "total_cost_usd": round(self._total_cost_usd, 4),
            "remaining_budget_usd": round(self.remaining_budget_usd, 4),
            "elapsed_ms": round(elapsed_ms, 1),
            "calls": [
                {
                    "purpose": u.purpose,
                    "model": u.model,
                    "tokens": u.input_tokens + u.output_tokens,
                    "cost_usd": round(u.estimated_cost_usd, 4),
                }
                for u in self._usage_log
            ],
        }
