"""Token meter + budget ChatMiddleware fed from REAL UsageDetails (B4 / D6).

The hard token cap is enforced off the provider's own ``UsageDetails`` — **never** a
word-count proxy (research 03 Rec 3; the Fase 1 ``_word_tokens`` theatre is retired). A
``BudgetMiddleware`` reads ``response.usage_details["total_token_count"]`` after each chat
call and short-circuits with ``BudgetExceeded`` the moment the cap is crossed.

``strict_usage`` (default ``True``) makes a missing usage a HARD FAIL (``UsageUnavailable``):
a usage regression must never silently disable the cap. Test doubles that legitimately supply
a synthetic ``UsageDetails`` do not trip it.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass

from agent_framework import ChatContext, ChatMiddleware


class BudgetExceeded(RuntimeError):
    """Raised the moment a runtime cap (tokens or rounds) is crossed (B4).

    Carries the breached ``kind`` ("tokens" | "rounds"), the ``limit`` set, and the
    ``observed`` value that crossed it — a structured stop event, never a silent hang.
    """

    def __init__(self, kind: str, limit: int, observed: int) -> None:
        self.kind = kind
        self.limit = limit
        self.observed = observed
        super().__init__(f"budget exceeded: {kind} limit={limit} observed={observed}")


class UsageUnavailable(RuntimeError):
    """Raised when ``strict_usage`` is on and a response carries no usable usage — the cap
    must fail closed rather than silently stop counting (research 03 Rec 3)."""


@dataclass(frozen=True)
class Budget:
    """Hard token + round/iteration caps, required at startup (A4 / D6).

    Refuses to construct without positive caps — fail-fast, never an unbounded loop.
    """

    max_tokens: int
    max_rounds: int

    def __post_init__(self) -> None:
        if self.max_tokens <= 0:
            raise ValueError(f"max_tokens must be positive, got {self.max_tokens}")
        if self.max_rounds <= 0:
            raise ValueError(f"max_rounds must be positive, got {self.max_rounds}")


class TokenMeter:
    """Accumulates token and round usage against a ``Budget``; raises the moment a cap is
    crossed."""

    def __init__(self, budget: Budget) -> None:
        self.budget = budget
        self.tokens = 0
        self.rounds = 0

    def charge(self, tokens: int) -> int:
        """Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap."""
        self.tokens += tokens
        if self.tokens > self.budget.max_tokens:
            raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
        return self.tokens

    def tick_round(self) -> int:
        """Increment the round counter; raise ``BudgetExceeded`` if over cap."""
        self.rounds += 1
        if self.rounds > self.budget.max_rounds:
            raise BudgetExceeded("rounds", self.budget.max_rounds, self.rounds)
        return self.rounds


class BudgetMiddleware(ChatMiddleware):
    """Chat middleware that charges a ``TokenMeter`` from each response's real
    ``UsageDetails`` and short-circuits when the cap is crossed."""

    def __init__(self, meter: TokenMeter, *, strict_usage: bool = True) -> None:
        self._meter = meter
        self._strict = strict_usage

    async def process(
        self, context: ChatContext, call_next: Callable[[], Awaitable[None]]
    ) -> None:
        await call_next()
        usage = getattr(context.result, "usage_details", None)
        total = usage.get("total_token_count") if usage is not None else None
        if total is None:
            if self._strict:
                raise UsageUnavailable(
                    "usage_details missing total_token_count; the cap must fail closed "
                    "(set strict_usage=False only for non-counting paths)"
                )
            return
        self._meter.charge(int(total))  # raises BudgetExceeded if over cap