"""按 mcp_servers.json 连接 MCP 服务，把所有服务的工具交给 agent

每一项是一个服务：键为服务名，transport / url / headers 原样交给 MultiServerMCPClient；
填写 module 的服务只能在本机，由应用以子进程拉起，不填表示服务由外部部署
各环境不同的值和密钥写成 ${变量名}，加载时从环境变量展开
"""

import asyncio
import json
import os
import re
import socket
import subprocess
import sys
from collections import Counter
from pathlib import Path
from urllib.parse import urlsplit

from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.interceptors import MCPToolCallRequest
from loguru import logger
from mcp.types import CallToolResult

from app.config import settings
from app.resources.component import get_catalog_cache

PROJECT_ROOT = Path(__file__).resolve().parents[2]
CONFIG_PATH = PROJECT_ROOT / "mcp_servers.json"
# 物料服务返回的目录版本须与主进程一致，否则模型查到的字段和校验依据不同
MATERIAL_SERVER = "materials"
LOCAL_HOSTS = {"localhost", "127.0.0.1"}
# 配置中引用环境变量的写法
ENV_REFERENCE = re.compile(r"\$\{(\w+)\}")


class MCPServiceError(RuntimeError):
    """MCP 配置错误、服务不可用，或物料目录版本不一致"""


def _expand_env(value):
    """只展开字符串值里的 ${NAME}；变量未设置时报错，不把字面量当成地址或密钥"""
    if isinstance(value, dict):
        return {key: _expand_env(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_expand_env(item) for item in value]
    if not isinstance(value, str):
        return value

    def lookup(match: re.Match) -> str:
        if (found := os.environ.get(match[1])) is None:
            raise MCPServiceError(f"MCP 服务配置引用的环境变量 {match[1]} 未设置")
        return found

    return ENV_REFERENCE.sub(lookup, value)


def load_mcp_config(path: Path = CONFIG_PATH) -> dict[str, dict]:
    """读取服务列表并展开环境变量；文件缺失或格式不对时报错"""
    if not path.exists():
        raise MCPServiceError(f"未找到 MCP 服务配置: {path}")
    try:
        config = _expand_env(json.loads(path.read_text(encoding="utf-8")))
    except json.JSONDecodeError as error:
        raise MCPServiceError(f"MCP 服务配置不是合法 JSON: {error}") from error
    if not isinstance(config, dict) or not all(
        isinstance(item, dict) and item.get("transport") and item.get("url")
        for item in config.values()
    ):
        raise MCPServiceError("MCP 服务配置的每一项都需要 transport 和 url")
    return config


async def _limit_time(request: MCPToolCallRequest, handler):
    try:
        return await asyncio.wait_for(
            handler(request), timeout=settings.mcp_tool_timeout
        )
    except TimeoutError as error:
        raise MCPServiceError(
            f"MCP 工具调用超时: {request.server_name}.{request.name}"
        ) from error


async def _check_catalog_version(request: MCPToolCallRequest, handler):
    result = await handler(request)
    if request.server_name != MATERIAL_SERVER:
        return result
    if isinstance(result, CallToolResult) and result.structuredContent:
        served = result.structuredContent.get("catalog_version")
        catalog = get_catalog_cache()
        expected = catalog.version if catalog is not None else None
        if served != expected:
            raise MCPServiceError(
                f"物料服务的目录版本 {served} 与主服务 {expected} 不一致，请重启服务"
            )
    return result


_tools: list[BaseTool] | None = None


async def get_mcp_tools() -> list[BaseTool]:
    """首次使用时拉取所有服务的工具并缓存；此前失败时重新拉取"""
    global _tools
    if _tools is None:
        config = load_mcp_config()
        client = MultiServerMCPClient(
            {
                name: {key: value for key, value in item.items() if key != "module"}
                for name, item in config.items()
            },
            tool_interceptors=[_limit_time, _check_catalog_version],
            # 参数不对等查询错误返回给模型自行修正；连接失败等仍直接抛出
            handle_tool_errors=True,
        )
        try:
            tools = await asyncio.wait_for(
                client.get_tools(), timeout=settings.mcp_tool_timeout
            )
        except Exception as error:
            raise MCPServiceError(f"MCP 服务不可用: {', '.join(config)}") from error
        # 不加服务名前缀，skill 里写的工具名保持不变，因此不同服务的工具不能重名
        if duplicated := [
            name for name, count in Counter(t.name for t in tools).items() if count > 1
        ]:
            raise MCPServiceError(f"不同 MCP 服务的工具重名: {duplicated}")
        _tools = tools
    return _tools


def _port_open(port: int) -> bool:
    with socket.socket() as probe:
        probe.settimeout(0.2)
        return probe.connect_ex(("127.0.0.1", port)) == 0


async def _start(name: str, item: dict) -> subprocess.Popen | None:
    """本机地址且尚未运行时拉起服务；已运行则直接复用"""
    target = urlsplit(item["url"])
    if target.hostname not in LOCAL_HOSTS:
        raise MCPServiceError(f"{name} 由应用拉起，只能使用本机地址")
    port = target.port or 80
    if _port_open(port):
        return None
    process = subprocess.Popen(
        [sys.executable, "-X", "utf8", "-m", item["module"]],
        cwd=PROJECT_ROOT,
        env={**os.environ, "PORT": str(port)},
    )
    for _ in range(100):
        if process.poll() is not None:
            raise MCPServiceError(f"{name} 启动失败，请查看终端日志")
        if _port_open(port):
            logger.info("MCP 服务 {} 已启动: {}", name, item["url"])
            return process
        await asyncio.sleep(0.1)
    await stop_mcp_servers([process])
    raise MCPServiceError(f"{name} 启动超时")


async def start_mcp_servers() -> list[subprocess.Popen]:
    """拉起配置了 module 的服务，返回进程供退出时关闭；单个服务失败只记录日志"""
    processes = []
    for name, item in load_mcp_config().items():
        if not item.get("module"):
            continue
        try:
            if (process := await _start(name, item)) is not None:
                processes.append(process)
        except MCPServiceError as error:
            logger.error("{}，用到它的请求将无法执行", error)
    return processes


async def stop_mcp_servers(processes: list[subprocess.Popen]) -> None:
    for process in processes:
        if process.poll() is not None:
            continue
        process.terminate()
        try:
            await asyncio.to_thread(process.wait, timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            await asyncio.to_thread(process.wait)