"""Token counter using OpenAI's tiktoken library."""

from __future__ import annotations

from context_manager.budget.base import TokenCounter


class TiktokenCounter(TokenCounter):
    """Count tokens for OpenAI / tiktoken-compatible models."""

    def __init__(self, model: str = "gpt-4") -> None:
        try:
            import tiktoken
        except ImportError as exc:
            raise ImportError(
                "tiktoken is required for OpenAI models. "
                "Install with: pip install llm-ctx-mgr[openai]"
            ) from exc

        try:
            self._enc = tiktoken.encoding_for_model(model)
        except KeyError:
            self._enc = tiktoken.get_encoding("cl100k_base")
        self._model = model

    def count(self, text: str) -> int:
        """Return token count using tiktoken."""
        return len(self._enc.encode(text))

    def model_name(self) -> str:
        return self._model
