"""Token counting using tiktoken.

Thin wrapper that handles encoding selection and caching.
"""

from __future__ import annotations

import tiktoken


class TokenCounter:
    """Counts tokens for a given model's tokenizer.

    Uses tiktoken with a fallback to cl100k_base (GPT-4/Claude approximation).
    Token counts are approximate across providers but good enough for budgeting.
    """

    def __init__(self, model: str = "gpt-4") -> None:
        try:
            self.encoding = tiktoken.encoding_for_model(model)
        except KeyError:
            # cl100k_base is a reasonable default for modern models
            self.encoding = tiktoken.get_encoding("cl100k_base")
        self.model = model

    def count(self, text: str) -> int:
        """Count tokens in a string."""
        if not text:
            return 0
        return len(self.encoding.encode(text))

    def count_messages(self, messages: list[dict[str, str]]) -> int:
        """Count tokens in a message array.

        Includes per-message overhead (~4 tokens for role/delimiters per the
        OpenAI tokenization spec). This is an approximation but errs on the
        side of overestimating, which is what we want for budgeting.
        """
        total = 0
        for msg in messages:
            total += 4  # role + delimiters overhead
            total += self.count(msg.get("content", ""))
        total += 2  # reply priming
        return total


# Module-level default for convenience
_default: TokenCounter | None = None


def get_counter(model: str = "gpt-4") -> TokenCounter:
    """Get or create a token counter. Reuses instance for same model."""
    global _default
    if _default is None or _default.model != model:
        _default = TokenCounter(model)
    return _default
