"""CostTracker — accumulate token usage and estimate costs per model."""

from __future__ import annotations

from dataclasses import dataclass, field

# Cost per 1M tokens (input, output) in USD — approximate as of early 2026
MODEL_PRICING: dict[str, tuple[float, float]] = {
    # OpenAI
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "gpt-4-turbo": (10.00, 30.00),
    "gpt-4": (30.00, 60.00),
    "gpt-3.5-turbo": (0.50, 1.50),
    "o1": (15.00, 60.00),
    "o1-mini": (3.00, 12.00),
    "o3-mini": (1.10, 4.40),
    # Anthropic
    "claude-opus-4-6": (15.00, 75.00),
    "claude-sonnet-4-6": (3.00, 15.00),
    "claude-haiku-4-5-20251001": (0.80, 4.00),
    "claude-3-5-sonnet-20241022": (3.00, 15.00),
    "claude-3-haiku-20240307": (0.25, 1.25),
}


@dataclass
class UsageRecord:
    """A single LLM call usage record."""

    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float = 0.0


@dataclass
class CostSummary:
    """Aggregated cost summary."""

    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_tokens: int = 0
    total_cost_usd: float = 0.0
    calls: int = 0
    by_model: dict[str, float] = field(default_factory=dict)


class CostTracker:
    """Accumulates token usage across LLM calls and estimates cost.

    Usage::

        tracker = CostTracker()
        tracker.record("gpt-4o", input_tokens=500, output_tokens=200)
        tracker.record("gpt-4o", input_tokens=300, output_tokens=100)

        summary = tracker.summary()
        print(f"Total cost: ${summary.total_cost_usd:.4f}")
        print(f"Total tokens: {summary.total_tokens}")
    """

    def __init__(self) -> None:
        self._records: list[UsageRecord] = []

    def record(
        self,
        model: str,
        *,
        input_tokens: int = 0,
        output_tokens: int = 0,
    ) -> UsageRecord:
        """Record a single LLM call's token usage.

        Returns the UsageRecord with estimated cost.
        """
        cost = self._estimate_cost(model, input_tokens, output_tokens)
        rec = UsageRecord(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost_usd=cost,
        )
        self._records.append(rec)
        return rec

    def summary(self) -> CostSummary:
        """Return aggregated cost summary across all recorded calls."""
        s = CostSummary()
        for r in self._records:
            s.total_input_tokens += r.input_tokens
            s.total_output_tokens += r.output_tokens
            s.total_cost_usd += r.cost_usd
            s.calls += 1
            s.by_model[r.model] = s.by_model.get(r.model, 0.0) + r.cost_usd

        s.total_tokens = s.total_input_tokens + s.total_output_tokens
        return s

    def reset(self) -> None:
        """Clear all recorded usage."""
        self._records.clear()

    @property
    def records(self) -> list[UsageRecord]:
        """Access all recorded usage entries."""
        return list(self._records)

    @staticmethod
    def _estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost in USD based on model pricing."""
        pricing = MODEL_PRICING.get(model)
        if pricing is None:
            return 0.0

        input_price_per_m, output_price_per_m = pricing
        return (input_tokens * input_price_per_m + output_tokens * output_price_per_m) / 1_000_000
