"""Google Gemini LLM provider for Patter's pipeline mode.

Backed by the ``google-genai`` SDK (``google.genai.Client``) which supports
both the Gemini Developer API and Vertex AI. The provider:

  * Translates OpenAI-formatted messages into Gemini ``Content`` turns
    (including ``function_call`` / ``function_response`` parts) so callers
    don't need a second message format.
  * Maps Gemini stream events (``text`` parts, ``function_call`` parts) to
    the Patter chunk protocol ``{"type": "text"|"tool_call"|"done"}``.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from enum import StrEnum
from typing import ClassVar, Any, AsyncIterator, Union

logger = logging.getLogger("getpatter")

__all__ = ["GoogleLLMProvider", "GoogleModel", "GoogleVertexLocation"]


class GoogleModel(StrEnum):
    """Known Google Gemini chat models."""

    GEMINI_2_5_FLASH = "gemini-2.5-flash"
    GEMINI_2_5_PRO = "gemini-2.5-pro"
    GEMINI_2_0_FLASH = "gemini-2.0-flash"
    GEMINI_2_0_FLASH_LITE = "gemini-2.0-flash-lite"
    GEMINI_1_5_FLASH = "gemini-1.5-flash"
    GEMINI_1_5_PRO = "gemini-1.5-pro"


class GoogleVertexLocation(StrEnum):
    """Common Vertex AI region codes accepted via the ``location`` arg."""

    US_CENTRAL1 = "us-central1"
    US_EAST1 = "us-east1"
    US_EAST4 = "us-east4"
    US_WEST1 = "us-west1"
    EUROPE_WEST1 = "europe-west1"
    EUROPE_WEST4 = "europe-west4"
    ASIA_NORTHEAST1 = "asia-northeast1"
    ASIA_SOUTHEAST1 = "asia-southeast1"


_DEFAULT_MODEL = GoogleModel.GEMINI_2_5_FLASH.value
_DEFAULT_VERTEX_LOCATION = GoogleVertexLocation.US_CENTRAL1.value


class GoogleLLMProvider:
    """LLM provider backed by Google Gemini (``google-genai`` SDK).

    Supports both the Gemini Developer API (with an API key) and
    Vertex AI (with Google Cloud credentials).  Streams chunks in
    Patter's ``{"type": "text" | "tool_call" | "done", ...}`` protocol.

    Args:
        api_key: Google API key for the Gemini Developer API. If
            omitted, ``GOOGLE_API_KEY`` is read from the environment.
            Ignored when ``vertexai=True``.
        model: Gemini model ID. Defaults to ``gemini-2.5-flash``.
        vertexai: If True, use Vertex AI instead of the Developer API.
        project: GCP project (Vertex AI only).
        location: GCP location (Vertex AI only). Defaults to
            ``us-central1``.
        temperature: Optional sampling temperature.
        max_output_tokens: Optional output token cap.
    """

    #: Stable pricing/dashboard key — read by stream-handler/metrics.
    provider_key: ClassVar[str] = "google"

    def __init__(
        self,
        api_key: str | None = None,
        model: Union[GoogleModel, str] = _DEFAULT_MODEL,
        vertexai: bool = False,
        project: str | None = None,
        location: str | None = None,
        temperature: float | None = None,
        max_output_tokens: int | None = None,
    ) -> None:
        try:
            from google.genai import Client
        except ImportError as e:
            raise RuntimeError(
                "The 'google-genai' package is required for GoogleLLMProvider. "
                "Install it with: pip install 'getpatter[google]'"
            ) from e

        use_vertexai = (
            vertexai
            if vertexai
            else os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "0").lower()
            in ["true", "1"]
        )

        resolved_key: str | None = None
        gcp_project = project or os.environ.get("GOOGLE_CLOUD_PROJECT")
        gcp_location = (
            location
            or os.environ.get("GOOGLE_CLOUD_LOCATION")
            or _DEFAULT_VERTEX_LOCATION
        )

        if use_vertexai:
            if not gcp_project:
                raise ValueError(
                    "Project is required for Vertex AI, either via the 'project' "
                    "argument or the GOOGLE_CLOUD_PROJECT environment variable."
                )
        else:
            resolved_key = api_key or os.environ.get("GOOGLE_API_KEY")
            if not resolved_key:
                raise ValueError(
                    "Google API key is required, either as the 'api_key' argument "
                    "or via the GOOGLE_API_KEY environment variable."
                )
            gcp_project = None
            gcp_location = None

        self._client = Client(
            api_key=resolved_key,
            vertexai=use_vertexai,
            project=gcp_project,
            location=gcp_location,
        )
        self._model = model
        self._temperature = temperature
        self._max_output_tokens = max_output_tokens

    async def warmup(self) -> None:
        """Pre-call DNS / TLS warmup for the Gemini API.

        Issues a lightweight ``GET https://generativelanguage.googleapis.com/v1beta/models``
        so DNS, TLS, and HTTP/2 are already up by the time the first
        ``generate_content_stream`` call lands. Best-effort: 5 s timeout,
        all exceptions swallowed at DEBUG. Skipped on Vertex AI which
        requires Application Default Credentials we don't want to mint
        for a probe.
        """
        try:
            # google-genai's Client doesn't expose the API key once it has
            # constructed its inner http client; fall back to env var.
            key = os.environ.get("GOOGLE_API_KEY") or ""
            if not key:
                return
            import httpx

            async with httpx.AsyncClient(timeout=5.0) as http:
                await http.get(
                    "https://generativelanguage.googleapis.com/v1beta/models",
                    params={"key": key},
                )
        except Exception as exc:  # noqa: BLE001 - best-effort
            logger.debug("Google LLM warmup failed (best-effort): %s", exc)

    async def stream(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        *,
        cancel_event: asyncio.Event | None = None,
        call_id: str | None = None,
    ) -> AsyncIterator[dict]:
        """Stream chunks from Gemini's ``generate_content_stream``.

        ``call_id`` is accepted for protocol parity with session-aware
        providers but ignored — Gemini has no per-call session model.

        ``cancel_event`` (set on barge-in by the stream handler) is checked
        between chunks and short-circuits the stream so the underlying
        request is freed immediately instead of blocking the next user
        transcript behind a long-running fetch.
        """
        from google.genai import types

        system_instruction, contents = _to_gemini_contents(messages)
        gemini_tools = _to_gemini_tools(tools) if tools else None

        config_kwargs: dict[str, Any] = {}
        if system_instruction:
            config_kwargs["system_instruction"] = [types.Part(text=system_instruction)]
        if gemini_tools is not None:
            config_kwargs["tools"] = gemini_tools
        if self._temperature is not None:
            config_kwargs["temperature"] = self._temperature
        if self._max_output_tokens is not None:
            config_kwargs["max_output_tokens"] = self._max_output_tokens

        config = types.GenerateContentConfig(**config_kwargs) if config_kwargs else None

        stream = await self._client.aio.models.generate_content_stream(
            model=self._model,
            contents=contents,
            config=config,
        )

        # Gemini does not provide a stable per-function-call index across
        # chunks, so we assign a monotonically increasing index per
        # function_call part that we see.
        next_index = 0

        last_usage = None
        async for response in stream:
            if cancel_event is not None and cancel_event.is_set():
                return
            # Gemini emits ``usage_metadata`` cumulatively on each chunk.
            # Capture only the most recent value so we yield ONE usage
            # event with the final totals (avoids double-counting).
            usage = getattr(response, "usage_metadata", None)
            if usage is not None:
                last_usage = usage

            if not getattr(response, "candidates", None):
                continue
            candidate = response.candidates[0]
            content = getattr(candidate, "content", None)
            if not content or not getattr(content, "parts", None):
                continue

            for part in content.parts:
                function_call = getattr(part, "function_call", None)
                if function_call:
                    args = getattr(function_call, "args", {}) or {}
                    fn_call_id = (
                        getattr(function_call, "id", None)
                        or f"gemini_call_{next_index}"
                    )
                    yield {
                        "type": "tool_call",
                        "index": next_index,
                        "id": fn_call_id,
                        "name": getattr(function_call, "name", "") or "",
                        "arguments": json.dumps(args),
                    }
                    next_index += 1
                    continue

                text = getattr(part, "text", None)
                if text:
                    yield {"type": "text", "content": text}

        if last_usage is not None:
            prompt_tokens = getattr(last_usage, "prompt_token_count", 0) or 0
            completion_tokens = getattr(last_usage, "candidates_token_count", 0) or 0
            cached_tokens = (
                getattr(last_usage, "cached_content_token_count", 0) or 0
            )
            self._record_completion_cost(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
            # ``prompt_token_count`` INCLUDES the cached portion; the usage
            # chunk contract bills ``input_tokens`` at the full input rate and
            # ``cache_read_tokens`` at the cache rate, so subtract to avoid
            # double-billing cached tokens (mirrors the OpenAI providers).
            yield {
                "type": "usage",
                "input_tokens": max(0, prompt_tokens - cached_tokens),
                "output_tokens": completion_tokens,
                "cache_read_tokens": cached_tokens,
            }

        yield {"type": "done"}

    def _record_completion_cost(
        self, *, prompt_tokens: int, completion_tokens: int
    ) -> None:
        """Stamp ``patter.cost.llm_*_tokens`` on the current span."""
        try:
            from getpatter.observability.attributes import record_patter_attrs

            record_patter_attrs(
                {
                    "patter.cost.llm_input_tokens": prompt_tokens,
                    "patter.cost.llm_output_tokens": completion_tokens,
                    "patter.llm.provider": "google",
                }
            )
        except Exception:  # pragma: no cover — defense in depth
            logger.debug("_record_completion_cost failed", exc_info=True)


# ---------------------------------------------------------------------------
# Message / tool translation (OpenAI format -> google.genai types)
# ---------------------------------------------------------------------------



# Keys Gemini's restricted proto ``Schema`` accepts. Anything else
# (``$schema``, ``additionalProperties`` — emitted by strict-mode tools and
# nearly every zod-derived MCP server — ``oneOf``, …) makes the request 400.
_GEMINI_SCHEMA_KEYS = frozenset(
    {
        "type",
        "description",
        "properties",
        "items",
        "enum",
        "required",
        "nullable",
        "format",
        "minimum",
        "maximum",
        "minLength",
        "maxLength",
        "minItems",
        "maxItems",
        "pattern",
        "anyOf",
        "default",
        "title",
    }
)


def _sanitize_gemini_schema(schema):
    """Recursively strip JSON-Schema keys Gemini's proto Schema rejects."""
    if isinstance(schema, dict):
        return {
            k: _sanitize_gemini_schema(v)
            for k, v in schema.items()
            if k in _GEMINI_SCHEMA_KEYS
        }
    if isinstance(schema, list):
        return [_sanitize_gemini_schema(v) for v in schema]
    return schema


def _to_gemini_tools(tools: list[dict]) -> list[Any]:
    """Convert OpenAI-style tool definitions to Gemini ``Tool`` objects."""
    from google.genai import types

    function_decls: list[types.FunctionDeclaration] = []
    for tool in tools:
        fn = tool.get("function", tool)
        function_decls.append(
            types.FunctionDeclaration(
                name=fn["name"],
                description=fn.get("description", ""),
                parameters=_sanitize_gemini_schema(
                    fn.get("parameters", {"type": "object", "properties": {}})
                ),
            )
        )
    if not function_decls:
        return []
    return [types.Tool(function_declarations=function_decls)]


def _to_gemini_contents(messages: list[dict]) -> tuple[str, list[Any]]:
    """Convert OpenAI-style messages to (system_instruction, Gemini Contents).

    Gemini expects:
      * ``system`` as a top-level ``GenerateContentConfig.system_instruction``.
      * ``user`` / ``model`` turns with ``Part`` lists.
      * Tool calls are ``Part(function_call=...)`` on ``model`` turns.
      * Tool results are ``Part(function_response=...)`` on ``user`` turns.
    """
    from google.genai import types

    system_parts: list[str] = []
    contents: list[types.Content] = []
    # tool_call_id → function name, harvested from assistant turns so the
    # paired FunctionResponse can carry the real function name. Gemini
    # requires ``FunctionResponse.name`` to match ``FunctionCall.name``;
    # the tool message itself only carries the call id.
    fn_name_by_call_id: dict[str, str] = {}

    for msg in messages:
        role = msg.get("role")

        if role == "system":
            content = msg.get("content")
            if isinstance(content, str) and content:
                system_parts.append(content)
            continue

        if role == "user":
            content = msg.get("content", "")
            if isinstance(content, str) and content:
                contents.append(
                    types.Content(role="user", parts=[types.Part(text=content)])
                )
            continue

        if role == "assistant":
            parts: list[types.Part] = []
            text = msg.get("content")
            if isinstance(text, str) and text:
                parts.append(types.Part(text=text))

            for tc in msg.get("tool_calls", []) or []:
                fn = tc.get("function", {})
                try:
                    args = json.loads(fn.get("arguments", "") or "{}")
                except json.JSONDecodeError:
                    args = {}
                if tc.get("id") and fn.get("name"):
                    fn_name_by_call_id[tc["id"]] = fn["name"]
                parts.append(
                    types.Part(
                        function_call=types.FunctionCall(
                            name=fn.get("name", ""),
                            args=args,
                            id=tc.get("id"),
                        )
                    )
                )

            if parts:
                contents.append(types.Content(role="model", parts=parts))
            continue

        if role == "tool":
            tool_call_id = msg.get("tool_call_id", "")
            raw = msg.get("content", "")
            try:
                response_dict = json.loads(raw) if isinstance(raw, str) else dict(raw)
                if not isinstance(response_dict, dict):
                    response_dict = {"result": response_dict}
            except (json.JSONDecodeError, TypeError):
                response_dict = {"result": raw}
            contents.append(
                types.Content(
                    role="user",
                    parts=[
                        types.Part(
                            function_response=types.FunctionResponse(
                                name=msg.get("name", "")
                                or fn_name_by_call_id.get(tool_call_id, "")
                                or tool_call_id,
                                response=response_dict,
                                id=tool_call_id or None,
                            )
                        )
                    ],
                )
            )
            continue

    # Gemini requires that all function_response parts for a single model turn
    # appear in ONE user Content entry (multiple parts). Mapping parallel
    # tool-call results one-per-message produces consecutive role='user'
    # Contents whose parts are exclusively function_response — which Gemini
    # rejects with a 400. Collapse those runs into a single Content so the
    # request is accepted. Single-tool and text turns are untouched.
    def _is_function_response_only(c: "types.Content") -> bool:
        parts = c.parts or []
        return (
            c.role == "user"
            and len(parts) > 0
            and all(getattr(p, "function_response", None) is not None for p in parts)
        )

    merged: list[types.Content] = []
    for entry in contents:
        if (
            merged
            and _is_function_response_only(merged[-1])
            and _is_function_response_only(entry)
        ):
            prev = merged[-1]
            merged[-1] = types.Content(
                role="user",
                parts=list(prev.parts or []) + list(entry.parts or []),
            )
        else:
            merged.append(entry)

    # Gemini expects the first content to be a user turn. Voice agents almost
    # always open with ``first_message`` (an assistant/model greeting) —
    # prepend a synthetic user turn so stricter backends don't 400.
    if merged and merged[0].role == "model":
        merged.insert(
            0,
            types.Content(role="user", parts=[types.Part(text="(call connected)")]),
        )

    return "\n\n".join(system_parts), merged
