"""Thin wrapper around the Google Gemini API via the google-genai SDK."""

from __future__ import annotations

import base64
import mimetypes
from pathlib import Path
import logging
from typing import Optional, Sequence, Union
from urllib.request import urlopen

from google import genai
from google.genai import types

from .types import (
    EmbeddingVector,
    ImageInput,
    LLMResponse,
    MessageSequence,
    TokenUsage,
    normalize_messages,
)
from .utils import clamp_retries, run_sync_in_thread, run_with_retries

logger = logging.getLogger(__name__)


_GEMINI_MIN_TIMEOUT_S = 10.0
_GEMINI_MIN_TIMEOUT_MS = int(_GEMINI_MIN_TIMEOUT_S * 1000)


def _normalize_gemini_timeout_ms(timeout_s: float) -> int:
    """Convert a seconds timeout into the millisecond value expected by google-genai HttpOptions."""
    # google-genai HttpOptions expects milliseconds, but our public API uses seconds.
    effective_timeout_s = max(_GEMINI_MIN_TIMEOUT_S, timeout_s)
    if effective_timeout_s != timeout_s:
        logger.warning(
            "Gemini timeout %ss is too short, clamping to %ss.",
            timeout_s,
            effective_timeout_s,
        )
    timeout_ms = int(effective_timeout_s * 1000)
    return max(_GEMINI_MIN_TIMEOUT_MS, timeout_ms)


class GeminiClient:
    """Convenience wrapper around the Google Gemini SDK."""

    def generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        """Generate a response from the specified Gemini model.

        Args:
            api_key: API key used to authenticate with the Gemini API.
            prompt: Natural-language instruction or query for the model.
            model: Identifier of the Gemini model to target (for example, ``"gemini-2.5-flash"``).
            max_tokens: Cap for tokens across the entire exchange, defaults to 32000.
            reasoning_effort: Included for API parity; currently unused by the Gemini SDK.
            temperature: Optional sampling temperature.
            top_p: Optional nucleus sampling value.
            images: Optional collection of image references (local paths, URLs, or data URLs).
            messages: Optional list of chat-style messages (role/content).
            request_id: Optional request identifier for tracing/logging.
            timeout_s: Optional request timeout in seconds.
            max_retries: Optional retry count for transient failures.
            retry_backoff_s: Base delay (seconds) for exponential backoff between retries.

        Returns:
            The text output produced by the model.

        Raises:
            ValueError: If required arguments are missing or the request payload is empty.
            URLError: If an image URL cannot be retrieved.
            google.genai.errors.APIError: If the underlying Gemini request fails.
        """
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt and not messages and not images:
            raise ValueError("At least one of prompt, messages, or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
        system_parts: list[str] = []
        contents: list[types.Content] = []
        for message in normalized_messages:
            if message["role"] == "system":
                if message["content"]:
                    system_parts.append(message["content"])
                continue
            parts: list[types.Part] = []
            if message["content"]:
                parts.append(types.Part.from_text(text=message["content"]))
            contents.append(types.Content(role=message["role"], parts=parts))

        if images:
            image_parts = [self._to_image_part(image) for image in images]
            target_index = next(
                (
                    index
                    for index in range(len(contents) - 1, -1, -1)
                    if contents[index].role == "user"
                ),
                None,
            )
            if target_index is None:
                contents.append(types.Content(role="user", parts=image_parts))
            else:
                existing_parts = list(contents[target_index].parts or [])
                existing_parts.extend(image_parts)
                contents[target_index] = types.Content(
                    role="user", parts=existing_parts
                )

        if not contents or not any(content.parts for content in contents):
            raise ValueError("No content provided for response generation.")

        config_kwargs: dict[str, object] = {"max_output_tokens": max_tokens}
        if system_parts:
            config_kwargs["system_instruction"] = "\n".join(system_parts)
        if reasoning_effort is not None:
            config_kwargs["thinking_config"] = types.ThinkingConfig(thinking_level=reasoning_effort)
        if temperature is not None:
            config_kwargs["temperature"] = float(temperature)
        if top_p is not None:
            config_kwargs["top_p"] = float(top_p)
        config = types.GenerateContentConfig(**config_kwargs)

        retry_count = clamp_retries(max_retries)

        def _build_client() -> genai.Client:
            client_kwargs: dict[str, object] = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["http_options"] = types.HttpOptions(
                    timeout=_normalize_gemini_timeout_ms(timeout_s)
                )
            return genai.Client(**client_kwargs)

        def _run_request() -> str:
            client = _build_client()
            try:
                try:
                    response = client.models.generate_content(
                        model=model,
                        contents=contents,
                        config=config,
                    )
                except Exception as exc:
                    logger.exception(
                        "Gemini generate_content failed: %s request_id=%s",
                        exc,
                        request_id,
                    )
                    raise
            finally:
                closer = getattr(client, "close", None)
                if callable(closer):
                    try:
                        closer()
                    except Exception:
                        pass

            if response.text:
                result_text = response.text
                logger.info(
                    "Gemini generate_content succeeded: model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(result_text or ""),
                    request_id,
                )
                return result_text

            candidate_texts: list[str] = []
            for candidate in getattr(response, "candidates", []) or []:
                content_obj = getattr(candidate, "content", None)
                if not content_obj:
                    continue
                for part in getattr(content_obj, "parts", []) or []:
                    text = getattr(part, "text", None)
                    if text:
                        candidate_texts.append(text)

            if candidate_texts:
                result_text = "\n".join(candidate_texts)
                logger.info(
                    "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(result_text or ""),
                    request_id,
                )
                return result_text

            # Treat successful calls without textual content as a successful, empty response
            # rather than raising. This aligns with callers that handle empty outputs gracefully.
            logger.info(
                "Gemini generate_content succeeded with no text: model=%s images=%d request_id=%s",
                model,
                len(images or []),
                request_id,
            )
            return ""

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    def generate_response_with_usage(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt and not messages and not images:
            raise ValueError("At least one of prompt, messages, or images must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        normalized_messages = normalize_messages(prompt=prompt, messages=messages)
        system_parts: list[str] = []
        contents: list[types.Content] = []
        for message in normalized_messages:
            if message["role"] == "system":
                if message["content"]:
                    system_parts.append(message["content"])
                continue
            parts: list[types.Part] = []
            if message["content"]:
                parts.append(types.Part.from_text(text=message["content"]))
            contents.append(types.Content(role=message["role"], parts=parts))

        if images:
            image_parts = [self._to_image_part(image) for image in images]
            target_index = next(
                (
                    index
                    for index in range(len(contents) - 1, -1, -1)
                    if contents[index].role == "user"
                ),
                None,
            )
            if target_index is None:
                contents.append(types.Content(role="user", parts=image_parts))
            else:
                existing_parts = list(contents[target_index].parts or [])
                existing_parts.extend(image_parts)
                contents[target_index] = types.Content(role="user", parts=existing_parts)

        if not contents or not any(content.parts for content in contents):
            raise ValueError("No content provided for response generation.")

        config_kwargs: dict[str, object] = {"max_output_tokens": max_tokens}
        if system_parts:
            config_kwargs["system_instruction"] = "\n".join(system_parts)
        if reasoning_effort is not None:
            config_kwargs["thinking_config"] = types.ThinkingConfig(thinking_level=reasoning_effort)
        if temperature is not None:
            config_kwargs["temperature"] = float(temperature)
        if top_p is not None:
            config_kwargs["top_p"] = float(top_p)
        config = types.GenerateContentConfig(**config_kwargs)

        retry_count = clamp_retries(max_retries)

        def _build_client() -> genai.Client:
            client_kwargs: dict[str, object] = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["http_options"] = types.HttpOptions(
                    timeout=_normalize_gemini_timeout_ms(timeout_s)
                )
            return genai.Client(**client_kwargs)

        def _run_request() -> LLMResponse:
            client = _build_client()
            try:
                try:
                    response = client.models.generate_content(
                        model=model,
                        contents=contents,
                        config=config,
                    )
                except Exception as exc:
                    logger.exception(
                        "Gemini generate_content failed: %s request_id=%s",
                        exc,
                        request_id,
                    )
                    raise
            finally:
                closer = getattr(client, "close", None)
                if callable(closer):
                    try:
                        closer()
                    except Exception:
                        pass

            usage = _extract_gemini_usage(response)

            if response.text:
                result_text = response.text
                logger.info(
                    "Gemini generate_content succeeded: model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(result_text or ""),
                    request_id,
                )
                return LLMResponse(text=result_text, usage=usage, provider="gemini", model=model)

            candidate_texts: list[str] = []
            for candidate in getattr(response, "candidates", []) or []:
                content_obj = getattr(candidate, "content", None)
                if not content_obj:
                    continue
                for part in getattr(content_obj, "parts", []) or []:
                    text = getattr(part, "text", None)
                    if text:
                        candidate_texts.append(text)

            if candidate_texts:
                result_text = "\n".join(candidate_texts)
                logger.info(
                    "Gemini generate_content succeeded (candidates): model=%s images=%d text_len=%d request_id=%s",
                    model,
                    len(images or []),
                    len(result_text or ""),
                    request_id,
                )
                return LLMResponse(text=result_text, usage=usage, provider="gemini", model=model)

            logger.info(
                "Gemini generate_content succeeded with no text: model=%s images=%d request_id=%s",
                model,
                len(images or []),
                request_id,
            )
            return LLMResponse(text="", usage=usage, provider="gemini", model=model)

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    async def async_generate_response(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> str:
        return await run_sync_in_thread(
            lambda: self.generate_response(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    async def async_generate_response_with_usage(
        self,
        *,
        api_key: str,
        prompt: Optional[str] = None,
        model: str,
        max_tokens: int = 32000,
        reasoning_effort: Optional[str] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        images: Optional[Sequence[ImageInput]] = None,
        messages: Optional[MessageSequence] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> LLMResponse:
        return await run_sync_in_thread(
            lambda: self.generate_response_with_usage(
                api_key=api_key,
                prompt=prompt,
                model=model,
                max_tokens=max_tokens,
                reasoning_effort=reasoning_effort,
                temperature=temperature,
                top_p=top_p,
                images=images,
                messages=messages,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    def generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> bytes:
        """Generate an image using Gemini 3 Pro Image.

        Args:
            api_key: API key used to authenticate with the Gemini API.
            prompt: Text prompt for image generation.
            model: Identifier of the Gemini model to target (e.g., "gemini-3-pro-image-preview").
            image_size: Size of the generated image (e.g., "2K", "4K"). Defaults to "2K".
            aspect_ratio: Aspect ratio of the generated image (e.g., "16:9", "4:3").
            image: Optional input image for editing tasks.

        Returns:
            The generated image data as bytes.

        Raises:
            ValueError: If required arguments are missing or no image is returned.
            google.genai.errors.APIError: If the underlying Gemini request fails.
        """
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not prompt:
            raise ValueError("prompt must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        config = types.GenerateContentConfig(
            tools=[{"google_search": {}}],
            image_config=types.ImageConfig(
                image_size=image_size or "2K",
                aspect_ratio=aspect_ratio,
            ),
        )

        contents = [prompt]
        if image:
            contents.append(self._to_image_part(image))

        retry_count = clamp_retries(max_retries)

        def _build_client() -> genai.Client:
            client_kwargs: dict[str, object] = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["http_options"] = types.HttpOptions(
                    timeout=_normalize_gemini_timeout_ms(timeout_s)
                )
            return genai.Client(**client_kwargs)

        def _run_request() -> bytes:
            client = _build_client()
            try:
                try:
                    response = client.models.generate_content(
                        model=model,
                        contents=contents,
                        config=config,
                    )
                except Exception as exc:
                    logger.exception(
                        "Gemini generate_image failed: %s request_id=%s",
                        exc,
                        request_id,
                    )
                    raise
            finally:
                closer = getattr(client, "close", None)
                if callable(closer):
                    try:
                        closer()
                    except Exception:
                        pass

            if not response.parts:
                raise ValueError("No content returned from Gemini.")

            for part in response.parts:
                if part.inline_data:
                    return part.inline_data.data

            raise ValueError("No image data found in response.")

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    async def async_generate_image(
        self,
        *,
        api_key: str,
        prompt: str,
        model: str,
        image_size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        image: Optional[ImageInput] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> bytes:
        return await run_sync_in_thread(
            lambda: self.generate_image(
                api_key=api_key,
                prompt=prompt,
                model=model,
                image_size=image_size,
                aspect_ratio=aspect_ratio,
                image=image,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    def list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[dict[str, Optional[str]]]:
        """Return the models available to the authenticated Gemini account."""
        if not api_key:
            raise ValueError("api_key must be provided.")

        retry_count = clamp_retries(max_retries)

        def _build_client() -> genai.Client:
            client_kwargs: dict[str, object] = {"api_key": api_key}
            if timeout_s is not None:
                client_kwargs["http_options"] = types.HttpOptions(
                    timeout=_normalize_gemini_timeout_ms(timeout_s)
                )
            return genai.Client(**client_kwargs)

        def _run_request() -> list[dict[str, Optional[str]]]:
            models: list[dict[str, Optional[str]]] = []
            client = _build_client()
            try:
                try:
                    iterator = client.models.list()
                except Exception as exc:
                    logger.exception(
                        "Gemini list models failed: %s request_id=%s",
                        exc,
                        request_id,
                    )
                    raise
                for model in iterator:
                    model_id = getattr(model, "name", None)
                    if model_id is None and isinstance(model, dict):
                        model_id = model.get("name")
                    if not model_id:
                        continue

                    # Normalize IDs like "models/<id>" -> "<id>"
                    if isinstance(model_id, str) and model_id.startswith("models/"):
                        model_id = model_id.split("/", 1)[1]

                    display_name = getattr(model, "display_name", None)
                    if display_name is None and isinstance(model, dict):
                        display_name = model.get("display_name")

                    models.append({"id": model_id, "display_name": display_name})
            finally:
                closer = getattr(client, "close", None)
                if callable(closer):
                    try:
                        closer()
                    except Exception:
                        pass

            logger.info(
                "Gemini list_models succeeded: count=%d request_id=%s",
                len(models),
                request_id,
            )
            return models

        return run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )

    async def async_list_models(
        self,
        *,
        api_key: str,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[dict[str, Optional[str]]]:
        return await run_sync_in_thread(
            lambda: self.list_models(
                api_key=api_key,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    def embed_content(
        self,
        *,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        if not api_key:
            raise ValueError("api_key must be provided.")
        if not model:
            raise ValueError("model must be provided.")

        if isinstance(contents, str):
            payload: Union[str, list[str]] = contents
        else:
            payload = list(contents)
            if not payload:
                raise ValueError("contents must not be empty.")

        retry_count = clamp_retries(max_retries)

        def _build_client() -> genai.Client:
            client_kwargs: dict[str, object] = {"api_key": api_key}
            if timeout_s is not None:
                http_options = getattr(types, "HttpOptions", None)
                if http_options is not None:
                    try:
                        client_kwargs["http_options"] = http_options(
                            timeout=_normalize_gemini_timeout_ms(timeout_s)
                        )
                    except Exception:
                        logger.debug("Gemini HttpOptions timeout not applied.", exc_info=True)
            return genai.Client(**client_kwargs)

        config_kwargs: dict[str, object] = {}
        if task_type is not None:
            config_kwargs["task_type"] = task_type
        if output_dimensionality is not None:
            config_kwargs["output_dimensionality"] = output_dimensionality
        config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None

        def _run_request() -> list[EmbeddingVector]:
            client = _build_client()
            try:
                result = client.models.embed_content(
                    model=model,
                    contents=payload,
                    config=config,
                )
                embeddings = getattr(result, "embeddings", None)
                if embeddings is None:
                    raise ValueError("Gemini embeddings response missing embeddings field.")
                vectors: list[EmbeddingVector] = []
                for embedding in embeddings:
                    values = getattr(embedding, "values", None)
                    if values is None:
                        raise ValueError("Gemini embedding missing values field.")
                    vectors.append(list(values))
                return vectors
            finally:
                closer = getattr(client, "close", None)
                if callable(closer):
                    try:
                        closer()
                    except Exception:
                        pass

        vectors = run_with_retries(
            func=_run_request,
            max_retries=retry_count,
            retry_backoff_s=retry_backoff_s,
            request_id=request_id,
        )
        logger.info(
            "Gemini embed_content succeeded: count=%d request_id=%s",
            len(vectors),
            request_id,
        )
        return vectors

    async def async_embed_content(
        self,
        *,
        api_key: str,
        model: str,
        contents: Union[str, Sequence[str]],
        task_type: Optional[str] = None,
        output_dimensionality: Optional[int] = None,
        request_id: Optional[str] = None,
        timeout_s: Optional[float] = None,
        max_retries: Optional[int] = None,
        retry_backoff_s: float = 0.5,
    ) -> list[EmbeddingVector]:
        return await run_sync_in_thread(
            lambda: self.embed_content(
                api_key=api_key,
                model=model,
                contents=contents,
                task_type=task_type,
                output_dimensionality=output_dimensionality,
                request_id=request_id,
                timeout_s=timeout_s,
                max_retries=max_retries,
                retry_backoff_s=retry_backoff_s,
            )
        )

    @staticmethod
    def _to_image_part(image: ImageInput) -> types.Part:
        """Convert an image reference into a Gemini SDK part."""
        if isinstance(image, Path):
            return _part_from_path(image)

        if image.startswith("data:"):
            return _part_from_data_url(image)

        if image.startswith(("http://", "https://")):
            return _part_from_url(image)

        return _part_from_path(Path(image))


def _part_from_path(path: Path) -> types.Part:
    """Create an image part from a local filesystem path."""
    # Ensure common audio types are recognized across platforms (used for transcription as well).
    mimetypes.add_type("audio/mp4", ".m4a")
    mimetypes.add_type("audio/mpeg", ".mp3")
    mimetypes.add_type("audio/wav", ".wav")
    mimetypes.add_type("audio/aac", ".aac")
    mimetypes.add_type("audio/webm", ".webm")

    expanded = path.expanduser()
    data = expanded.read_bytes()
    mime_type = mimetypes.guess_type(expanded.name)[0] or "application/octet-stream"
    return types.Part.from_bytes(data=data, mime_type=mime_type)


def _part_from_url(url: str) -> types.Part:
    """Create an image part by downloading content from a URL."""
    with urlopen(url) as response:
        data = response.read()
        mime_type = response.info().get_content_type()

    if not mime_type or mime_type == "application/octet-stream":
        mime_type = mimetypes.guess_type(url)[0] or "application/octet-stream"

    return types.Part.from_bytes(data=data, mime_type=mime_type)


def _part_from_data_url(data_url: str) -> types.Part:
    """Create an image part from a data URL."""
    header, encoded = data_url.split(",", 1)
    metadata = header[len("data:") :]
    mime_type = "application/octet-stream"

    if ";" in metadata:
        mime_type, _, metadata = metadata.partition(";")
    elif metadata:
        mime_type = metadata

    if "base64" in metadata:
        data = base64.b64decode(encoded)
    else:
        data = encoded.encode("utf-8")

    return types.Part.from_bytes(data=data, mime_type=mime_type or "application/octet-stream")


def _extract_gemini_usage(response: object) -> TokenUsage | None:
    usage_obj = getattr(response, "usage_metadata", None)
    if usage_obj is None:
        usage_obj = getattr(response, "usage", None)
    if usage_obj is None:
        return None

    input_tokens = getattr(usage_obj, "prompt_token_count", None)
    output_tokens = getattr(usage_obj, "candidates_token_count", None)
    total_tokens = getattr(usage_obj, "total_token_count", None)

    if input_tokens is None:
        input_tokens = getattr(usage_obj, "input_tokens", None)
    if output_tokens is None:
        output_tokens = getattr(usage_obj, "output_tokens", None)
    if total_tokens is None:
        total_tokens = getattr(usage_obj, "total_tokens", None)

    if isinstance(usage_obj, dict):
        input_tokens = usage_obj.get("prompt_token_count", usage_obj.get("input_tokens"))
        output_tokens = usage_obj.get("candidates_token_count", usage_obj.get("output_tokens"))
        total_tokens = usage_obj.get("total_token_count", usage_obj.get("total_tokens"))

    return TokenUsage(
        input_tokens=int(input_tokens) if isinstance(input_tokens, int) else None,
        output_tokens=int(output_tokens) if isinstance(output_tokens, int) else None,
        total_tokens=int(total_tokens) if isinstance(total_tokens, int) else None,
    )
