"""MCP 协议网关：注册算法型工具并经 NATS 分发调用。

与 Go 侧 sundynix-shared/bus 的 ServeTool 同契约（core NATS request-reply）：
  请求体  {"tool": str, "args": {...}, "task_id": str}
  应答体  {"ok": bool, "content": str, "error": str}
订阅 sundynix.tools.py.>，队列组 mcp-py-workers（多副本负载均衡）。
"""

from __future__ import annotations

import asyncio
import json
import logging
import os

import nats

from .interpreter import CodeInterpreter
from .mineru import MultimodalParser
from .sandbox import SecureSandbox

log = logging.getLogger("mcp_py")

# 与 contract.SubjectToolsPyAll / QueueToolsPy 保持一致。
SUBJECT_PY_ALL = "sundynix.tools.py.>"
QUEUE_PY = "mcp-py-workers"

# 工具元信息（与 mcp-go 注册表同形）：list_tools 上报给管理端展示 + dispatcher 动态构建自主 agent 工具集。
# agent=是否暴露给自主 agent；params=模型可填参数；inject=服务端注入参数（不暴露给模型）。
TOOL_META = {
    "echo": {"cn": "回显", "desc": "原样返回入参（调试用）"},
    "run_code": {
        "cn": "代码执行",
        "desc": "在隔离沙箱里执行 Python 代码（静态守卫 + Docker，256m/10s）。用于计算、数据处理、逻辑验证。",
        "agent": True,
        "params": [{"name": "code", "type": "string", "desc": "要执行的 Python 代码", "required": True}],
    },
    "parse_document": {"cn": "文档解析", "desc": "文件 → 纯文本（MinerU / PaddleOCR）"},
    "secure_sandbox": {"cn": "安全沙箱", "desc": "更严资源档（128m/5s）的隔离执行，用于高风险代码"},
}


class McpGateway:
    def __init__(self) -> None:
        self.sandbox = SecureSandbox()       # gVisor / KataVM + Static Code Guard
        self.parser = MultimodalParser()     # MinerU / PaddleOCR
        self.interpreter = CodeInterpreter() # Docker 隔离沙箱
        self._nc: nats.NATS | None = None
        self._tools: dict[str, callable] = {}
        # 单实例工具并发上限（垂直扩）；队列组多副本再叠加水平扩。
        self._sem = asyncio.Semaphore(int(os.getenv("MCP_TOOL_CONCURRENCY", "16")))

    async def register_tools(self) -> None:
        """注册工具名 → 处理协程。"""
        self._tools = {
            "echo": self._echo,
            "run_code": self._run_code,
            "parse_document": self._parse_document,
            "secure_sandbox": self._secure_sandbox,
            "list_tools": self._list_tools,
        }

    async def serve(self, url: str | None = None) -> None:
        url = url or os.getenv("NATS_URL", "nats://localhost:4222")
        # 容忍服务先于 NATS 启动：无限重连 + 等待间隔。
        self._nc = await nats.connect(
            url,
            allow_reconnect=True,
            max_reconnect_attempts=-1,
            reconnect_time_wait=1,
            connect_timeout=5,
        )
        log.info("[mcp_py] connected %s", url)
        await self._nc.subscribe(SUBJECT_PY_ALL, queue=QUEUE_PY, cb=self._on_call)
        log.info(
            "[mcp_py] tools ready on %s (queue=%s): %s",
            SUBJECT_PY_ALL, QUEUE_PY, ", ".join(self._tools),
        )
        await asyncio.Event().wait()

    async def _on_call(self, msg) -> None:
        """立即派发到独立 task 并发处理，不阻塞订阅消息循环（nats-py 单订阅回调本是串行 await）。"""
        asyncio.create_task(self._handle(msg))

    async def _handle(self, msg) -> None:
        """解析 ToolCall → 路由到工具 → Respond ToolResult（限并发）。"""
        async with self._sem:
            try:
                req = json.loads(msg.data)
            except Exception as e:  # noqa: BLE001
                await self._reply(msg, ok=False, error=f"bad tool call: {e}")
                return
            tool = req.get("tool", "")
            task_id = req.get("task_id", "")
            args = req.get("args") or {}
            log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args)
            fn = self._tools.get(tool)
            if fn is None:
                await self._reply(msg, ok=False, error=f"unknown tool: {tool}")
                return
            try:
                content = await fn(args)
                await self._reply(msg, ok=True, content=content)
            except Exception as e:  # noqa: BLE001
                await self._reply(msg, ok=False, error=str(e))

    @staticmethod
    async def _reply(msg, *, ok: bool, content: str = "", error: str = "") -> None:
        payload = json.dumps({"ok": ok, "content": content, "error": error})
        await msg.respond(payload.encode())

    # ---- 工具实现（算法层目前为桩，但调用链路做真）----

    async def _echo(self, args: dict) -> str:
        return str(args.get("text", ""))

    async def _list_tools(self, args: dict) -> str:
        """自省：上报业务工具清单（名称/中文名/作用 + agent 暴露元信息），供管理端展示 + dispatcher 动态发现。"""
        tools = [
            {
                "name": n,
                "cn": m["cn"],
                "desc": m["desc"],
                "agent_exposed": m.get("agent", False),
                "agent_name": m.get("agent_name", ""),
                "params": m.get("params", []),
                "inject": m.get("inject", []),
            }
            for n, m in TOOL_META.items()
            if n in self._tools  # 仅上报真正注册的业务工具（list_tools 自身不计入）
        ]
        return json.dumps({"service": "mcp-py", "tools": tools})

    async def _run_code(self, args: dict) -> str:
        """静态守卫 → Docker 隔离执行（标准档 256m/0.5cpu/10s）。"""
        code = str(args.get("code", ""))
        if not code.strip():
            return "（空代码）"
        ok, reason = self.sandbox.static_guard(code)
        if not ok:
            return f"🚫 静态守卫拒绝：{reason}"
        return _fmt_result(await self.interpreter.execute(code))

    async def _parse_document(self, args: dict) -> str:
        """文件 → 纯文本。content_b64=文件内容(base64)，filename 决定解析器。"""
        import base64

        from . import parsers

        filename = str(args.get("filename", ""))
        content_b64 = str(args.get("content_b64", ""))
        if not content_b64:
            return str(args.get("text", ""))
        data = base64.b64decode(content_b64)
        # 解析是 CPU 密集，丢到线程池避免阻塞事件循环。
        return await asyncio.to_thread(parsers.parse, filename, data)

    async def _secure_sandbox(self, args: dict) -> str:
        """同 run_code 的守卫 + 隔离，更紧资源档（128m/5s），用于高风险代码。"""
        code = str(args.get("code", ""))
        if not code.strip():
            return "（空代码）"
        ok, reason = self.sandbox.static_guard(code)
        if not ok:
            return f"🚫 静态守卫拒绝：{reason}"
        return "[secure] " + _fmt_result(await self.interpreter.execute(code, mem="128m", timeout=5))

    async def close(self) -> None:
        if self._nc is not None:
            await self._nc.drain()


def _fmt_result(r: dict) -> str:
    """把执行结果 {ok,stdout,stderr,exit,degraded} 整理成可读文本。"""
    if r.get("degraded"):
        return f"⚠️ {r.get('stderr', 'Docker 不可用')}"
    parts = []
    if r.get("stdout", "").strip():
        parts.append("stdout:\n" + r["stdout"].strip())
    if r.get("stderr", "").strip():
        parts.append("stderr:\n" + r["stderr"].strip())
    parts.append(f"exit={r.get('exit')}")
    return "\n".join(parts)
