"""Token and USD accounting with a hard cap. One tracker per run (and a global one per process)."""
from __future__ import annotations

import threading
from dataclasses import dataclass, field


class BudgetExceeded(RuntimeError):
    pass


@dataclass
class Usage:
    calls: int = 0
    cached_hits: int = 0
    input_tokens: int = 0
    cached_input_tokens: int = 0
    output_tokens: int = 0
    reasoning_tokens: int = 0
    usd: float = 0.0
    failures: int = 0
    invalid: int = 0
    latency_s: float = 0.0
    by_model: dict[str, dict[str, float]] = field(default_factory=dict)

    def as_dict(self) -> dict:
        d = {k: getattr(self, k) for k in ("calls", "cached_hits", "input_tokens", "cached_input_tokens",
                                             "output_tokens", "reasoning_tokens", "usd", "failures", "invalid", "latency_s")}
        d["by_model"] = self.by_model
        return d


class BudgetTracker:
    def __init__(self, usd_cap: float, token_cap: int | None, prices: dict[str, dict[str, float]]) -> None:
        self.usd_cap = usd_cap
        self.token_cap = token_cap
        self.prices = prices
        self.usage = Usage()
        self._lock = threading.Lock()

    def price(self, model: str, input_tokens: int, cached_tokens: int, output_tokens: int) -> float:
        p = self.prices.get(model)
        if p is None:
            # unknown model: use the most expensive known price so the cap is conservative
            p = max(self.prices.values(), key=lambda d: d["output"]) if self.prices else {"input": 5, "cached": 0.5, "output": 25}
        uncached = max(input_tokens - cached_tokens, 0)
        return (uncached * p["input"] + cached_tokens * p["cached"] + output_tokens * p["output"]) / 1e6

    def check(self) -> None:
        if self.usage.usd >= self.usd_cap:
            raise BudgetExceeded(f"USD cap reached: {self.usage.usd:.4f} >= {self.usd_cap}")
        if self.token_cap is not None and self.usage.input_tokens + self.usage.output_tokens >= self.token_cap:
            raise BudgetExceeded("token cap reached")

    def record(self, model: str, input_tokens: int, cached_tokens: int, output_tokens: int,
               reasoning_tokens: int = 0, latency_s: float = 0.0, cached_hit: bool = False) -> float:
        with self._lock:
            u = self.usage
            u.calls += 1
            if cached_hit:
                u.cached_hits += 1
                return 0.0
            cost = self.price(model, input_tokens, cached_tokens, output_tokens)
            u.input_tokens += input_tokens
            u.cached_input_tokens += cached_tokens
            u.output_tokens += output_tokens
            u.reasoning_tokens += reasoning_tokens
            u.usd += cost
            u.latency_s += latency_s
            m = u.by_model.setdefault(model, {"calls": 0, "input_tokens": 0, "output_tokens": 0, "usd": 0.0})
            m["calls"] += 1
            m["input_tokens"] += input_tokens
            m["output_tokens"] += output_tokens
            m["usd"] += cost
            return cost