"""
Token usage tracking for evaluation runs.

TokenUsage captures the input, output, and cache token counts from a single
LLM call and supports aggregation via __add__ so callers can accumulate totals
across nodes and cases without manual bookkeeping.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class TokenUsage:
    """Token counts from a single LLM call or an aggregated group of calls.

    Attributes:
        input_tokens:           Prompt / context tokens billed as input.
        output_tokens:          Completion tokens generated by the model.
        cache_read_tokens:      Tokens read from a prompt cache (e.g. Anthropic
                                cache_read_input_tokens or OpenAI cached_tokens).
        cache_creation_tokens:  Tokens written to create a new cache entry.
    """

    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        """Sum of input and output tokens (cache tokens not double-counted)."""
        return self.input_tokens + self.output_tokens

    def __add__(self, other: TokenUsage) -> TokenUsage:
        """Combine two TokenUsage instances field-by-field."""
        return TokenUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            cache_read_tokens=self.cache_read_tokens + other.cache_read_tokens,
            cache_creation_tokens=self.cache_creation_tokens + other.cache_creation_tokens,
        )

    def __radd__(self, other: object) -> TokenUsage:
        """Support sum() starting from 0."""
        if other == 0:
            return self
        return NotImplemented  # type: ignore[return-value]

    def to_dict(self) -> dict[str, int]:
        """Serialise to a plain dict for JSON reporters."""
        return {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cache_read_tokens": self.cache_read_tokens,
            "cache_creation_tokens": self.cache_creation_tokens,
            "total_tokens": self.total_tokens,
        }
