"""Per-run token usage accumulation, keyed by model name.

RAGAS 0.4.3's `ragas.metrics.collections` + instructor code path does not
expose real token counts (`ragas/cost.py` only serves the legacy langchain
`evaluate()` path). This module provides a context-scoped accumulator that
the HTTP response hook in `rag_eval.metrics.factory` feeds into, so token
counts survive across the AsyncOpenAI client caching used by InlineScorer
(see docs/superpowers/specs/2026-07-02-token-usage-tracking-design.md).
"""

from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Iterator


@dataclass
class TokenUsageTracker:
    """Accumulates input/output token counts and call counts, grouped by model name."""

    _totals: dict[str, dict[str, int]] = field(default_factory=dict)

    def record(self, model: str, input_tokens: int, output_tokens: int) -> None:
        """Add one API call's usage to the running total for `model`."""
        key = model or "unknown"
        bucket = self._totals.setdefault(
            key, {"input_tokens": 0, "output_tokens": 0, "calls": 0}
        )
        bucket["input_tokens"] += int(input_tokens)
        bucket["output_tokens"] += int(output_tokens)
        bucket["calls"] += 1

    def as_dict(self) -> dict[str, dict[str, int]]:
        """Return a plain-dict snapshot: {model: {input_tokens, output_tokens, calls}}."""
        return {model: dict(usage) for model, usage in self._totals.items()}

    def merge_into(self, existing: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]:
        """Return a new dict combining `existing` accumulated totals with this tracker's totals.

        Used by session-scoped scoring (one call at a time) to keep a running
        total across multiple calls instead of overwriting with just the latest call.
        Does not mutate `existing`.
        """
        merged: dict[str, dict[str, int]] = {
            model: dict(usage) for model, usage in existing.items()
        }
        for model, usage in self.as_dict().items():
            bucket = merged.setdefault(
                model, {"input_tokens": 0, "output_tokens": 0, "calls": 0}
            )
            bucket["input_tokens"] += usage["input_tokens"]
            bucket["output_tokens"] += usage["output_tokens"]
            bucket["calls"] += usage["calls"]
        return merged


_current_tracker: ContextVar[TokenUsageTracker | None] = ContextVar(
    "_current_tracker", default=None
)


@contextmanager
def track_token_usage() -> Iterator[TokenUsageTracker]:
    """Activate a fresh TokenUsageTracker for the duration of the `with` block.

    Any AsyncOpenAI client with `attach_usage_hook()` applied that makes a
    call while this context is active will have its usage recorded here.
    Safe to nest; the innermost tracker is active within its own block.
    """
    tracker = TokenUsageTracker()
    token = _current_tracker.set(tracker)
    try:
        yield tracker
    finally:
        _current_tracker.reset(token)


def get_current_tracker() -> TokenUsageTracker | None:
    """Return the currently active tracker, or None if no `track_token_usage()` block is active."""
    return _current_tracker.get()