"""Google Generative AI provider for Gemini models.

Examples:
    >>> from ember.models.providers.google import GoogleProvider
"""

import logging
from typing import Any, Dict, Iterable, Mapping, Optional, Sequence

from tenacity import retry, stop_after_attempt, wait_exponential

from ember._internal.exceptions import (
    ModelProviderError,
    ProviderAPIError,
)
from ember.models.discovery.provider_api import DiscoveryProvider
from ember.models.discovery.registry import register_provider as register_discovery_provider
from ember.models.discovery.types import DiscoveredModel
from ember.models.providers.base import BaseProvider
from ember.models.providers.utils import (
    PromptComponents,
    ResponsesRequest,
    build_prompt_components,
    parse_responses_request,
    text_turns,
)
from ember.models.schemas import ChatResponse, EmbeddingResponse, UsageStats

genai: Any = None


_MODEL_ALIASES: Dict[str, str] = {
    # Legacy Gemini 1.0 identifiers → Gemini 1.5 latest endpoints
    "models/gemini-pro": "models/gemini-1.5-pro-latest",
    "models/gemini-pro-vision": "models/gemini-1.5-pro-latest",
    "models/gemini-1.5-pro": "models/gemini-1.5-pro-latest",
    # Ensure bare aliases without "-latest" resolve to the newest release
    "models/gemini-1.5-flash": "models/gemini-1.5-flash-latest",
}

_MODEL_CONTEXT_WINDOWS: Dict[str, int] = {
    "models/gemini-1.5-pro-latest": 2_000_000,
    "models/gemini-1.5-flash-latest": 1_000_000,
}


def _load_genai_module() -> Any:
    """Load the Google Generative AI module.

    Prefers the legacy ``google-generativeai`` SDK (which exposes
    ``configure``, ``GenerativeModel``, and ``embed_content``) because the
    provider implementation is written against that API surface.  Falls back
    to the newer ``google-genai`` SDK when the legacy package is absent.
    """
    global genai
    if genai is not None:
        return genai

    try:
        import google.generativeai as genai_import  # type: ignore
    except ImportError:
        try:
            from google import genai as genai_import  # type: ignore
        except ImportError as exc:  # pragma: no cover
            raise ModelProviderError(
                "Google Gen AI library not installed. "
                "Install with: pip install google-generativeai"
            ) from exc

    globals()["genai"] = genai_import  # type: ignore[assignment]
    return genai_import


logger = logging.getLogger(__name__)


def _deep_convert_proto(value: Any) -> Any:
    """Recursively convert a protobuf Struct/MapComposite to plain Python types.

    ``dict(protobuf_struct)`` does a *shallow* conversion — nested lists remain
    as ``RepeatedComposite`` containers that ``json.dumps`` cannot serialize.
    This helper walks the tree and produces only built-in Python dicts, lists,
    strings, numbers, booleans, and None.
    """
    if value is None:
        return {}
    if isinstance(value, (str, int, float, bool)):
        return value
    if hasattr(value, "items"):
        return {k: _deep_convert_proto(v) for k, v in value.items()}
    if hasattr(value, "__iter__"):
        return [_deep_convert_proto(v) for v in value]
    return value


def _convert_tools_to_gemini(tools: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
    """Convert OpenAI-style function tools to Gemini format.

    Handles both nested Chat Completions and flat Responses formats:
    - Nested: {"type": "function", "function": {"name": "X", ...}}
    - Flat: {"type": "function", "name": "X", ...}

    Gemini only supports function-type tools. Non-function tools (web_search,
    file_search, etc.) are skipped with a warning since they would cause
    malformed function_declarations and fail at runtime.

    Args:
        tools: Sequence of OpenAI-style tool definitions.

    Returns:
        List of Gemini-compatible function declarations wrapped in a Tool.
    """
    function_declarations = []
    for tool in tools:
        tool_type = tool.get("type")
        if tool_type == "function":
            # Check flat format first (Responses API), then nested (Chat Completions)
            name = tool.get("name")
            description = tool.get("description")
            parameters = tool.get("parameters")

            if not name:
                # Nested format: {"type": "function", "function": {...}}
                func = tool.get("function", {})
                name = func.get("name", "")
                description = func.get("description", "")
                parameters = func.get("parameters")

            declaration = {
                "name": name or "",
                "description": description or "",
            }
            if parameters:
                declaration["parameters"] = parameters
            function_declarations.append(declaration)
        elif tool_type is None and "name" in tool:
            # Flat format without explicit type (assume function)
            declaration = {
                "name": tool.get("name", ""),
                "description": tool.get("description", ""),
            }
            parameters = tool.get("parameters")
            if parameters:
                declaration["parameters"] = parameters
            function_declarations.append(declaration)
        else:
            # Gemini only supports function tools; skip others with a warning
            logger.warning(
                "Gemini does not support tool type '%s'; skipping tool. "
                "Only 'function' tools are supported.",
                tool_type,
            )

    if not function_declarations:
        return []

    # Gemini expects tools as a list containing dicts with function_declarations
    return [{"function_declarations": function_declarations}]


def _convert_tool_choice_to_gemini(
    tool_choice: Mapping[str, Any] | str | None,
) -> dict[str, Any] | None:
    """Convert OpenAI tool_choice to Gemini tool_config format.

    Handles both nested Chat Completions and flat Responses formats:
    - Nested: {"type": "function", "function": {"name": "X"}}
    - Flat: {"type": "function", "name": "X"}

    Args:
        tool_choice: OpenAI-style tool_choice.

    Returns:
        Gemini tool_config dict or None.
    """
    if tool_choice is None:
        return None

    if isinstance(tool_choice, str):
        if tool_choice == "auto":
            return {"function_calling_config": {"mode": "AUTO"}}
        if tool_choice == "none":
            return {"function_calling_config": {"mode": "NONE"}}
        if tool_choice == "required":
            return {"function_calling_config": {"mode": "ANY"}}
        return None

    # Handle specific function selection
    choice_type = tool_choice.get("type")
    if choice_type == "function":
        # Check flat format first, then nested
        name = tool_choice.get("name")
        if not name:
            func_info = tool_choice.get("function", {})
            if isinstance(func_info, Mapping):
                name = func_info.get("name")
        if name:
            return {
                "function_calling_config": {
                    "mode": "ANY",
                    "allowed_function_names": [name],
                }
            }

    return None


def _resolve_google_api_key() -> str:
    """Resolve Google API key from the active Ember context."""
    from ember._internal.context.runtime import EmberContext

    return EmberContext.current().get_credential("google")


class GoogleProvider(BaseProvider):
    """Google Generative AI provider for Gemini chat models."""

    supports_responses_api = True

    def __init__(self, api_key: str | None = None):
        """Initialize Google provider.

        Args:
            api_key: Google API key override.
        """
        super().__init__(api_key)
        self._apply_pyasn1_patch()
        self._ensure_imports()
        # The legacy SDK uses module-level configure(); the new SDK uses
        # Client(api_key=...) and has no configure function.
        if hasattr(genai, "configure"):
            genai.configure(api_key=self.api_key)  # type: ignore[attr-defined]

    def _apply_pyasn1_patch(self) -> None:
        """Apply a pyasn1 compatibility patch.

        This defers the import and avoids issues during static type checking.

        Returns:
            None
        """
        try:
            import importlib

            mod = importlib.import_module("ember._internal.patches.pyasn1")
            func = getattr(mod, "ensure_pyasn1_compatibility", None)
            if callable(func):
                func()
        except Exception as exc:
            logger.warning("pyasn1 compatibility patch failed", exc_info=exc)

    def _ensure_imports(self) -> None:
        """Ensure Google Generative AI imports are available.

        Uses lazy loading to avoid protobuf issues at module import time.

        Returns:
            None
        """
        global genai
        if genai is not None:
            return

        try:
            module = _load_genai_module()
        except TypeError as e:
            if "got multiple values for keyword argument '_options'" in str(e):
                raise ModelProviderError(
                    "Google Gen AI import failed due to protobuf version conflict. "
                    "Install protobuf<3.20 or pin google-genai to a compatible version."
                ) from e
            raise

        globals()["genai"] = module  # type: ignore[assignment]

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=60),
        reraise=True,
    )
    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        """Complete a prompt using Google's API.

        Args:
            prompt: The prompt text.
            model: Model name (e.g., "gemini-1.5-pro-latest", "gemini-1.5-flash-latest").
            **kwargs: Additional parameters like temperature, max_tokens, etc.

        Returns:
            ChatResponse with completion and usage information.

        Raises:
            ProviderAPIError: For API errors.
            AuthenticationError: For invalid API key.
        """
        requested_model = model

        try:
            model = self._normalize_model_name(model)
        except ModelProviderError as exc:
            logger.error("Invalid Google model '%s': %s", requested_model, exc)
            raise ProviderAPIError(
                f"Unsupported Google model: {requested_model}",
                context={"model": requested_model},
            ) from exc

        # Get the model
        try:
            gemini_model = genai.GenerativeModel(model)  # type: ignore[attr-defined]
        except Exception as e:
            logger.error(f"Failed to create Gemini model: {e}")
            raise ProviderAPIError(
                f"Failed to create Gemini model: {str(e)}", context={"model": model}
            ) from e

        # Build generation config
        generation_config = {}

        # Map common parameters
        if "temperature" in kwargs:
            generation_config["temperature"] = kwargs.pop("temperature")
        if "max_tokens" in kwargs:
            generation_config["max_output_tokens"] = kwargs.pop("max_tokens")
        if "top_p" in kwargs:
            generation_config["top_p"] = kwargs.pop("top_p")
        if "stop" in kwargs:
            generation_config["stop_sequences"] = kwargs.pop("stop")

        # Support JSON output mode via response_format
        response_format = kwargs.pop("response_format", None)
        if response_format:
            if isinstance(response_format, dict):
                fmt_type = response_format.get("type", "")
                if fmt_type in ("json_object", "json"):
                    generation_config["response_mime_type"] = "application/json"
                    # If schema provided, pass it through
                    json_schema = response_format.get("json_schema")
                    if json_schema:
                        generation_config["response_schema"] = json_schema
            elif response_format == "json":
                generation_config["response_mime_type"] = "application/json"

        # Handle context/system message
        context = kwargs.pop("context", None)
        if context:
            prompt = f"{context}\n\n{prompt}"

        # Handle tools and tool_choice for function calling
        tools_raw = kwargs.pop("tools", None)
        tool_choice_raw = kwargs.pop("tool_choice", None)

        gemini_tools = None
        tool_config = None
        if tools_raw:
            gemini_tools = _convert_tools_to_gemini(tools_raw)
        if tool_choice_raw is not None:
            tool_config = _convert_tool_choice_to_gemini(tool_choice_raw)

        try:
            # Make API call
            logger.debug(f"Google API call: model={model}")

            # Build generate_content kwargs
            generate_kwargs: Dict[str, Any] = {}
            if generation_config:
                generate_kwargs["generation_config"] = generation_config
            if gemini_tools:
                generate_kwargs["tools"] = gemini_tools
            if tool_config:
                generate_kwargs["tool_config"] = tool_config

            response = gemini_model.generate_content(  # type: ignore[no-untyped-call]
                prompt,
                **generate_kwargs,
            )

            text = self._coerce_response_text(response)

            if not text:
                finish_reasons = [
                    getattr(candidate, "finish_reason", None)
                    for candidate in getattr(response, "candidates", []) or []
                ]
                logger.debug(
                    "Google response contained no text; finish_reasons=%s",
                    finish_reasons,
                )

            # Extract tool_calls from Gemini function calls
            tool_calls = self._extract_tool_calls(response)

            # Build usage stats
            usage = None
            try:
                um = getattr(response, "usage_metadata", None)
                if um is not None:
                    # Prefer SDK-provided metadata when available
                    pt = int(getattr(um, "prompt_token_count", 0) or 0)
                    ct = int(getattr(um, "candidates_token_count", 0) or 0)
                    tt = int(getattr(um, "total_token_count", 0) or (pt + ct))
                    usage = UsageStats(prompt_tokens=pt, completion_tokens=ct, total_tokens=tt)
            except Exception:
                # Fall back to estimate below
                usage = None

            if usage is None:
                # Estimate when metadata is unavailable
                pt_est = len(prompt.split()) * 2
                ct_est = len(text.split()) * 2
                usage = UsageStats(
                    prompt_tokens=pt_est,
                    completion_tokens=ct_est,
                    total_tokens=pt_est + ct_est,
                )

            return ChatResponse(
                data=text,
                usage=usage,
                model_id=model,
                raw_output=response,
                tool_calls=tool_calls if tool_calls else None,
            )

        except Exception as e:
            error_str = str(e)

            if "API_KEY_INVALID" in error_str or "invalid api key" in error_str.lower():
                logger.error(f"Google authentication error: {e}")
                raise ProviderAPIError(
                    "Invalid Google API key",
                    context={"model": model, "error_type": "authentication"},
                ) from e

            elif "RATE_LIMIT_EXCEEDED" in error_str:
                logger.warning(f"Google rate limit: {e}")
                raise ProviderAPIError(
                    "Google rate limit exceeded",
                    context={"model": model, "error_type": "rate_limit"},
                ) from e

            else:
                logger.error(f"Google API error: {e}")
            raise ProviderAPIError(
                f"Google API error: {error_str}", context={"model": model}
            ) from e

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=30),
        reraise=True,
    )
    def embed(
        self,
        inputs: Sequence[str],
        model: str,
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Generate embeddings using Google's embedding models.

        Args:
            inputs: Sequence of text strings to embed.
            model: Embedding model name (e.g., "gemini-embedding-001").
            **kwargs: Additional parameters like dimensions.

        Returns:
            EmbeddingResponse with embedding vectors and usage info.
        """
        dimensions = kwargs.pop("dimensions", None)

        # Normalize model name - embedding models use different naming
        if not model.startswith("models/"):
            model = f"models/{model}"

        try:
            # Google API supports batch embedding - pass list directly
            # Max ~100 items per request, so chunk if needed
            chunk_size = 100
            embeddings = []
            total_tokens = 0

            for i in range(0, len(inputs), chunk_size):
                chunk = list(inputs[i : i + chunk_size])
                result = genai.embed_content(  # type: ignore[attr-defined]
                    model=model,
                    content=chunk,
                    output_dimensionality=dimensions,
                )

                # Extract embeddings from batch result
                chunk_embeddings = result.get("embedding", [])
                if not chunk_embeddings and hasattr(result, "embedding"):
                    chunk_embeddings = result.embedding

                # Handle both single and batch responses
                if chunk_embeddings and not isinstance(chunk_embeddings[0], list):
                    # Single embedding returned as flat list
                    embeddings.append(chunk_embeddings)
                else:
                    embeddings.extend(chunk_embeddings)

                # Estimate tokens
                total_tokens += sum(len(t.split()) * 2 for t in chunk)

            usage = UsageStats(
                prompt_tokens=total_tokens,
                completion_tokens=0,
                total_tokens=total_tokens,
            )

            return EmbeddingResponse(
                embeddings=embeddings,
                usage=usage,
                model_id=model,
                raw_output=None,
            )

        except Exception as e:
            error_str = str(e)
            if "API_KEY_INVALID" in error_str or "invalid api key" in error_str.lower():
                raise ProviderAPIError(
                    "Invalid Google API key",
                    context={"model": model, "error_type": "authentication"},
                ) from e
            raise ProviderAPIError(
                f"Google embedding error: {error_str}",
                context={"model": model},
            ) from e

    def complete_responses_payload(
        self,
        payload: Mapping[str, Any],
        **overrides: Any,
    ) -> ChatResponse:
        request = self._parse_responses_payload(payload)
        if overrides:
            unsupported = ", ".join(sorted(overrides))
            raise ProviderAPIError(
                f"Unsupported Gemini Responses overrides: {unsupported}",
                context={"model": request.model},
            )

        self._validate_responses_payload(request)
        turns = text_turns(request)
        if not turns:
            raise ProviderAPIError(
                "Responses payload must include at least one user message",
                context={"model": request.model},
            )

        components = build_prompt_components(request)
        params: Dict[str, Any] = {}
        if request.max_output_tokens is not None:
            params["max_tokens"] = request.max_output_tokens
        if request.temperature is not None:
            params["temperature"] = request.temperature
        if request.top_p is not None:
            params["top_p"] = request.top_p
        if request.stop_sequences:
            params["stop"] = list(request.stop_sequences)
        if request.response_format:
            params["response_format"] = request.response_format

        # Pass tools and tool_choice for Gemini function calling
        if request.tools:
            params["tools"] = list(request.tools)
        if request.tool_choice is not None:
            params["tool_choice"] = request.tool_choice

        context_blob = self._compose_context_text(components)

        return self.complete(
            components.prompt,
            request.model,
            context=context_blob,
            **params,
        )

    def stream_responses_payload(
        self,
        payload: Mapping[str, Any],
        **_: Any,
    ) -> Iterable[str]:
        raise ProviderAPIError(
            "Streaming is not implemented for Gemini Responses payloads",
            context={"model": payload.get("model")},
        )

    @staticmethod
    def _parse_responses_payload(payload: Mapping[str, Any]) -> ResponsesRequest:
        if not isinstance(payload, Mapping):
            raise ProviderAPIError("Responses payload must be a JSON object")
        try:
            return parse_responses_request(dict(payload))
        except ValueError as exc:
            raise ProviderAPIError(str(exc)) from exc

    @staticmethod
    def _validate_responses_payload(request: ResponsesRequest) -> None:
        # Tools and tool_choice are now supported via Gemini function calling
        # response_format is now supported for JSON output mode
        # parallel_tool_calls is NOT supported by Gemini - warn if explicitly disabled
        if request.parallel_tool_calls is False:
            logger.warning(
                "Gemini does not support parallel_tool_calls=False; parameter ignored "
                "(model=%s)",
                request.model,
            )
        if request.response_format:
            fmt = request.response_format
            if isinstance(fmt, dict):
                fmt_type = fmt.get("type", "")
                if fmt_type not in ("json_object", "json", "text"):
                    raise ProviderAPIError(
                        f"Gemini Responses adapter only supports json/json_object/text "
                        f"response_format types, got: {fmt_type}",
                        context={"model": request.model},
                    )
        if request.reasoning is not None:
            raise ProviderAPIError(
                "Gemini Responses adapter does not support reasoning controls",
                context={"model": request.model},
            )

    @staticmethod
    def _compose_context_text(components: PromptComponents) -> Optional[str]:
        segments: list[str] = []
        if components.system:
            segments.append(components.system)
        for entry in components.context:
            role = entry.get("role", "context")
            content = entry.get("content", "")
            if content:
                segments.append(f"{role}: {content}")

        if not segments:
            return None
        return "\n\n".join(segment.strip() for segment in segments if segment)

    def _coerce_response_text(self, response: Any) -> str:
        """Return best-effort text from a Gemini response.

        The Google SDK raises when ``response.text`` is accessed for candidates that
        terminate with ``MAX_TOKENS`` or safety blocks. We prefer the quick accessor
        when available and otherwise fold over the candidates manually.
        """

        try:
            text_value = response.text
        except Exception as text_error:  # pragma: no cover
            logger.debug(
                "Google response text accessor failed; falling back to candidate parts",
                exc_info=text_error,
            )
        else:
            if isinstance(text_value, str) and text_value:
                return text_value

        return self._extract_text_from_candidates(response)

    def _extract_tool_calls(self, response: Any) -> list[dict[str, Any]]:
        """Extract function calls from Gemini response and convert to OpenAI tool_calls format.

        Gemini returns function calls as parts with function_call attribute containing
        name and args. This converts them to OpenAI's tool_calls format with id, type,
        and function (name + arguments as JSON string).
        """
        import json
        import uuid

        tool_calls: list[dict[str, Any]] = []
        candidates = getattr(response, "candidates", None)
        if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)):
            return []

        for candidate in candidates:
            parts = self._coerce_candidate_parts(candidate)
            if not parts:
                continue

            for part in parts:
                # Check for function_call attribute (Gemini function calling)
                function_call = getattr(part, "function_call", None)
                if function_call is None and isinstance(part, dict):
                    function_call = part.get("function_call")

                if function_call is not None:
                    # Extract name and args from function_call
                    fc_name = getattr(function_call, "name", None)
                    fc_args = getattr(function_call, "args", None)
                    if fc_name is None and isinstance(function_call, dict):
                        fc_name = function_call.get("name")
                        fc_args = function_call.get("args")

                    if fc_name:
                        # Convert args to JSON string for OpenAI compatibility.
                        # Gemini args is a protobuf Struct whose dict() does a
                        # *shallow* conversion — nested lists remain as
                        # RepeatedComposite and fail json.dumps.  Deep-convert.
                        args_dict = _deep_convert_proto(fc_args)
                        arguments = json.dumps(args_dict) if args_dict else "{}"

                        # Generate a unique ID (Gemini doesn't provide one)
                        call_id = f"call_{uuid.uuid4().hex[:24]}"

                        tool_calls.append({
                            "id": call_id,
                            "type": "function",
                            "function": {
                                "name": fc_name,
                                "arguments": arguments,
                            },
                        })

        return tool_calls

    @staticmethod
    def _extract_text_from_candidates(response: Any) -> str:
        """Safely materialize text from candidate parts when quick accessor fails."""

        candidates = getattr(response, "candidates", None) or []
        collected: list[str] = []

        for candidate in candidates:
            parts = GoogleProvider._coerce_candidate_parts(candidate)
            if not parts:
                continue

            for part in parts:
                text = GoogleProvider._coerce_part_text(part)
                if text:
                    collected.append(text)

        return "\n".join(segment.strip() for segment in collected if segment).strip()

    @staticmethod
    def _coerce_candidate_parts(candidate: Any) -> Iterable[Any]:
        """Return the most likely ``parts`` iterable from a candidate."""

        content = getattr(candidate, "content", None)
        parts = None
        if content is not None:
            parts = getattr(content, "parts", None)
            if parts is None and isinstance(content, dict):
                parts = content.get("parts")

        if parts is None and hasattr(candidate, "parts"):
            parts = getattr(candidate, "parts", None)

        if not parts:
            return ()

        if isinstance(parts, (list, tuple)):
            return parts

        if isinstance(parts, Iterable) and not isinstance(parts, (str, bytes)):
            return parts

        return ()

    @staticmethod
    def _coerce_part_text(part: Any) -> Optional[str]:
        """Extract the ``text`` field from a part-like object."""

        if isinstance(part, str):
            return part

        text = getattr(part, "text", None)
        if text is None and isinstance(part, dict):
            text = part.get("text")

        return text if text else None

    def _get_api_key_from_env(self) -> str:
        """Resolve Google API key from the active Ember configuration."""
        from ember._internal.context.runtime import EmberContext

        return EmberContext.current().get_credential("google")

    def validate_model(self, model: str) -> bool:
        """Check if this provider supports the given model.

        Args:
            model: Model name to validate.

        Returns:
            True if model is supported.
        """
        try:
            normalized = self._normalize_model_name(model)
        except ModelProviderError:
            return False

        return normalized.startswith("models/gemini")

    def _normalize_model_name(self, model: str) -> str:
        """Normalize user-supplied model identifiers to Google SDK expected ids."""

        trimmed = (model or "").strip()
        if not trimmed:
            raise ModelProviderError("Model identifier cannot be empty.")

        candidate = trimmed if trimmed.startswith("models/") else f"models/{trimmed}"
        canonical = _MODEL_ALIASES.get(candidate, candidate)

        if not canonical.startswith("models/gemini"):
            raise ModelProviderError(f"Model '{model}' is not a supported Google Gemini model.")

        return canonical

    def get_model_info(self, model: str) -> Dict[str, Any]:
        """Get information about a specific model.

        Args:
            model: Model name.

        Returns:
            Dictionary with model information.
        """
        info = super().get_model_info(model)

        try:
            canonical = self._normalize_model_name(model)
        except ModelProviderError:
            canonical = model if model.startswith("models/") else f"models/{model}"

        context_window = _MODEL_CONTEXT_WINDOWS.get(canonical)
        if context_window is None:
            if "-flash" in canonical:
                context_window = 1_000_000
            elif "-pro" in canonical or "-ultra" in canonical:
                context_window = 2_000_000
            else:
                context_window = 32768

        supports_vision = any(token in canonical for token in ("-flash", "-pro", "-vision"))

        # Add Google-specific information
        info.update(
            {
                "context_window": context_window,
                "supports_vision": supports_vision,
                "supports_functions": True,  # Gemini supports function calling
                "canonical_model": canonical,
            }
        )

        return info


class GoogleDiscoveryAdapter(DiscoveryProvider):
    """Discovery adapter for Google Gemini models."""

    name = "google"

    def __init__(
        self,
        *,
        api_key_resolver=_resolve_google_api_key,
        module_loader=_load_genai_module,
    ) -> None:
        self._resolve_api_key = api_key_resolver
        self._load_module = module_loader

    def list_models(
        self, *, region: Optional[str] = None, project_hint: Optional[str] = None
    ) -> Iterable[DiscoveredModel]:
        api_key = self._resolve_api_key()
        if not api_key:
            logger.debug("Skipping Google discovery: no API key configured")
            return []

        try:
            module = self._load_module()
        except ModelProviderError as exc:  # pragma: no cover
            logger.debug("Google discovery unavailable: %s", exc)
            return []

        if region:
            logger.debug("Google discovery region hint '%s' not currently applied", region)

        try:
            models = [
                self._to_discovered(model)
                for model in self._iter_models(module, api_key, region, project_hint)
            ]
        except Exception as exc:
            logger.warning("Google discovery failed: %s", exc)
            return []

        return models

    def _iter_models(
        self,
        module: Any,
        api_key: str,
        region: Optional[str],
        project_hint: Optional[str],
    ) -> Iterable[Any]:
        if hasattr(module, "Client"):
            client = module.Client(api_key=api_key)  # type: ignore[no-untyped-call]
            iterator = client.models.list(parent=project_hint)  # type: ignore[arg-type,no-untyped-call]
        else:
            if hasattr(module, "configure"):
                module.configure(api_key=api_key)
            list_kwargs: Dict[str, Any] = {}
            if project_hint:
                list_kwargs["project"] = project_hint
            iterator = module.list_models(**list_kwargs)  # type: ignore[call-arg]

        for model in iterator:
            # Capability filtering (generate/ embed) ensures runnable models.
            capabilities = _extract_sequence(
                model,
                "supported_generation_methods",
                "supportedGenerationMethods",
                "supported_actions",
                "supportedActions",
            )
            if capabilities and not {
                "generateContent",
                "embedContent",
            }.intersection(set(capabilities)):
                continue
            yield model

    def _to_discovered(self, model: Any) -> DiscoveredModel:
        raw_name = (
            getattr(model, "name", None)
            or getattr(model, "model", None)
            or getattr(model, "id", None)
            or ""
        )
        vendor_id = raw_name.split("/")[-1] if raw_name else getattr(model, "id", "")

        display_name = getattr(model, "display_name", None) or getattr(model, "displayName", None)
        description = getattr(model, "description", None)
        context_in = _extract_int(model, "input_token_limit", "inputTokenLimit")
        context_out = _extract_int(model, "output_token_limit", "outputTokenLimit")
        capabilities = tuple(
            _extract_sequence(
                model,
                "supported_generation_methods",
                "supportedGenerationMethods",
                "supported_actions",
                "supportedActions",
            )
        )
        region_scope = tuple(_extract_sequence(model, "supported_regions", "supportedRegions"))
        raw_payload = _extract_payload(model)

        return DiscoveredModel(
            provider="google",
            id=vendor_id,
            display_name=display_name,
            description=description,
            context_window_in=context_in,
            context_window_out=context_out,
            capabilities=capabilities,
            region_scope=region_scope,
            raw_payload=raw_payload,
        )


def _extract_sequence(model: Any, *names: str) -> list[str]:
    for name in names:
        value = getattr(model, name, None)
        if not value:
            continue
        if isinstance(value, (list, tuple)):
            return [str(item) for item in value if item]
        return [str(value)]
    return []


def _extract_int(model: Any, *names: str) -> Optional[int]:
    for name in names:
        value = getattr(model, name, None)
        if value is None:
            continue
        try:
            return int(value)
        except (TypeError, ValueError):
            continue
    return None


def _extract_payload(model: Any) -> Optional[Dict[str, Any]]:
    to_dict = getattr(model, "to_dict", None)
    if callable(to_dict):
        try:
            payload = to_dict()
            if isinstance(payload, dict):
                return payload
        except Exception:
            logger.debug(
                "to_dict() failed for %s; skipping payload extraction",
                type(model).__name__,
            )
            return None
    if isinstance(model, dict):
        return dict(model)
    return None


register_discovery_provider(GoogleDiscoveryAdapter())
