from __future__ import annotations

import threading

from leshy.core.models import TokenUsage


class TokenTracker:
    """Thread-safe token usage tracker for LLM API calls.

    Used as a module-level singleton. LLM evaluators and buffs call `add()`
    after each API call; the runner reads `usage` after the suite completes.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._prompt_tokens = 0
        self._completion_tokens = 0

    def add(self, prompt_tokens: int, completion_tokens: int) -> None:
        with self._lock:
            self._prompt_tokens += prompt_tokens
            self._completion_tokens += completion_tokens

    @property
    def usage(self) -> TokenUsage:
        with self._lock:
            return TokenUsage(
                prompt_tokens=self._prompt_tokens,
                completion_tokens=self._completion_tokens,
                total_tokens=self._prompt_tokens + self._completion_tokens,
            )

    def reset(self) -> None:
        with self._lock:
            self._prompt_tokens = 0
            self._completion_tokens = 0


token_tracker = TokenTracker()
