"""
MCP (Model Context Protocol) 客户端 —— 基于官方 \"mcp\" Python SDK (Phase 1.E 标准化)

设计基准（标准对照表见 STANDARDS_AUDIT_2026-07-24.md §E）：
- 采用 Anthropic 官方 \"mcp\" SDK 处理 JSON-RPC 帧 + 协议版本协商，
  不再手写解析、不再写死 MCP_PROTOCOL_VERSION 常量（旧实现写死 "2024-11-05"，
  上游协议一升级即碎）。
- 支持 streamable_http（默认）与 sse 两种传输，按 server['transport_type'] 选择。
- 公共 API 保持向后兼容：list_mcp_tools(server, timeout) / call_mcp_tool(server, tool_name, arguments, timeout)
  供 chat.py / chat_service.py / chat_agents.py / chat_skill.py / chat_mcp.py /
  chat_context.py / chat_knowledge.py / chat_media.py / chat_utils.py / core/aca.py
  共 10+ 处调用，签名与返回结构（tools 数组 / 文本结果字符串）不变。
- 工具列表缓存（按 url）与错误处理保持原行为，避免重复握手、保证降级安全。
"""

import asyncio
import logging
from typing import Optional

from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamablehttp_client

logger = logging.getLogger(__name__)

# 工具列表缓存（url -> (timestamp, tools)），避免同一请求内 / 短时间重复握手
_TOOLS_CACHE: dict = {}
_TOOLS_CACHE_TTL = 120  # 秒


def _transport_cm(server: dict, timeout: float):
    """按 server['transport_type'] 选择官方 SDK 传输层。

    返回异步上下文管理器，进入后产出 (read_stream, write_stream, ...)。
    """
    url = (server.get("url") or "").strip()
    ttype = (server.get("transport_type") or "streamable_http").lower()
    headers = server.get("headers") or {}
    if ttype == "sse":
        return sse_client(url, headers=headers, timeout=timeout)
    # 默认 streamable_http（官方 SDK 标准传输）
    return streamablehttp_client(url, headers=headers, timeout=timeout)


async def _session_request(
    server: dict, method: str, params: dict, timeout: float = 30.0,
    raise_on_error: bool = False
):
    """用官方 ClientSession 完成一次标准 MCP 调用（initialize -> 目标请求）。

    返回解析后的 dict：{"result": {...}} 或 None（失败）。
    协议版本由 SDK 自动协商，不再硬编码。
    raise_on_error=True 时连接/网络错误（如 ExceptionGroup）会向上抛，不再吞掉。
    """
    url = (server.get("url") or "").strip()
    if not url:
        return None
    try:
        async with _transport_cm(server, timeout) as streams:
            read, write, *_ = streams
            async with ClientSession(read, write) as session:
                # 标准握手：initialize 会发送 initialize 请求并自动补 notifications/initialized
                await session.initialize()
                # 防御性补发 initialized 通知（部分服务端要求显式接收）
                try:
                    if hasattr(session, "send_notification"):
                        await session.send_notification("notifications/initialized", {})
                except Exception:
                    pass

                if method == "tools/list":
                    result = await session.list_tools()
                    tools = [t.model_dump() for t in (result.tools or [])]
                    return {"result": {"tools": tools}}
                elif method == "tools/call":
                    res = await session.call_tool(
                        params.get("name"), params.get("arguments") or {}
                    )
                    content = []
                    for blk in (res.content or []):
                        btype = getattr(blk, "type", "text")
                        if btype == "text":
                            content.append({"type": "text", "text": getattr(blk, "text", "")})
                        else:
                            content.append({"type": btype, "text": str(blk)})
                    return {
                        "result": {
                            "content": content,
                            "isError": bool(getattr(res, "isError", False)),
                        }
                    }
    except Exception as e:
        logger.warning(f"[MCP] {method} 异常 {url}: {e}")
        if raise_on_error:
            raise
        return None
    return None


async def list_mcp_tools(server: dict, timeout: float = 30.0, raise_on_error: bool = False) -> list:
    """
    拉取 MCP 服务器的工具列表。
    server: {url, transport_type, headers}
    返回 MCP tools 数组（每个含 name/description/inputSchema）。
    raise_on_error=True 时连接/网络错误不再吞掉，直接向上抛（供测试诊断用）。
    """
    url = (server.get("url") or "").strip()
    if not url:
        return []
    # 缓存（仅在正常返回 tools 时缓存；raise_on_error 模式不缓存，便于实时重试）
    if not raise_on_error:
        now = asyncio.get_event_loop().time()
        cached = _TOOLS_CACHE.get(url)
        if cached and (now - cached[0]) < _TOOLS_CACHE_TTL:
            return cached[1]
    try:
        data = await _session_request(server, "tools/list", {}, timeout, raise_on_error=raise_on_error)
        tools = []
        if data and "result" in data:
            tools = data["result"].get("tools", []) or []
        if not raise_on_error:
            _TOOLS_CACHE[url] = (asyncio.get_event_loop().time(), tools)
        return tools
    except Exception as e:
        logger.warning(f"[MCP] 拉取工具失败 {url}: {e}")
        if raise_on_error:
            raise
        return []


async def call_mcp_tool(
    server: dict, tool_name: str, arguments: dict, timeout: float = 60.0
) -> str:
    """
    在 MCP 服务器上执行一个工具，返回可供 LLM 消费的文本结果。
    """
    url = (server.get("url") or "").strip()
    if not url:
        return f"[MCP 工具 {tool_name} 执行失败：服务端地址为空]"
    logger.info(f"[MCP] 执行工具调用 tool={tool_name} server={server.get('name')} url={url}")
    try:
        data = await _session_request(
            server,
            "tools/call",
            {"name": tool_name, "arguments": arguments or {}},
            timeout,
        )
        if not data:
            return f"[MCP 工具 {tool_name} 调用失败：无响应]"
        if "error" in data:
            err = data["error"]
            return f"[MCP 工具 {tool_name} 返回错误] {__import__('json').dumps(err, ensure_ascii=False)}"
        result = data.get("result", {})
        content = result.get("content", []) or []
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
            else:
                parts.append(__import__("json").dumps(item, ensure_ascii=False))
        text = "\n".join(parts).strip()
        if result.get("isError"):
            return f"[MCP 工具 {tool_name} 执行报错]\n{text}"
        return text or f"[MCP 工具 {tool_name} 返回空结果]"
    except Exception as e:
        logger.warning(f"[MCP] 工具执行异常 {tool_name} @ {url}: {e}")
        return f"[MCP 工具 {tool_name} 执行异常：{e}]"