"""Google Gemini 适配器（ARCH §4.2/§4.4）。

- 能力：原生结构化输出（`response_mime_type=application/json` + `response_schema`）、思考。
  前缀缓存走 Gemini 隐式/显式缓存，机制差异大，原型保守标 `prefix_cache=False`。
- 客户端注入（`GeminiClient` Protocol，`google.genai.Client().aio` 满足）——测试注入替身，
  绝不联网；真实 SDK（`google-genai`）仅在 apps/api 构造客户端时需要，本模块不硬 import。
- 瞬时故障翻译为 `TransientProviderError`，交网关退避/回退。
"""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any, Protocol

from ..errors import TransientProviderError
from ..types import LlmRequest
from .base import (
    Capabilities,
    ProviderResult,
    ProviderUsage,
    StreamChunk,
    is_transient_by_name,
)

_TRANSIENT_NAMES = frozenset(
    {
        "ResourceExhausted",
        "ServiceUnavailable",
        "DeadlineExceeded",
        "InternalServerError",
        "ServerError",
        "APIError",
    }
)


class GeminiModels(Protocol):
    async def generate_content(self, **kwargs: Any) -> Any: ...

    async def generate_content_stream(self, **kwargs: Any) -> Any: ...


class GeminiAio(Protocol):
    models: GeminiModels


class GeminiClient(Protocol):
    """`google.genai.Client()` 满足此 Protocol（用其 `.aio.models`）。"""

    aio: GeminiAio


def _is_transient(exc: Exception) -> bool:
    # Gemini SDK 部分错误用 `exc.code` 携带状态码 → extra_codes=True 一并检查。
    return is_transient_by_name(exc, _TRANSIENT_NAMES, extra_codes=True)


def _contents(req: LlmRequest) -> str:
    if isinstance(req.input, str):
        return req.input
    return "\n\n".join(b.text for b in req.input)


def _system_text(req: LlmRequest) -> str:
    return "\n\n".join(b.text for b in req.system)


def _usage_from(meta: Any) -> ProviderUsage:
    if meta is None:
        return ProviderUsage(input_tokens=0, output_tokens=0)
    return ProviderUsage(
        input_tokens=int(getattr(meta, "prompt_token_count", 0) or 0),
        output_tokens=int(getattr(meta, "candidates_token_count", 0) or 0),
        cache_read_tokens=int(getattr(meta, "cached_content_token_count", 0) or 0),
    )


def _config(req: LlmRequest) -> dict[str, Any]:
    cfg: dict[str, Any] = {}
    system = _system_text(req)
    if system:
        cfg["system_instruction"] = system
    if req.max_tokens is not None:
        cfg["max_output_tokens"] = req.max_tokens
    if req.output_schema is not None:
        cfg["response_mime_type"] = "application/json"
        cfg["response_schema"] = req.output_schema
    return cfg


class GeminiAdapter:
    def __init__(self, provider: str, client: GeminiClient) -> None:
        self.provider = provider
        self._client = client

    def capabilities(self) -> Capabilities:
        return Capabilities(structured_output=True, prefix_cache=False, thinking=True)

    async def complete(self, req: LlmRequest, model: str) -> ProviderResult:
        try:
            resp = await self._client.aio.models.generate_content(
                model=model, contents=_contents(req), config=_config(req)
            )
        except Exception as exc:
            if _is_transient(exc):
                raise TransientProviderError(str(exc), provider=self.provider) from exc
            raise
        text = getattr(resp, "text", "") or ""
        usage = _usage_from(getattr(resp, "usage_metadata", None))
        if req.output_schema is not None:
            parsed = req.output_schema.model_validate_json(text)
            return ProviderResult(text=text, usage=usage, parsed=parsed)
        return ProviderResult(text=text, usage=usage)

    async def stream(self, req: LlmRequest, model: str) -> AsyncIterator[StreamChunk]:
        try:
            stream = await self._client.aio.models.generate_content_stream(
                model=model, contents=_contents(req), config=_config(req)
            )
            async for chunk in stream:
                text = getattr(chunk, "text", "") or ""
                if text:
                    yield StreamChunk(text=text)
                meta = getattr(chunk, "usage_metadata", None)
                if meta is not None:
                    yield StreamChunk(usage=_usage_from(meta))
        except Exception as exc:
            if _is_transient(exc):
                raise TransientProviderError(str(exc), provider=self.provider) from exc
            raise
