"""Budget tracking: token counting and cost estimation."""

import time
from dataclasses import dataclass
from typing import Any


# Pricing table ($ per 1M tokens)
# Source: https://openai.com/api/pricing/
MODEL_PRICING: dict[str, tuple[float, float]] = {
    # model_id: (input_per_1M, output_per_1M)
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "gpt-4.1": (2.00, 8.00),
    "gpt-4.1-mini": (0.40, 1.60),
    "gpt-4.1-nano": (0.10, 0.40),
    "o3": (2.00, 8.00),
    "o3-mini": (1.10, 4.40),
    "o4-mini": (1.10, 4.40),
}

_DEFAULT_PRICING = (2.50, 10.00)  # assume gpt-4o pricing for unknown models


@dataclass
class BudgetConfig:
    """Budget configuration for standalone mode.

    All limits are optional (0 = unlimited).
    In standalone mode budget check is post-hoc (after run completes).
    Pre-emptive enforcement is scope of durable mode.
    """

    max_tokens: int = 0
    max_cost_usd: float = 0.0
    max_tool_calls: int = 0
    max_duration_s: float = 0.0


class BudgetTracker:
    """Tracks token consumption and cost."""

    def __init__(self, model: str, config: BudgetConfig | None = None) -> None:
        self._model = model
        self._config = config or BudgetConfig()
        self._input_tokens = 0
        self._output_tokens = 0
        self._total_tokens = 0
        self._cost_usd = 0.0
        self._tool_calls = 0
        self._start_time = time.monotonic()

    def record_usage(self, usage: Any) -> None:
        """Record usage from OpenAI SDK (result.context_wrapper.usage).

        Usage object expected attributes: input_tokens, output_tokens, total_tokens.
        """
        if usage is None:
            return
        self._input_tokens += getattr(usage, "input_tokens", 0) or 0
        self._output_tokens += getattr(usage, "output_tokens", 0) or 0
        self._total_tokens += getattr(usage, "total_tokens", 0) or 0
        self._cost_usd = self._calculate_cost()

    def record_tool_call(self) -> None:
        """Record a tool invocation."""
        self._tool_calls += 1

    def check_budget(self) -> str | None:
        """Check limits. Returns reason for exceeded budget or None."""
        c = self._config
        if c.max_tokens > 0 and self._total_tokens > c.max_tokens:
            return f"token limit {c.max_tokens} exceeded ({self._total_tokens})"
        if c.max_cost_usd > 0 and self._cost_usd > c.max_cost_usd:
            return f"cost limit ${c.max_cost_usd:.2f} exceeded (${self._cost_usd:.4f})"
        if c.max_tool_calls > 0 and self._tool_calls > c.max_tool_calls:
            return f"tool call limit {c.max_tool_calls} exceeded ({self._tool_calls})"
        elapsed = time.monotonic() - self._start_time
        if c.max_duration_s > 0 and elapsed > c.max_duration_s:
            return f"duration limit {c.max_duration_s}s exceeded ({elapsed:.1f}s)"
        return None

    def _calculate_cost(self) -> float:
        input_price, output_price = MODEL_PRICING.get(
            self._model, _DEFAULT_PRICING
        )
        return (
            self._input_tokens * input_price / 1_000_000
            + self._output_tokens * output_price / 1_000_000
        )

    @property
    def summary(self) -> dict:
        elapsed = time.monotonic() - self._start_time
        return {
            "input_tokens": self._input_tokens,
            "output_tokens": self._output_tokens,
            "total_tokens": self._total_tokens,
            "cost_usd": self._cost_usd,
            "tool_calls": self._tool_calls,
            "duration_s": elapsed,
        }
