"""Google Gemini generateContent REST protocol."""

from typing import Any

from agenteval.models import ModelCapability, ModelProtocol

from .base import ModelProtocolAdapter, ProtocolAdapterError


class GoogleGeminiAdapter(ModelProtocolAdapter):
    protocol = ModelProtocol.GOOGLE_GEMINI
    supported_capabilities = frozenset({ModelCapability.CHAT})

    def headers(self, api_key: str | None) -> dict[str, str]:
        headers = super().headers(api_key)
        if api_key:
            headers["x-goog-api-key"] = api_key
        return headers

    def chat_payload(
        self,
        model_name: str | None,
        messages: list[dict[str, str]],
        temperature: float,
    ) -> dict[str, Any]:
        self.require_model(model_name)
        system_parts: list[dict[str, str]] = []
        contents: list[dict[str, Any]] = []
        for message in messages:
            role = message.get("role", "user")
            part = {"text": message.get("content", "")}
            if role in {"system", "developer"}:
                system_parts.append(part)
            else:
                contents.append({"role": "model" if role == "assistant" else "user", "parts": [part]})
        payload: dict[str, Any] = {
            "contents": contents,
            "generationConfig": {"temperature": temperature},
        }
        if system_parts:
            payload["systemInstruction"] = {"parts": system_parts}
        return payload

    def parse_chat(self, data: dict[str, Any]) -> str:
        try:
            parts = data["candidates"][0]["content"]["parts"]
        except (KeyError, IndexError, TypeError) as exc:
            raise ProtocolAdapterError("Gemini 接口返回格式不正确") from exc
        content = "\n".join(
            str(part.get("text", ""))
            for part in parts
            if isinstance(part, dict) and part.get("text")
        )
        if not content:
            raise ProtocolAdapterError("模型接口返回内容为空")
        return content