"""Google provider client (Gemini).

基于 google-genai Python SDK 实现，遵循 gemini.md 示例：
- 文本生成: client.models.generate_content(_stream)
- 工具调用: Tool(function_declarations=[FunctionDeclaration(...)]),
  工具结果通过 Part(function_response=FunctionResponse(...)) 回传。

多轮工具调用的关键约束：
1. Gemini 要求 user/model 角色严格交替，不能有连续同 role 消息
2. model 的工具调用响应必须包含 Part(function_call=...) 部件
3. 工具结果以 function_response 发送，角色为 user
"""

from __future__ import annotations

from typing import Any, AsyncGenerator, Dict, List, Optional, Set

from core.logger import log_error, log_llm_request, log_llm_response, log_token_usage
from .base import (
    BaseModelClient,
    CompletionResponse,
    Message,
    MessageRole,
    StreamChunk,
    ToolCall,
    ToolDefinition,
    TokenUsage,
)


# ============================================================================
# Gemini Schema 兼容性处理
#
# Gemini FunctionDeclaration.parameters 接受类似 OpenAPI Schema 的对象，
# 但比 OpenAI 更严格，会拒绝 pydantic v2 自动生成的多种字段。
#
# 已知不兼容字段及处理策略：
# - title / default / $defs / $ref / $schema / examples → 直接移除
# - additional_properties / additionalProperties → 直接移除（Gemini 不支持）
# - anyOf: [{type: X}, {type: "null"}] → 简化为 {type: X}（Optional 语义）
#
# 注意：Gemini SDK 内部使用 pydantic，传入 dict 会被转为 Schema 对象。
# 如果 dict 中有 additionalProperties，SDK 会转为 additional_properties（snake_case），
# 然后 API 返回 400 "Unknown name additional_properties"。所以必须在传入前移除。
# ============================================================================

# Gemini function declaration 不支持的 JSON Schema 关键字
_UNSUPPORTED_KEYS: Set[str] = {
    "title",
    "default",
    "$defs",
    "$ref",
    "$schema",
    "examples",
    "additional_properties",   # pydantic v2 snake_case 输出
    "additionalProperties",    # 标准 JSON Schema，但 Gemini 不支持
}


def _normalize_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
    """标准化 JSON Schema，使其能被 Gemini function_declarations 接受。

    处理：
    1. 移除不支持的关键字
    2. 解析 $ref → $defs 引用
    3. 简化 anyOf Optional 模式：anyOf[{type:X}, {type:null}] → {type:X}
    4. 递归处理所有嵌套结构
    """
    if not schema:
        return {}
    defs = schema.get("$defs", schema.get("defs", {}))
    return _walk(schema, defs)


def _simplify_any_of(any_of: List[Any], defs: Dict[str, Any]) -> Dict[str, Any]:
    """简化 pydantic 生成的 anyOf Optional 模式。

    pydantic 将 Optional[str] 编码为：
        {"anyOf": [{"type": "string"}, {"type": "null"}]}
    Gemini 不一定支持 null 类型，简化为非 null 的那个类型。

    如果 anyOf 中有多个非 null 类型，保留 anyOf 结构。
    """
    non_null = []
    for item in any_of:
        walked = _walk(item, defs)
        if isinstance(walked, dict) and walked.get("type") == "null":
            continue
        non_null.append(walked)

    if len(non_null) == 1:
        res = non_null[0]
        return res if isinstance(res, dict) else {"type": "string"}
    elif len(non_null) == 0:
        return {"type": "string"}
    else:
        return {"anyOf": non_null}


def _walk(node: Any, defs: Dict[str, Any]) -> Any:
    """递归处理 schema 节点"""
    if isinstance(node, list):
        return [_walk(x, defs) for x in node]
    if not isinstance(node, dict):
        return node

    # 处理 $ref 引用
    ref = node.get("$ref")
    if ref and isinstance(ref, str) and ref.startswith("#/$defs/"):
        ref_name = ref[len("#/$defs/"):]
        if ref_name in defs:
            return _walk(defs[ref_name], defs)
        return {"type": "object"}

    # 处理 anyOf（在构建 out 之前，因为可能要替换整个节点）
    if "anyOf" in node:
        simplified = _simplify_any_of(node["anyOf"], defs)
        if isinstance(simplified, dict):
            for k, v in node.items():
                if k in _UNSUPPORTED_KEYS or k == "anyOf":
                    continue
                if k not in simplified:
                    simplified[k] = _walk(v, defs)
        return simplified

    out: Dict[str, Any] = {}
    for k, v in node.items():
        if k in _UNSUPPORTED_KEYS:
            continue
        out[k] = _walk(v, defs)

    return out


class GoogleClient(BaseModelClient):
    """Google Gemini API 客户端"""

    def __init__(
        self,
        model: str,
        api_key: str,
        base_url: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(model, api_key, base_url, **kwargs)

        try:
            from google import genai  # type: ignore
            from google.genai import types  # type: ignore
        except Exception as e:  # pragma: no cover
            raise ImportError(
                "缺少依赖: google-genai。安装: pip install google-genai"
            ) from e

        self._types = types

        # 创建 Gemini 客户端（异步版本）
        # SDK 也支持环境变量: GEMINI_API_KEY / GOOGLE_API_KEY
        if api_key:
            self._client = genai.Client(api_key=api_key).aio
        else:
            self._client = genai.Client().aio

    @property
    def provider(self) -> str:
        return "google"

    @staticmethod
    def _content_to_log(content: Any) -> Dict[str, Any]:
        """将 types.Content 对象转换为可 JSON 序列化的 dict（用于日志）"""
        role = getattr(content, "role", "?")
        parts_log = []
        for part in (getattr(content, "parts", None) or []):
            if getattr(part, "text", None):
                text = part.text
                parts_log.append({"text": text[:200] + "..." if len(text) > 200 else text})
            elif getattr(part, "function_call", None):
                fc = part.function_call
                parts_log.append({"function_call": {"name": getattr(fc, "name", ""), "args": dict(getattr(fc, "args", None) or {})}})
            elif getattr(part, "function_response", None):
                fr = part.function_response
                resp = getattr(fr, "response", {})
                content_str = str(resp.get("content", ""))
                parts_log.append({"function_response": {"name": getattr(fr, "name", ""), "content": content_str[:200] + "..." if len(content_str) > 200 else content_str}})
            else:
                parts_log.append({"type": "other"})
        return {"role": role, "parts": parts_log}

    def _merge_system_prompts(self, system_prompts: Optional[List[Message]]) -> str:
        """合并系统提示词为单个字符串"""
        if not system_prompts:
            return ""
        parts = [m.content for m in system_prompts if m.content]
        return "\n\n".join(parts).strip()

    def _to_function_declarations(self, tools: List[ToolDefinition]) -> List[Any]:
        """将 ToolDefinition 转换为 Gemini FunctionDeclaration"""
        decls = []
        for t in tools:
            decls.append(
                self._types.FunctionDeclaration(
                    name=t.name,
                    description=t.description,
                    parameters=_normalize_schema(t.parameters),
                )
            )
        return decls

    def _to_gemini_contents(
        self,
        messages: List[Message],
    ) -> List[Any]:
        """将统一消息格式转换为 Gemini 多轮对话格式。

        返回 types.Content 对象列表，SDK 要求使用类型化对象而非原始 dict。

        关键处理：
        1. assistant + tool_calls → model 消息包含 functionCall parts
        2. tool result → user 消息包含 functionResponse parts
        3. 合并连续同 role 消息（Gemini 要求 user/model 严格交替）

        注意：系统提示词通过 GenerateContentConfig.system_instruction 传递，
        不再嵌入到 user 消息中。
        """
        types = self._types

        # 逐条转换消息为 (role, parts) 元组
        raw_entries: List[tuple] = []  # [(role, [Part, ...])]

        for msg in messages:
            # 跳过系统角色消息（通过 system_instruction 传递）
            if msg.role in (
                MessageRole.SYSTEM,
                MessageRole.AGENT,
                MessageRole.ENVIRONMENT,
                MessageRole.CUSTOM,
            ):
                continue

            if msg.role == MessageRole.USER:
                text = msg.content or ""
                if text:
                    raw_entries.append(("user", [types.Part(text=text)]))
                continue

            if msg.role == MessageRole.ASSISTANT:
                parts: List[Any] = []

                # 添加文本部分
                if msg.content:
                    parts.append(types.Part(text=msg.content))

                # 添加 function_call 部分（多轮工具调用的关键）
                for tc in (msg.tool_calls or []):
                    fc_part_kwargs: Dict[str, Any] = {
                        "function_call": types.FunctionCall(
                            name=tc.name,
                            args=tc.arguments,
                        )
                    }
                    # 恢复 thought_signature（Gemini thinking 模型必需）
                    thought_sig_b64 = (tc.metadata or {}).get("thought_signature")
                    if thought_sig_b64:
                        import base64
                        fc_part_kwargs["thought_signature"] = base64.b64decode(thought_sig_b64)
                    parts.append(types.Part(**fc_part_kwargs))

                if parts:
                    raw_entries.append(("model", parts))
                continue

            if msg.role == MessageRole.TOOL and msg.tool_result:
                # 工具结果以 functionResponse 发送
                fr = types.FunctionResponse(
                    name=msg.tool_result.name,
                    response={
                        "content": msg.tool_result.content,
                        "is_error": bool(msg.tool_result.is_error),
                    },
                )
                fr_part = types.Part(function_response=fr)
                raw_entries.append(("user", [fr_part]))
                continue

        # 合并连续同 role 消息（Gemini 要求 user/model 严格交替）
        merged: List[tuple] = []
        for role, parts in raw_entries:
            if merged and merged[-1][0] == role:
                merged[-1][1].extend(parts)
            else:
                merged.append((role, parts))

        # 构建 types.Content 对象列表
        return [
            types.Content(role=role, parts=parts)
            for role, parts in merged
        ]

    def _extract_tool_calls(self, response: Any) -> List[ToolCall]:
        """从 Gemini 响应中提取工具调用

        对于 Gemini 3.x thinking 模型，Part 上会携带 thought_signature，
        必须在后续请求中原样带回，否则 API 返回 400。
        这里将 thought_signature (bytes) base64 编码后存入 ToolCall.metadata。
        """
        import base64
        tool_calls: List[ToolCall] = []
        try:
            candidates = getattr(response, "candidates", None) or []
            call_idx = 0
            for cand in candidates:
                content = getattr(cand, "content", None)
                parts = getattr(content, "parts", None) or []
                for part in parts:
                    fc = getattr(part, "function_call", None)
                    if not fc:
                        continue
                    call_idx += 1
                    args = getattr(fc, "args", None) or {}
                    metadata: Dict[str, Any] = {}
                    # 保存 thought_signature（Gemini thinking 模型必需）
                    thought_sig = getattr(part, "thought_signature", None)
                    if thought_sig is not None:
                        if isinstance(thought_sig, bytes):
                            metadata["thought_signature"] = base64.b64encode(thought_sig).decode("ascii")
                        else:
                            metadata["thought_signature"] = str(thought_sig)
                    tool_calls.append(
                        ToolCall(
                            id=f"google_fc_{call_idx}",
                            name=getattr(fc, "name", ""),
                            arguments=dict(args),
                            metadata=metadata,
                        )
                    )
        except Exception as e:
            log_error(f"google_extract_tool_calls_error: {e}")
        return tool_calls

    def _extract_usage(self, chunk: Any) -> Optional[TokenUsage]:
        """从 Gemini chunk 中提取 token 使用量"""
        usage = getattr(chunk, "usage_metadata", None)
        if not usage:
            return None
            
        prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0
        candidates_tokens = getattr(usage, "candidates_token_count", 0) or 0
        cached_tokens = getattr(usage, "cached_content_token_count", 0) or 0
        
        # 为了与 TokenUsage 统一语义（input_tokens 为非缓存输入），需要减去 cached_tokens
        # Gemini 的 prompt_token_count 包含了 cached_content_token_count
        input_tokens = max(0, prompt_tokens - cached_tokens)
        
        return TokenUsage(
            input_tokens=input_tokens,
            output_tokens=candidates_tokens,
            cache_read_input_tokens=cached_tokens,
        )

    async def chat_stream(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs,
    ) -> AsyncGenerator[StreamChunk, None]:
        """流式调用 Gemini API"""
        effective_temperature = temperature if temperature is not None else self.temperature
        effective_max_tokens = max_tokens if max_tokens is not None else self.max_tokens

        contents = self._to_gemini_contents(messages)

        req: Dict[str, Any] = {
            "model": self.model,
            "contents": contents,
        }

        # 构建生成配置
        config_kwargs: Dict[str, Any] = {"temperature": effective_temperature}
        if effective_max_tokens:
            config_kwargs["max_output_tokens"] = effective_max_tokens
        # 系统提示词通过原生 system_instruction 传递（比嵌入 user 消息权重更高）
        system_text = self._merge_system_prompts(system_prompts)
        if system_text:
            config_kwargs["system_instruction"] = system_text
        if tools:
            decls = self._to_function_declarations(tools)
            config_kwargs["tools"] = [self._types.Tool(function_declarations=decls)]
        req["config"] = self._types.GenerateContentConfig(**config_kwargs)

        # 构建日志用请求体（与 OpenAI/Anthropic 格式对齐）
        request_log: Dict[str, Any] = {
            "model": self.model,
            "messages": [self._content_to_log(c) for c in contents],
            "temperature": effective_temperature,
        }
        if effective_max_tokens:
            request_log["max_output_tokens"] = effective_max_tokens
        if system_prompts:
            request_log["system"] = self._merge_system_prompts(system_prompts)
        if tools:
            request_log["tools"] = [
                {"name": t.name, "description": t.description, "parameters": t.parameters}
                for t in tools
            ]

        log_llm_request(
            provider=self.provider,
            model=self.model,
            request_body=request_log,
        )

        accumulated_text = ""
        all_tool_calls: List[ToolCall] = []
        last_usage: Optional[TokenUsage] = None

        try:
            stream = await self._client.models.generate_content_stream(**req)
            async for chunk in stream:
                # 提取文本增量
                text = getattr(chunk, "text", None)
                if text:
                    # google-genai SDK 每个 chunk.text 是增量文本
                    accumulated_text += text
                    yield StreamChunk(text=text)

                # 收集所有 chunk 中的工具调用
                tcs = self._extract_tool_calls(chunk)
                if tcs:
                    all_tool_calls.extend(tcs)

                # 提取 usage（通常在最后一个 chunk）
                usage = self._extract_usage(chunk)
                if usage:
                    last_usage = usage

            yield StreamChunk(text_complete=True)

            # 日志：记录响应（与 OpenAI/Anthropic 对齐）
            if last_usage:
                log_token_usage(
                    provider=self.provider,
                    model=self.model,
                    input_tokens=last_usage.input_tokens,
                    output_tokens=last_usage.output_tokens,
                    cache_read_input_tokens=last_usage.cache_read_input_tokens,
                )
            log_llm_response(
                provider=self.provider,
                model=self.model,
                response_body={
                    "content": accumulated_text,
                    "tool_calls": [{"name": tc.name, "arguments": tc.arguments} for tc in all_tool_calls] if all_tool_calls else None,
                    "usage": last_usage.__dict__ if last_usage else None,
                },
            )

            # 工具调用必须包含在 is_finished 的最终 chunk 中
            # base.py 只从 is_finished chunk 提取 tool_calls
            yield StreamChunk(
                is_finished=True,
                finish_reason="tool_use" if all_tool_calls else "stop",
                tool_calls=all_tool_calls,
                usage=last_usage,
            )

        except Exception as e:
            import traceback
            log_error(
                f"google_stream_error: [{type(e).__name__}] {str(e) or repr(e)}\n{traceback.format_exc()}",
                model=self.model,
            )
            raise

    async def chat(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs,
    ) -> CompletionResponse:
        """非流式调用 Gemini API"""
        content = ""
        tool_calls: List[ToolCall] = []
        usage = None
        async for chunk in self.chat_stream(
            messages=messages,
            tools=tools,
            system_prompts=system_prompts,
            temperature=temperature,
            max_tokens=max_tokens,
            **kwargs,
        ):
            if chunk.text:
                content += chunk.text
            if chunk.tool_calls:
                tool_calls.extend(chunk.tool_calls)
            if chunk.usage:
                usage = chunk.usage
        return CompletionResponse(
            content=content,
            tool_calls=tool_calls,
            finish_reason="tool_use" if tool_calls else "stop",
            usage=usage,
        )

    async def count_tokens(
        self,
        messages: List[Message],
        system_prompts: Optional[List[Message]] = None,
        tools: Optional[List[ToolDefinition]] = None,
    ) -> int:
        """
        使用 Gemini SDK 精确计算 token 数

        调用 client.models.count_tokens() API。
        失败时 fallback 到基类的字符估算。
        """
        try:
            contents = self._to_gemini_contents(messages)

            req: Dict[str, Any] = {
                "model": self.model,
                "contents": contents,
            }

            # 系统提示词和工具通过 CountTokensConfig 传递
            config_kwargs: Dict[str, Any] = {}
            system_text = self._merge_system_prompts(system_prompts)
            if system_text:
                config_kwargs["system_instruction"] = system_text

            # 工具定义
            if tools:
                decls = self._to_function_declarations(tools)
                config_kwargs["tools"] = [self._types.Tool(function_declarations=decls)]

            if config_kwargs:
                req["config"] = self._types.CountTokensConfig(**config_kwargs)

            response = await self._client.models.count_tokens(**req)
            return getattr(response, "total_tokens", 0) or 0

        except Exception as e:
            log_error(f"google_count_tokens_error: {e}")
            return await super().count_tokens(messages, system_prompts, tools)
