# -*- coding: utf-8 -*-
"""A Google Gemini provider implementation using AgentScope's native
GeminiChatModel."""

from __future__ import annotations

import copy
import logging
import time
from typing import Any, List

from agentscope.model import ChatModelBase
from google import genai
from google.genai import errors as genai_errors
from google.genai import types as genai_types
from pydantic import Field

from qwenpaw.providers.multimodal_prober import (
    ProbeResult,
    _PROBE_IMAGE_B64,
    _IMAGE_PROBE_PROMPT,
    _PROBE_VIDEO_URL,
    _is_media_keyword_error,
    evaluate_image_probe_answer,
)
from qwenpaw.providers.provider import ModelInfo, Provider
from .capping_formatter import _CappingGeminiFormatter
from .capping_formatter import MAX_INLINE_MEDIA_BYTES

logger = logging.getLogger(__name__)


# TODO: Remove _flatten_json_schema and _sanitize_schema_for_gemini once
#  agentscope >= 2.0.5 is released (these are upstreamed into
#  GeminiChatModel._format_tools in newer versions).


def _flatten_json_schema(schema: dict) -> dict:
    """Flatten a JSON schema by resolving all ``$ref`` references.

    Gemini API does not support ``$defs`` and ``$ref`` in JSON schemas.
    """
    schema = copy.deepcopy(schema)
    defs = schema.pop("$defs", {})

    def _resolve_ref(obj: Any, visited: set | None = None) -> Any:
        if visited is None:
            visited = set()
        if not isinstance(obj, dict):
            if isinstance(obj, list):
                return [_resolve_ref(item, visited.copy()) for item in obj]
            return obj
        if "$ref" in obj:
            ref_path = obj["$ref"]
            if ref_path.startswith("#/$defs/"):
                def_name = ref_path[len("#/$defs/") :]
                if def_name in visited:
                    logger.warning(
                        "Circular reference detected for '%s' in tool schema",
                        def_name,
                    )
                    return {
                        "type": "object",
                        "description": f"(circular: {def_name})",
                    }
                visited.add(def_name)
                if def_name in defs:
                    resolved = _resolve_ref(defs[def_name], visited.copy())
                    for key, value in obj.items():
                        if key != "$ref":
                            resolved[key] = _resolve_ref(
                                value,
                                visited.copy(),
                            )
                    return resolved
            return obj
        result = {}
        for key, value in obj.items():
            result[key] = _resolve_ref(value, visited.copy())
        return result

    return _resolve_ref(schema)


# pylint: disable=too-many-branches
def _sanitize_schema_for_gemini(schema: Any) -> Any:
    """Sanitize a JSON schema to be compatible with the Gemini API.

    Removes or rewrites constructs that Gemini does not support:

    - ``additionalProperties``: removed entirely.
    - ``anyOf`` containing ``{"type": "null"}``: simplified to the single
      non-null type (i.e. ``Optional[X]`` becomes just ``X``).
    - ``{"type": "null"}`` appearing as a standalone schema: replaced with
      ``{"type": "object"}`` because the Gemini API (and many third-party
      proxies) do not accept ``null`` as a functionDeclaration property type.
    - ``const``: not supported by the Gemini SDK's ``Schema`` pydantic
      model (raises ``extra_forbidden``); rewritten to an ``enum`` with a
      single value, which is semantically equivalent and is supported.
    - All nested sub-schemas are processed recursively.
    """
    if not isinstance(schema, dict):
        if isinstance(schema, list):
            return [_sanitize_schema_for_gemini(v) for v in schema]
        return schema

    schema = dict(schema)

    # Replace standalone "type": "null" with "type": "object".
    # Some MCP servers emit ``{"type": "null"}`` for parameters that
    # accept ``None``; the Gemini API and many third-party proxies
    # reject ``null`` as a valid functionDeclaration schema type.
    if schema.get("type") == "null":
        schema["type"] = "object"

    schema.pop("additionalProperties", None)

    # JSON Schema ``const`` is not a recognized field on the Gemini SDK's
    # ``Schema`` pydantic model and triggers an ``extra_forbidden``
    # validation error. Rewrite it as a single-value ``enum``, which is
    # semantically equivalent and is supported by Gemini.
    if "const" in schema:
        const_value = schema.pop("const")
        schema.setdefault("enum", [const_value])

    if "anyOf" in schema and isinstance(schema["anyOf"], list):
        any_of = schema["anyOf"]
        non_null = [v for v in any_of if v != {"type": "null"}]
        if len(non_null) < len(any_of):
            if len(non_null) == 1:
                merged = dict(_sanitize_schema_for_gemini(non_null[0]))
                for k, v in schema.items():
                    if k != "anyOf":
                        merged.setdefault(k, v)
                return merged
            elif non_null:
                schema["anyOf"] = [
                    _sanitize_schema_for_gemini(v) for v in non_null
                ]
            else:
                del schema["anyOf"]

    for key in ["properties", "patternProperties", "$defs"]:
        if key in schema and isinstance(schema[key], dict):
            schema[key] = {
                k: _sanitize_schema_for_gemini(v)
                for k, v in schema[key].items()
            }

    for key in ["items", "not", "if", "then", "else"]:
        if key in schema:
            schema[key] = _sanitize_schema_for_gemini(schema[key])

    for key in ["allOf", "oneOf", "anyOf"]:
        if key in schema and isinstance(schema[key], list):
            schema[key] = [_sanitize_schema_for_gemini(v) for v in schema[key]]

    return schema


class GeminiProvider(Provider):
    """Provider implementation for Google Gemini API."""

    max_inline_media_bytes: int = Field(
        default=MAX_INLINE_MEDIA_BYTES,
        ge=0,
        description=(
            "Maximum size (in bytes) of a local media file inlined as "
            "base64 into the model request body. Media above this is "
            "replaced with a text placeholder to avoid oversized requests "
            "when large files (e.g. generated videos) persist in "
            "conversation history. 0 disables capping."
        ),
    )

    def _build_default_headers(self) -> dict:
        return dict(self.custom_headers) if self.custom_headers else {}

    def _client(self, timeout: float = 10) -> Any:
        headers = self._build_default_headers() or None
        return genai.Client(
            api_key=self.api_key,
            http_options=genai_types.HttpOptions(
                timeout=int(timeout * 1000),
                headers=headers,
            ),
        )

    @staticmethod
    def _normalize_models_payload(payload: Any) -> List[ModelInfo]:
        models: List[ModelInfo] = []
        for row in payload or []:
            model_id = str(getattr(row, "name", "") or "").strip()

            if not model_id:
                continue

            # Gemini API returns model names like "models/gemini-2.5-flash"
            # Strip the "models/" prefix for cleaner IDs
            if model_id.startswith("models/"):
                model_id = model_id[len("models/") :]

            display_name = str(
                getattr(row, "display_name", "") or model_id,
            ).strip()

            if not display_name or display_name.startswith("models/"):
                display_name = model_id

            models.append(ModelInfo(id=model_id, name=display_name))

        deduped: List[ModelInfo] = []
        seen: set[str] = set()
        for model in models:
            if model.id in seen:
                continue
            seen.add(model.id)
            deduped.append(model)
        return deduped

    async def check_connection(self, timeout: float = 10) -> tuple[bool, str]:
        """Check if Google Gemini provider is reachable."""
        try:
            client = self._client(timeout=timeout)
            # Use the async list models endpoint to verify connectivity
            async for _ in await client.aio.models.list():
                break
            return True, ""
        except genai_errors.APIError:
            return (
                False,
                "Failed to connect to Google Gemini API. "
                "Check your API key.",
            )
        except Exception:
            return (
                False,
                "Unknown exception when connecting to Google Gemini API.",
            )

    async def fetch_models(self, timeout: float = 10) -> List[ModelInfo]:
        """Fetch available models from Gemini API."""
        try:
            client = self._client(timeout=timeout)
            payload = []
            async for model in await client.aio.models.list():
                payload.append(model)
            models = self._normalize_models_payload(payload)
            return models
        except genai_errors.APIError:
            return []
        except Exception:
            return []

    async def check_model_connection(
        self,
        model_id: str,
        timeout: float = 10,
    ) -> tuple[bool, str]:
        """Check if a specific Gemini model is reachable/usable."""
        target = (model_id or "").strip()
        if not target:
            return False, "Empty model ID"

        try:
            client = self._client(timeout=timeout)
            response = await client.aio.models.generate_content_stream(
                model=target,
                contents="ping",
            )
            async for _ in response:
                break
            return True, ""
        except genai_errors.APIError:
            return (
                False,
                f"Model '{model_id}' is not reachable or usable",
            )
        except Exception:
            return (
                False,
                f"Unknown exception when connecting to model '{model_id}'",
            )

    @staticmethod
    def _adapt_generate_kwargs_for_gemini(
        kwargs: dict,
    ) -> dict:
        """Translate OpenAI-style keys to Gemini's GenerateContentConfig
        schema.

        google-genai's GenerateContentConfig forbids extra fields, so
        ``max_tokens`` must be renamed to ``max_output_tokens``.  If both are
        present, the explicit ``max_output_tokens`` wins.
        """
        adapted = dict(kwargs)
        max_tokens = adapted.pop("max_tokens", None)
        if max_tokens is not None and "max_output_tokens" not in adapted:
            adapted["max_output_tokens"] = max_tokens
        return adapted

    def get_chat_model_instance(self, model_id: str) -> ChatModelBase:
        from agentscope.credential import GeminiCredential
        from agentscope.model import GeminiChatModel

        credential = GeminiCredential(
            id=f"qwenpaw-{self.id}",
            api_key=self.api_key,
        )

        gen_kwargs = self._adapt_generate_kwargs_for_gemini(
            self.get_effective_generate_kwargs(model_id),
        )
        parameters = GeminiChatModel.Parameters(
            max_tokens=gen_kwargs.pop("max_output_tokens", None),
            temperature=gen_kwargs.pop("temperature", None),
            top_p=gen_kwargs.pop("top_p", None),
        )

        headers = self._build_default_headers()

        return _GeminiChatModelCompat(
            credential=credential,
            model=model_id,
            parameters=parameters,
            stream=True,
            default_headers=headers or None,
            extra_config_kwargs=gen_kwargs or None,
            context_size=self._get_context_size(model_id),
            formatter=_CappingGeminiFormatter(
                max_bytes=self.max_inline_media_bytes,
                relay_reasoning_content=self._get_relay_reasoning(model_id),
            ),
        )

    async def probe_model_multimodal(
        self,
        model_id: str,
        timeout: float = 60,
        image_only: bool = False,
    ) -> ProbeResult:
        """Probe multimodal support using Gemini generateContent API.

        Gemini supports both image and video via inline_data.  Each
        modality is probed independently with a minimal payload.
        """
        img_ok, img_msg = await self._probe_image_support(model_id, timeout)
        if image_only:
            return ProbeResult(
                supports_image=img_ok,
                supports_video=False,
                image_message=img_msg,
                video_message="Skipped: image_only=True",
            )
        vid_ok, vid_msg = await self._probe_video_support(model_id, timeout)
        return ProbeResult(
            supports_image=img_ok,
            supports_video=vid_ok,
            image_message=img_msg,
            video_message=vid_msg,
        )

    async def _probe_image_support(
        self,
        model_id: str,
        timeout: float = 15,
    ) -> tuple[bool, str]:
        """Probe image support via Gemini generateContent with inline_data.

        Sends a solid-red 16x16 PNG and asks the model to name the colour.
        """
        import base64

        logger.info(
            "Image probe start: model=%s url=%s",
            model_id,
            self.base_url,
        )
        start_time = time.monotonic()
        client = self._client(timeout=timeout)
        try:
            image_bytes = base64.b64decode(_PROBE_IMAGE_B64)
            response = await client.aio.models.generate_content(
                model=model_id,
                contents=[
                    genai_types.Part(
                        inline_data=genai_types.Blob(
                            mime_type="image/png",
                            data=image_bytes,
                        ),
                    ),
                    genai_types.Part(text=_IMAGE_PROBE_PROMPT),
                ],
                config=genai_types.GenerateContentConfig(
                    max_output_tokens=20,
                ),
            )
            answer = response.text or ""
            return evaluate_image_probe_answer(
                answer,
                model_id,
                start_time,
            )
        except genai_errors.APIError as e:
            elapsed = time.monotonic() - start_time
            logger.warning(
                "Image probe error: model=%s type=%s msg=%s %.2fs",
                model_id,
                type(e).__name__,
                e,
                elapsed,
            )
            status = getattr(e, "code", None)
            if status == 400 or _is_media_keyword_error(e):
                return False, f"Image not supported: {e}"
            return False, f"Probe inconclusive: {e}"
        except Exception as e:
            elapsed = time.monotonic() - start_time
            logger.warning(
                "Image probe error: model=%s type=%s msg=%s %.2fs",
                model_id,
                type(e).__name__,
                e,
                elapsed,
            )
            return False, f"Probe failed: {e}"

    async def _probe_video_support(
        self,
        model_id: str,
        timeout: float = 30,
    ) -> tuple[bool, str]:
        """Probe video support via Gemini generateContent with a video URL.

        Asks the model whether the video contains moving content.
        """
        logger.info(
            "Video probe start: model=%s url=%s",
            model_id,
            self.base_url,
        )
        start_time = time.monotonic()
        client = self._client(timeout=timeout)
        try:
            response = await client.aio.models.generate_content(
                model=model_id,
                contents=[
                    genai_types.Part(
                        file_data=genai_types.FileData(
                            file_uri=_PROBE_VIDEO_URL,
                            mime_type="video/mp4",
                        ),
                    ),
                    genai_types.Part(
                        text=(
                            "Does this contain moving content? "
                            "Reply with ONLY 'yes' or 'no', nothing else."
                        ),
                    ),
                ],
                config=genai_types.GenerateContentConfig(
                    max_output_tokens=10,
                ),
            )
            answer = (response.text or "").lower().strip()
            if "yes" in answer:
                result = True, f"Video supported (answer={answer!r})"
                elapsed = time.monotonic() - start_time
                logger.info(
                    "Video probe done: model=%s result=%s %.2fs",
                    model_id,
                    result[0],
                    elapsed,
                )
                return result
            result = (
                False,
                f"Model did not recognise video (answer={answer!r})",
            )
            elapsed = time.monotonic() - start_time
            logger.info(
                "Video probe done: model=%s result=%s %.2fs",
                model_id,
                result[0],
                elapsed,
            )
            return result
        except genai_errors.APIError as e:
            elapsed = time.monotonic() - start_time
            logger.warning(
                "Video probe error: model=%s type=%s msg=%s %.2fs",
                model_id,
                type(e).__name__,
                e,
                elapsed,
            )
            status = getattr(e, "code", None)
            if status == 400 or _is_media_keyword_error(e):
                return False, f"Video not supported: {e}"
            return False, f"Probe inconclusive: {e}"
        except Exception as e:
            elapsed = time.monotonic() - start_time
            logger.warning(
                "Video probe error: model=%s type=%s msg=%s %.2fs",
                model_id,
                type(e).__name__,
                e,
                elapsed,
            )
            return False, f"Probe failed: {e}"


class _GeminiChatModelCompat:
    """Factory that creates a ``GeminiChatModel`` subclass with custom headers
    and extra config kwargs injected into every API call."""

    def __new__(cls, **kwargs: Any) -> Any:
        from agentscope.model import GeminiChatModel

        default_headers = kwargs.pop("default_headers", None)
        extra_config_kwargs = kwargs.pop("extra_config_kwargs", None) or {}

        class _Compat(GeminiChatModel):
            _qp_default_headers = default_headers
            _qp_extra_config_kwargs = extra_config_kwargs

            # TODO: Remove this override once agentscope >= 2.0.5 is
            #  released (upstream _format_tools will handle sanitization).
            def _format_tools(self, tools, tool_choice):
                if tools:
                    sanitized = []
                    for schema in tools:
                        if "function" not in schema:
                            sanitized.append(schema)
                            continue
                        func = schema["function"].copy()
                        if "parameters" in func:
                            func["parameters"] = _sanitize_schema_for_gemini(
                                _flatten_json_schema(func["parameters"]),
                            )
                        sanitized.append({**schema, "function": func})
                    tools = sanitized
                return super()._format_tools(tools, tool_choice)

            async def _call_api(
                self,
                model_name,
                messages,
                tools=None,
                tool_choice=None,
                **config_kwargs,
            ):
                # Translate the neutral ``disable_thinking`` flag
                if config_kwargs.pop("disable_thinking", False):
                    self.parameters.thinking_enable = False
                merged = {**self._qp_extra_config_kwargs, **config_kwargs}
                if self._qp_default_headers:
                    from datetime import datetime

                    client = genai.Client(
                        api_key=self.credential.api_key.get_secret_value(),
                        http_options=genai_types.HttpOptions(
                            headers=self._qp_default_headers,
                        ),
                    )

                    formatted = await self.formatter.format(messages)
                    config: dict[str, Any] = {**merged}
                    if self.parameters.max_tokens is not None:
                        config[
                            "max_output_tokens"
                        ] = self.parameters.max_tokens
                    if self.parameters.temperature is not None:
                        config["temperature"] = self.parameters.temperature
                    if self.parameters.top_p is not None:
                        config["top_p"] = self.parameters.top_p
                    if self.parameters.thinking_enable:
                        config["thinking_config"] = {
                            "include_thoughts": True,
                            "thinking_budget": (
                                self.parameters.thinking_budget or 1024
                            ),
                        }

                    fmt_tools, fmt_tc = self._format_tools(
                        tools,
                        tool_choice,
                    )
                    if fmt_tools is not None:
                        config["tools"] = fmt_tools
                    if fmt_tc is not None:
                        config["tool_config"] = fmt_tc

                    call_kwargs = {
                        "model": model_name,
                        "contents": formatted,
                        "config": config,
                    }
                    start = datetime.now()
                    if self.stream:
                        response = (
                            await client.aio.models.generate_content_stream(
                                **call_kwargs,
                            )
                        )
                        return self._parse_stream_response(
                            start,
                            response,
                            client,
                        )
                    response = await client.aio.models.generate_content(
                        **call_kwargs,
                    )
                    return self._parse_completion_response(start, response)

                return await super()._call_api(
                    model_name,
                    messages,
                    tools,
                    tool_choice,
                    **merged,
                )

        return _Compat(**kwargs)
