"""Gemini Provider - 原生 Google Gemini API 支持"""

import json
import logging
import time

import httpx

from agentkit.core.exceptions import LLMProviderError
from agentkit.llm.protocol import (
    LLMProvider,
    LLMRequest,
    LLMResponse,
    StreamChunk,
    TokenUsage,
    ToolCall,
)
from agentkit.llm.retry import (
    CircuitBreaker,
    CircuitBreakerConfig,
    RetryConfig,
    RetryPolicy,
)

logger = logging.getLogger(__name__)


class _GeminiStreamContext:
    """Wraps an httpx streaming response context manager for use with retry/circuit breaker."""

    def __init__(self, response_ctx, response):
        self._response_ctx = response_ctx
        self._response = response

    async def __aenter__(self):
        return self._response

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return await self._response_ctx.__aexit__(exc_type, exc_val, exc_tb)


class GeminiProvider(LLMProvider):
    """Google Gemini API 原生 Provider"""

    def __init__(
        self,
        api_key: str,
        model: str = "gemini-2.0-flash",
        max_output_tokens: int = 4096,
        base_url: str = "https://generativelanguage.googleapis.com",
        timeout: float = 120.0,
        safety_settings: list | None = None,
        retry_config: RetryConfig | None = None,
        circuit_breaker_config: CircuitBreakerConfig | None = None,
        max_connections: int = 100,
        max_keepalive_connections: int = 20,
        keepalive_expiry: float = 30.0,
    ):
        self._api_key = api_key
        self._model = model
        self._max_output_tokens = max_output_tokens
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._safety_settings = safety_settings
        self._limits = httpx.Limits(
            max_connections=max_connections,
            max_keepalive_connections=max_keepalive_connections,
            keepalive_expiry=keepalive_expiry,
        )
        self._client: httpx.AsyncClient | None = None
        self._retry_policy = RetryPolicy(retry_config) if retry_config else None
        self._circuit_breaker = (
            CircuitBreaker(circuit_breaker_config, provider="gemini")
            if circuit_breaker_config
            else None
        )

    def _get_client(self) -> httpx.AsyncClient:
        """Lazy client initialization"""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=self._timeout, limits=self._limits)
        return self._client

    async def close(self) -> None:
        """关闭 HTTP 客户端连接池"""
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    def _convert_messages(
        self, messages: list[dict[str, str]]
    ) -> tuple[dict[str, object] | None, list[dict[str, object]]]:
        """将 OpenAI 风格消息转换为 Gemini 格式

        Returns:
            (system_instruction, contents)
        """
        system_instruction: dict[str, object] | None = None
        contents: list[dict[str, object]] = []

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

            if role == "system":
                system_instruction = {"parts": [{"text": content}]}
                continue

            if role == "user":
                # Check if this is a tool result message
                if msg.get("tool_call_id"):
                    # Tool response: role="user" with functionResponse part
                    tool_name = msg.get("name", "")
                    # If name not at top level, try to extract from content
                    if not tool_name and isinstance(content, str):
                        try:
                            parsed = json.loads(content)
                            tool_name = parsed.get("name", "")
                        except (json.JSONDecodeError, AttributeError):
                            pass
                    contents.append(
                        {
                            "role": "user",
                            "parts": [
                                {
                                    "functionResponse": {
                                        "name": tool_name,
                                        "response": {
                                            "content": content,
                                        },
                                    },
                                }
                            ],
                        }
                    )
                else:
                    contents.append(
                        {
                            "role": "user",
                            "parts": [{"text": content}],
                        }
                    )
                continue

            if role == "assistant":
                tool_calls = msg.get("tool_calls")
                if tool_calls:
                    parts: list[dict[str, object]] = []
                    if content:
                        parts.append({"text": content})
                    for tc in tool_calls:
                        func = tc.get("function", {})
                        arguments = func.get("arguments", "{}")
                        if isinstance(arguments, str):
                            try:
                                arguments = json.loads(arguments)
                            except json.JSONDecodeError:
                                arguments = {"raw": arguments}
                        parts.append(
                            {
                                "functionCall": {
                                    "name": func.get("name", ""),
                                    "args": arguments,
                                },
                            }
                        )
                    contents.append({"role": "model", "parts": parts})
                else:
                    contents.append(
                        {
                            "role": "model",
                            "parts": [{"text": content}],
                        }
                    )
                continue

            if role == "tool":
                # OpenAI format: {"role": "tool", "tool_call_id": "...", "content": "..."}
                tool_name = msg.get("name", "")
                tool_content = msg.get("content", "")
                contents.append(
                    {
                        "role": "user",
                        "parts": [
                            {
                                "functionResponse": {
                                    "name": tool_name,
                                    "response": {
                                        "content": tool_content,
                                    },
                                },
                            }
                        ],
                    }
                )

        return system_instruction, contents

    def _convert_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]:
        """将 OpenAI function 格式转换为 Gemini functionDeclarations"""
        declarations = []
        for tool in tools:
            if tool.get("type") == "function":
                func = tool.get("function", {})
                declarations.append(
                    {
                        "name": func.get("name", ""),
                        "description": func.get("description", ""),
                        "parameters": func.get("parameters", {"type": "object", "properties": {}}),
                    }
                )
        if not declarations:
            return []
        return [{"functionDeclarations": declarations}]

    def _convert_tool_choice(self, tool_choice: str) -> dict[str, object] | None:
        """将 OpenAI tool_choice 格式转换为 Gemini toolConfig"""
        if tool_choice == "auto":
            return {"functionCallingConfig": {"mode": "AUTO"}}
        elif tool_choice == "required":
            return {"functionCallingConfig": {"mode": "ANY"}}
        elif tool_choice and tool_choice not in ("none",):
            return {"functionCallingConfig": {"mode": "AUTO"}}
        if tool_choice == "none":
            return {"functionCallingConfig": {"mode": "NONE"}}
        return None

    def _parse_response(self, data: dict[str, object], model: str) -> LLMResponse:
        """将 Gemini 响应转换为 LLMResponse"""
        candidates = data.get("candidates", [])
        text_parts: list[str] = []
        tool_calls: list[ToolCall] = []
        tool_call_index = 0

        if candidates:
            content = candidates[0].get("content", {})
            parts = content.get("parts", [])
            for part in parts:
                if "text" in part:
                    text_parts.append(part["text"])
                elif "functionCall" in part:
                    fc = part["functionCall"]
                    tool_calls.append(
                        ToolCall(
                            id=f"call_{tool_call_index}",
                            name=fc.get("name", ""),
                            arguments=fc.get("args", {}),
                        )
                    )
                    tool_call_index += 1

        usage_metadata = data.get("usageMetadata", {})
        usage = TokenUsage(
            prompt_tokens=usage_metadata.get("promptTokenCount", 0),
            completion_tokens=usage_metadata.get("candidatesTokenCount", 0),
        )

        return LLMResponse(
            content="".join(text_parts),
            model=data.get("modelVersion", model),
            usage=usage,
            tool_calls=tool_calls,
        )

    def _handle_error(self, status_code: int, resp_body: bytes) -> None:
        """处理 Gemini API 错误响应"""
        try:
            error_data = json.loads(resp_body)
            error_info = error_data.get("error", {})
            error_msg = error_info.get("message", f"HTTP {status_code}")
        except (json.JSONDecodeError, AttributeError):
            error_msg = f"HTTP {status_code}"

        raise LLMProviderError("gemini", f"HTTP {status_code}: {error_msg}")

    async def chat(self, request: LLMRequest) -> LLMResponse:
        """发送 chat 请求（带 retry + circuit breaker）"""
        if self._circuit_breaker and self._retry_policy:
            return await self._circuit_breaker.execute(
                self._retry_policy.execute, self._chat_impl, request
            )
        if self._retry_policy:
            return await self._retry_policy.execute(self._chat_impl, request)
        if self._circuit_breaker:
            return await self._circuit_breaker.execute(self._chat_impl, request)
        return await self._chat_impl(request)

    async def _chat_impl(self, request: LLMRequest) -> LLMResponse:
        client = self._get_client()
        model = request.model or self._model
        url = f"{self._base_url}/v1beta/models/{model}:generateContent?key={self._api_key}"

        system_instruction, contents = self._convert_messages(request.messages)

        payload: dict[str, object] = {
            "contents": contents,
            "generationConfig": {
                "temperature": request.temperature,
                "maxOutputTokens": request.max_tokens or self._max_output_tokens,
            },
        }

        if system_instruction is not None:
            payload["systemInstruction"] = system_instruction

        if request.tools:
            gemini_tools = self._convert_tools(request.tools)
            if gemini_tools:
                payload["tools"] = gemini_tools
            tool_config = self._convert_tool_choice(request.tool_choice)
            if tool_config is not None:
                payload["toolConfig"] = tool_config

        if self._safety_settings:
            payload["safetySettings"] = self._safety_settings

        start = time.monotonic()

        try:
            resp = await client.post(url, json=payload)
        except httpx.HTTPError as e:
            raise LLMProviderError("gemini", str(e)) from e

        latency_ms = (time.monotonic() - start) * 1000

        if resp.status_code != 200:
            self._handle_error(resp.status_code, resp.content)

        data = resp.json()
        response = self._parse_response(data, model)
        response.latency_ms = latency_ms

        return response

    async def chat_stream(self, request: LLMRequest):
        """Stream chat response using SSE（带 retry + circuit breaker）"""
        if self._circuit_breaker and self._retry_policy:
            ctx = await self._circuit_breaker.execute(
                self._retry_policy.execute, self._open_stream, request
            )
        elif self._retry_policy:
            ctx = await self._retry_policy.execute(self._open_stream, request)
        elif self._circuit_breaker:
            ctx = await self._circuit_breaker.execute(self._open_stream, request)
        else:
            ctx = await self._open_stream(request)

        async with ctx as response:
            async for chunk in self._iterate_stream(response, request):
                yield chunk

    async def _open_stream(self, request: LLMRequest):
        """Open the streaming HTTP connection; returns an async context manager."""
        client = self._get_client()
        model = request.model or self._model
        url = f"{self._base_url}/v1beta/models/{model}:streamGenerateContent?key={self._api_key}&alt=sse"

        system_instruction, contents = self._convert_messages(request.messages)

        payload: dict[str, object] = {
            "contents": contents,
            "generationConfig": {
                "temperature": request.temperature,
                "maxOutputTokens": request.max_tokens or self._max_output_tokens,
            },
        }

        if system_instruction is not None:
            payload["systemInstruction"] = system_instruction

        if request.tools:
            gemini_tools = self._convert_tools(request.tools)
            if gemini_tools:
                payload["tools"] = gemini_tools
            tool_config = self._convert_tool_choice(request.tool_choice)
            if tool_config is not None:
                payload["toolConfig"] = tool_config

        if self._safety_settings:
            payload["safetySettings"] = self._safety_settings

        response_ctx = client.stream("POST", url, json=payload)
        response = await response_ctx.__aenter__()

        if response.status_code != 200:
            error_body = await response.aread()
            await response_ctx.__aexit__(None, None, None)
            self._handle_error(response.status_code, error_body)

        return _GeminiStreamContext(response_ctx, response)

    async def _iterate_stream(self, response, request: LLMRequest):
        """Iterate over an already-open SSE stream and yield StreamChunks."""
        accumulated_tool_calls: list[dict[str, object]] = []
        model = request.model or self._model

        async for line in response.aiter_lines():
            line = line.strip()
            if not line or not line.startswith("data: "):
                continue

            data_str = line[6:]
            try:
                data = json.loads(data_str)
            except json.JSONDecodeError:
                continue

            candidates = data.get("candidates", [])
            if not candidates:
                # Usage-only chunk
                usage_metadata = data.get("usageMetadata")
                if usage_metadata:
                    usage = TokenUsage(
                        prompt_tokens=usage_metadata.get("promptTokenCount", 0),
                        completion_tokens=usage_metadata.get("candidatesTokenCount", 0),
                    )
                    if accumulated_tool_calls:
                        tool_calls = [
                            ToolCall(
                                id=tc["id"],
                                name=tc["name"],
                                arguments=tc["arguments"],
                            )
                            for tc in accumulated_tool_calls
                        ]
                        yield StreamChunk(
                            content="",
                            model=data.get("modelVersion", model),
                            tool_calls=tool_calls,
                            usage=usage,
                            is_final=True,
                        )
                        accumulated_tool_calls = []
                    else:
                        yield StreamChunk(
                            content="",
                            model=data.get("modelVersion", model),
                            usage=usage,
                            is_final=True,
                        )
                continue

            content = candidates[0].get("content", {})
            parts = content.get("parts", [])

            for part in parts:
                if "text" in part:
                    text = part["text"]
                    if text:
                        yield StreamChunk(
                            content=text,
                            model=data.get("modelVersion", model),
                        )
                elif "functionCall" in part:
                    fc = part["functionCall"]
                    accumulated_tool_calls.append(
                        {
                            "id": f"call_{len(accumulated_tool_calls)}",
                            "name": fc.get("name", ""),
                            "arguments": fc.get("args", {}),
                        }
                    )

            # Check for finish reason
            finish_reason = candidates[0].get("finishReason", "")
            if finish_reason in ("STOP", "MAX_TOKENS") and accumulated_tool_calls:
                tool_calls = [
                    ToolCall(
                        id=tc["id"],
                        name=tc["name"],
                        arguments=tc["arguments"],
                    )
                    for tc in accumulated_tool_calls
                ]
                yield StreamChunk(
                    content="",
                    model=data.get("modelVersion", model),
                    tool_calls=tool_calls,
                    is_final=True,
                )
                accumulated_tool_calls = []

    def get_model_info(self) -> dict[str, object]:
        """返回 Provider 和模型信息"""
        return {
            "provider": "gemini",
            "model": self._model,
            "max_output_tokens": self._max_output_tokens,
        }