"""Google Gemini provider adapter for agentbudget."""

from typing import Any

from agentbudget.providers.base import BaseProvider, UsageData


class GeminiProvider(BaseProvider):
    """Provider adapter for Google Gemini API responses.

    Supports both the newer ``google-genai`` SDK (which exposes
    ``response.model_version``) and the older ``google-generativeai``
    SDK (which does not include model info in the response).

    When no model name can be found in the response, the model is set to
    ``"gemini"`` — pricing lookup will return $0.00 with a warning, and
    users should pass ``model=`` explicitly to ``session.track()``.
    """

    def extract_usage(self, response: Any) -> UsageData:
        """Extract token usage from a Gemini API response.

        Args:
            response: The raw API response object from the Gemini SDK.

        Returns:
            A UsageData instance with model, token counts, and provider name.

        Raises:
            ValueError: If the response has no usage_metadata.
        """
        usage_metadata = getattr(response, "usage_metadata", None)
        if usage_metadata is None:
            raise ValueError(
                "[agentbudget] Gemini response has no usage_metadata. "
                "Cannot extract token usage."
            )

        model: str = (
            getattr(response, "model", None)
            or getattr(response, "model_version", None)
            or "gemini"
        )

        input_tokens: int = getattr(usage_metadata, "prompt_token_count", 0) or 0
        output_tokens: int = getattr(usage_metadata, "candidates_token_count", 0) or 0

        return UsageData(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            provider="gemini",
        )
