"""Token counter using Google's GenAI SDK (Gemini models)."""

from __future__ import annotations

from context_manager.budget.base import TokenCounter


class GoogleCounter(TokenCounter):
    """Count tokens for Google Gemini / Gemma models via the GenAI API.

    Requires the ``google-genai`` package and a valid API key
    (passed explicitly or set as ``GOOGLE_API_KEY`` env var).
    """

    def __init__(self, model: str = "gemini-2.0-flash", api_key: str | None = None) -> None:
        try:
            from google import genai
        except ImportError as exc:
            raise ImportError(
                "google-genai is required for Google models. "
                "Install with: pip install llm-ctx-mgr[google]"
            ) from exc

        self._model = model
        self._client = genai.Client(api_key=api_key) if api_key else genai.Client()

    def count(self, text: str) -> int:
        """Return token count via the Google GenAI count_tokens API."""
        response = self._client.models.count_tokens(
            model=self._model,
            contents=text,
        )
        return response.total_tokens

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