"""
=============================================================================
CS329Z · 前置讲 P2 共享模块：mcp_lab.py —— JSON-RPC 2.0 与一个迷你 MCP 服务器 / 客户端
=============================================================================

对应讲义：lectures/prereq2_http_apis_protocols.md §7–§8

【这个文件要回答的问题】
    MCP（Model Context Protocol）是今天把工具接给模型的事实标准。它听起来很大，其实只是：
        **JSON-RPC 2.0 消息** + **一次握手** + **几个约定好的方法名**（tools/list、tools/call……）
    这个文件把它从零实现一遍，两种传输都做：
        stdio           ：客户端把服务器**当子进程启动**，一行一个 JSON 消息（MCP 推荐、最常用）
        Streamable HTTP ：服务器是一个独立的进程，客户端 POST 消息过去（用于远程服务器）

【这个文件里有什么】
    JSON-RPC 的三种消息：request（有 id，要回复）、notification（没有 id，不回复）、response
    rpc_request / rpc_notification / rpc_result / rpc_error：构造这些消息
    ToolBox            ：几个玩具工具（查员工、查邮件、算日期差），每个带 JSON Schema
    MiniMCPServer      ：处理 initialize / notifications/initialized / tools/list / tools/call / ping
    serve_stdio()      ：把服务器跑成"从 stdin 读、往 stdout 写"的程序（本文件可以直接被当作服务器启动）
    StdioClient        ：启动一个子进程服务器，完成握手，调用工具
    HttpMCPServer / serve_http()：同一个服务器，跑在 HTTP 上（带 Mcp-Session-Id、协议版本头、Origin 检查）
    HttpMCPClient      ：HTTP 版的客户端

【协议的关键点（都来自 MCP 规范 2025-06-18）】
    - stdio：消息**用换行分隔，消息里不能有换行**；服务器**不能**往 stdout 写任何非 MCP 的内容
      （日志要写 stderr）。关闭时客户端先关 stdin，再等，必要时 SIGTERM / SIGKILL（前置讲 P1 §7.3）。
    - 握手：客户端先发 `initialize`（带 protocolVersion、capabilities、clientInfo），
      服务器回自己的 capabilities 和 serverInfo，客户端再发 `notifications/initialized` 通知。
    - HTTP：服务器在 `InitializeResult` 的响应里用 `Mcp-Session-Id` 头给出会话 id，
      客户端之后每个请求都要带上它，并且带 `MCP-Protocol-Version` 头。
    - 两种错误要分清：**协议错误**（不存在的工具、参数不合法）用 JSON-RPC 的 error 字段；
      **工具执行错误**（API 失败、业务逻辑错误）用正常的 result，但 `isError: true`。
      前者是"你问错了"，后者是"你问对了，但事情没办成"——模型要看到后者才能自己纠正。

【运行】
    python3 mcp_lab.py              # 自检：用 stdio 和 HTTP 各跑一遍握手 + 调用
    python3 mcp_lab.py --stdio-server   # 把自己当成一个 MCP 服务器跑（被 StdioClient 这样启动）
【依赖】只用 Python 标准库；import 同目录的 net_lab.py（复用 JSON Schema 校验）。
"""

import json
import os
import subprocess
import sys
import threading
import urllib.error
import urllib.parse
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from net_lab import toy_data, validate

PROTOCOL_VERSION = "2025-06-18"
SERVER_INFO = {"name": "cardinal-mini-mcp", "title": "Cardinal Mini MCP Server", "version": "0.1.0"}
CLIENT_INFO = {"name": "cs329z-demo-client", "version": "0.1.0"}

# JSON-RPC 2.0 规定的几个标准错误码
PARSE_ERROR, INVALID_REQUEST, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR = \
    -32700, -32600, -32601, -32602, -32603


# ============================================================================ 1. JSON-RPC 2.0 的消息
def rpc_request(id_, method, params=None):
    """一个**请求**：有 id，对方必须回一条带同样 id 的响应。"""
    msg = {"jsonrpc": "2.0", "id": id_, "method": method}
    if params is not None:
        msg["params"] = params
    return msg


def rpc_notification(method, params=None):
    """一个**通知**：没有 id，对方**不回**任何东西（也就无法知道它处理得怎么样）。"""
    msg = {"jsonrpc": "2.0", "method": method}
    if params is not None:
        msg["params"] = params
    return msg


def rpc_result(id_, result):
    """一个成功的**响应**。"""
    return {"jsonrpc": "2.0", "id": id_, "result": result}


def rpc_error(id_, code, message, data=None):
    """一个失败的**响应**（协议层面的失败）。"""
    err = {"code": code, "message": message}
    if data is not None:
        err["data"] = data
    return {"jsonrpc": "2.0", "id": id_, "error": err}


# ============================================================================ 2. 玩具工具
def _employees(team=None, limit=10):
    rows = toy_data()["employees"]
    if team:
        rows = [r for r in rows if r["team"] == team]
    return rows[:limit]


def _emails(urgent=None, limit=10):
    rows = toy_data()["emails"]
    if urgent is not None:
        rows = [r for r in rows if r["urgent"] is urgent]
    return rows[:limit]


def _office_days(employee_id):
    rows = {r["id"]: r for r in toy_data()["employees"]}
    if employee_id not in rows:
        raise LookupError(f"no employee with id {employee_id!r}")   # → 工具执行错误（isError）
    r = rows[employee_id]
    return f"{r['name']} ({r['team']}) is in the office {r['office_days']} days per week"


class ToolBox:
    """
    服务器提供的工具集：每个工具有名字、描述、inputSchema，和一个 Python 函数。

    【思路】和第 1 讲的工具循环一样：**校验和执行要分开**。validate() 先检查参数合不合 schema，
        合法了才真的调用函数——这样一个参数写错的调用在执行之前就被挡住（也在请用户批准之前）。
    """

    def __init__(self):
        self.tools = [
            {"name": "list_employees",
             "title": "List employees",
             "description": "List Cardinal employees, optionally filtered by team.",
             "inputSchema": {"type": "object",
                             "properties": {"team": {"type": "string", "enum": ["firmware", "analytics", "legal"]},
                                            "limit": {"type": "integer", "minimum": 1, "maximum": 50}},
                             "additionalProperties": False},
             "fn": _employees},
            {"name": "list_emails",
             "title": "List emails",
             "description": "List emails from the archive. Set urgent=true for urgent ones only.",
             "inputSchema": {"type": "object",
                             "properties": {"urgent": {"type": "boolean"},
                                            "limit": {"type": "integer", "minimum": 1, "maximum": 50}},
                             "additionalProperties": False},
             "fn": _emails},
            {"name": "office_days",
             "title": "Office days for one employee",
             "description": "How many days per week one employee is in the office. Needs an employee id like 'e-01'.",
             "inputSchema": {"type": "object",
                             "properties": {"employee_id": {"type": "string", "pattern": r"^e-\d{2}$"}},
                             "required": ["employee_id"],
                             "additionalProperties": False},
             "fn": _office_days},
        ]

    def schemas(self):
        """给模型（和 tools/list）看的部分：不含 Python 函数。"""
        return [{k: v for k, v in t.items() if k != "fn"} for t in self.tools]

    def find(self, name):
        """按名字找工具。【返回】工具 dict 或 None"""
        return next((t for t in self.tools if t["name"] == name), None)

    def call(self, name, arguments):
        """
        调用一个工具。

        【返回】MCP 的 `tools/call` 结果：{"content": [{"type": "text", "text": ...}], "isError": bool}
        【边界】
            - 工具不存在 / 参数不合 schema → 抛 ValueError，由服务器转成**协议错误**
            - 工具自己抛异常（例如查不到这个员工）→ 返回 isError=true 的**正常结果**，
              把错误信息交给模型，让它自己改正
        """
        tool = self.find(name)
        if tool is None:
            raise ValueError(f"unknown tool: {name}")
        ok, why = validate(tool["inputSchema"], arguments)
        if not ok:
            raise ValueError(f"invalid arguments: {why}")
        try:
            out = tool["fn"](**arguments)
        except Exception as e:                                  # 工具执行错误
            return {"content": [{"type": "text", "text": f"{type(e).__name__}: {e}"}], "isError": True}
        text = out if isinstance(out, str) else json.dumps(out, ensure_ascii=False, sort_keys=True)
        return {"content": [{"type": "text", "text": text}], "isError": False}


# ============================================================================ 3. 服务器的核心逻辑（和传输无关）
class MiniMCPServer:
    """
    MCP 服务器的"大脑"：收一条 JSON-RPC 消息，返回一条响应（通知则返回 None）。

    【状态】initialized：客户端有没有发过 notifications/initialized
    【支持的方法】initialize、notifications/initialized、tools/list、tools/call、ping
    """

    def __init__(self):
        self.box = ToolBox()
        self.initialized = False
        self.log = []                 # 收到的方法名，方便演示时回看

    def handle(self, msg):
        """
        处理一条消息。

        【参数】msg：已经解析好的 dict
        【返回】要回的 dict，或者 None（收到的是通知）
        【边界】
            - 不是合法的 JSON-RPC 2.0（缺 jsonrpc / method）→ INVALID_REQUEST
            - 握手之前调用 tools/list 之类 → INVALID_REQUEST（规范要求先 initialize）
            - 不认识的方法 → METHOD_NOT_FOUND
        """
        if msg.get("jsonrpc") != "2.0" or "method" not in msg:
            return rpc_error(msg.get("id"), INVALID_REQUEST, "not a valid JSON-RPC 2.0 message")
        method, id_, params = msg["method"], msg.get("id"), msg.get("params") or {}
        self.log.append(method)

        if method == "initialize":
            asked = params.get("protocolVersion")
            # 版本协商：支持就回同一个版本，不支持就回我们自己支持的最新版本
            version = PROTOCOL_VERSION if asked == PROTOCOL_VERSION else PROTOCOL_VERSION
            return rpc_result(id_, {"protocolVersion": version,
                                    "capabilities": {"tools": {"listChanged": False}},
                                    "serverInfo": SERVER_INFO,
                                    "instructions": "Use list_employees / list_emails to look things up."})
        if method == "notifications/initialized":
            self.initialized = True
            return None                                 # 通知不回复
        if method == "ping":
            return rpc_result(id_, {})
        if not self.initialized:
            return rpc_error(id_, INVALID_REQUEST, "not initialized: send initialize first")
        if method == "tools/list":
            return rpc_result(id_, {"tools": self.box.schemas()})
        if method == "tools/call":
            name = params.get("name")
            args = params.get("arguments") or {}
            try:
                return rpc_result(id_, self.box.call(name, args))
            except ValueError as e:                     # 协议错误：工具不存在或参数不合法
                code = METHOD_NOT_FOUND if str(e).startswith("unknown tool") else INVALID_PARAMS
                return rpc_error(id_, code, str(e))
        return rpc_error(id_, METHOD_NOT_FOUND, f"unknown method: {method}")


# ============================================================================ 4. stdio 传输
def serve_stdio(stdin=None, stdout=None, stderr=None):
    """
    把服务器跑成一个"从 stdin 读、往 stdout 写"的程序。

    【思路】每行一条 JSON 消息（MCP 规定：消息用换行分隔，消息内部**不能**有换行，
        所以序列化时不能用缩进）。日志只能写 stderr——往 stdout 写一句闲话就会把协议搞坏。
    【参数】三条流，默认用进程自己的（便于测试时替换）
    【返回】None（读到 stdin 结束就退出，退出码 0）
    """
    stdin = stdin or sys.stdin
    stdout = stdout or sys.stdout
    stderr = stderr or sys.stderr
    server = MiniMCPServer()
    print("mini MCP server ready (logging to stderr)", file=stderr, flush=True)
    for line in stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError as e:
            out = rpc_error(None, PARSE_ERROR, f"invalid JSON: {e}")
        else:
            out = server.handle(msg)
        if out is not None:
            stdout.write(json.dumps(out, ensure_ascii=False) + "\n")
            stdout.flush()                              # 不刷新的话消息会卡在缓冲区里，客户端一直等
        print(f"handled {msg.get('method') if isinstance(msg, dict) else '?'}", file=stderr, flush=True)


class StdioClient:
    """
    stdio 版的 MCP 客户端：把服务器当子进程启动，握手，然后调用工具。

    【用法】
        with StdioClient([sys.executable, "mcp_lab.py", "--stdio-server"]) as c:
            c.list_tools(); c.call_tool("list_employees", {"team": "legal"})
    【边界】关闭时按规范来：先关 stdin（服务器读到结束就会退出），等一会儿，还不退就 terminate / kill。
    """

    def __init__(self, argv, timeout=5.0):
        self.argv = list(argv)
        self.timeout = timeout
        self.proc = None
        self._stderr = ""
        self._closed = False
        self._id = 0
        self.server_info = None
        self.capabilities = None

    def __enter__(self):
        self.proc = subprocess.Popen(self.argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE, text=True, bufsize=1,
                                     start_new_session=True)
        self.initialize()
        return self

    def __exit__(self, *exc):
        self.close()
        return False

    def _next_id(self):
        self._id += 1
        return self._id

    def send(self, msg, expect_reply=True):
        """
        发一条消息；expect_reply 为真时读一行回复。

        【返回】回复的 dict（不期待回复时返回 None）
        【边界】服务器死了或者没按时回话 → 抛 RuntimeError，并附上它 stderr 的内容（最常见的调试线索）
        """
        self.proc.stdin.write(json.dumps(msg, ensure_ascii=False) + "\n")
        self.proc.stdin.flush()
        if not expect_reply:
            return None
        line = self.proc.stdout.readline()
        if not line:
            raise RuntimeError("server closed its stdout without replying")
        return json.loads(line)

    def initialize(self):
        """完成握手：initialize → （记下服务器信息）→ notifications/initialized。"""
        reply = self.send(rpc_request(self._next_id(), "initialize", {
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {},
            "clientInfo": CLIENT_INFO,
        }))
        res = reply["result"]
        self.server_info = res["serverInfo"]
        self.capabilities = res["capabilities"]
        self.send(rpc_notification("notifications/initialized"), expect_reply=False)
        return res

    def list_tools(self):
        """tools/list。【返回】工具 schema 列表"""
        return self.send(rpc_request(self._next_id(), "tools/list"))["result"]["tools"]

    def call_tool(self, name, arguments=None):
        """tools/call。【返回】完整的 JSON-RPC 回复（可能是 result，也可能是 error）"""
        return self.send(rpc_request(self._next_id(), "tools/call",
                                     {"name": name, "arguments": arguments or {}}))

    def stderr_text(self):
        """服务器写到 stderr 的日志（close() 之后才有内容：关闭时一次性读完）。"""
        return getattr(self, "_stderr", "")

    def close(self):
        """
        按 MCP 规范关闭：关 stdin → 等 → terminate → kill。

        【返回】服务器的退出码（没启动过则为 None）
        【边界】可以重复调用：第二次直接返回退出码（所以在 with 块里手动调一次也不会出错）。
        """
        if self.proc is None or self._closed:        # 重复调用是安全的（with 块里手动 close 过也没事）
            return None if self.proc is None else self.proc.returncode
        self._closed = True
        try:
            self.proc.stdin.close()
        except Exception:
            pass
        try:
            self.proc.wait(timeout=self.timeout)
        except subprocess.TimeoutExpired:
            self.proc.terminate()
            try:
                self.proc.wait(timeout=1.0)
            except subprocess.TimeoutExpired:
                self.proc.kill()
                self.proc.wait(timeout=1.0)
        self._stderr = self.proc.stderr.read() if self.proc.stderr and not self.proc.stderr.closed else ""
        for f in (self.proc.stdout, self.proc.stderr):     # 我们这一头的管道要自己关（P1 §7.2）
            if f is not None and not f.closed:
                f.close()
        return self.proc.returncode


# ============================================================================ 5. Streamable HTTP 传输
class MCPHandler(BaseHTTPRequestHandler):
    """
    MCP 的 Streamable HTTP 传输（最简版）：一个端点 /mcp，POST 收消息。

    【实现了规范里的哪些要求】
        - 校验 Origin 头（防 DNS rebinding 攻击）
        - initialize 的响应里用 Mcp-Session-Id 头给出会话 id；之后的请求必须带上它，否则 400
        - 收到通知 / 响应时返回 202 Accepted、没有响应体
        - 认不出的会话 id → 404（客户端应该重新 initialize）
        - DELETE 结束会话
    【没实现】SSE 流（规范允许服务器用单个 JSON 响应，我们就用这种）、断线重连、分页。
    """

    server_version = "MiniMCP/0.1"
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        self.server.access_log.append(f"{self.command} {self.path}")

    def _json(self, obj, status=200, extra=()):
        body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        for k, v in extra:
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(body)

    def _empty(self, status):
        self.send_response(status)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_POST(self):
        if self.path != "/mcp":
            return self._empty(404)
        origin = self.headers.get("Origin")
        if origin and origin not in self.server.allowed_origins:
            self.server.rejected_origins.append(origin)
            return self._json({"error": "origin not allowed"}, status=403)
        n = int(self.headers.get("Content-Length") or 0)
        try:
            msg = json.loads(self.rfile.read(n).decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as e:
            return self._json(rpc_error(None, PARSE_ERROR, f"invalid JSON: {e}"), status=400)
        session_id = self.headers.get("Mcp-Session-Id")
        is_init = isinstance(msg, dict) and msg.get("method") == "initialize"

        if is_init:
            new_id = f"sess-{len(self.server.sessions) + 1:04d}"      # 真实实现要用随机、不可猜的 id
            self.server.sessions[new_id] = MiniMCPServer()
            out = self.server.sessions[new_id].handle(msg)
            return self._json(out, extra=[("Mcp-Session-Id", new_id)])
        if session_id is None:
            return self._json(rpc_error(msg.get("id"), INVALID_REQUEST, "missing Mcp-Session-Id header"), status=400)
        if session_id not in self.server.sessions:
            return self._json({"error": "unknown session"}, status=404)
        out = self.server.sessions[session_id].handle(msg)
        if out is None:                                 # 通知：202 + 空响应体
            return self._empty(202)
        return self._json(out)

    def do_DELETE(self):
        sid = self.headers.get("Mcp-Session-Id")
        if sid in self.server.sessions:
            del self.server.sessions[sid]
            return self._empty(204)
        return self._empty(404)

    def do_GET(self):
        """规范允许客户端用 GET 开一条 SSE 流（接收服务器主动发的消息）；我们不支持，按规范返回 405。"""
        return self._empty(405)


class MCPHttpServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sessions = {}
        self.access_log = []
        self.rejected_origins = []
        self.allowed_origins = {"http://localhost", "http://127.0.0.1"}


class serve_http:
    """在本机起一个 HTTP 版的 MCP 服务器（用法同 net_lab.serve）。"""

    def __init__(self, host="127.0.0.1"):
        self.httpd = MCPHttpServer((host, 0), MCPHandler)
        self.base = f"http://{host}:{self.httpd.server_address[1]}"
        self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)

    def __enter__(self):
        self.thread.start()
        return self.base

    def __exit__(self, *exc):
        self.httpd.shutdown()
        self.httpd.server_close()
        self.thread.join(timeout=2)
        return False


class HttpMCPClient:
    """
    HTTP 版的 MCP 客户端：POST 消息到 /mcp，带上会话 id 和协议版本头。

    【用法】c = HttpMCPClient(base); c.initialize(); c.list_tools()
    """

    def __init__(self, base, timeout=3.0, origin=None):
        self.base = base.rstrip("/") + "/mcp"
        self.timeout = timeout
        self.origin = origin
        self.session_id = None
        self._id = 0
        self.last_status = None

    def _post(self, msg):
        """
        发一条消息。【返回】(status, headers, 解析好的响应或 None)
        【边界】HTTP 错误不抛异常，而是把状态码交回去——协议错误和 HTTP 错误要分开看。
        """
        self._id += 1
        hdrs = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
        if self.session_id:
            hdrs["Mcp-Session-Id"] = self.session_id
            hdrs["MCP-Protocol-Version"] = PROTOCOL_VERSION
        if self.origin:
            hdrs["Origin"] = self.origin
        req = urllib.request.Request(self.base, data=json.dumps(msg).encode("utf-8"),
                                     headers=hdrs, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                payload = resp.read()
                self.last_status = resp.status
                obj = json.loads(payload.decode("utf-8")) if payload else None
                return resp.status, dict(resp.headers), obj
        except urllib.error.HTTPError as e:
            payload = e.read()
            self.last_status = e.code
            try:
                obj = json.loads(payload.decode("utf-8")) if payload else None
            except json.JSONDecodeError:
                obj = None
            return e.code, dict(e.headers), obj

    def initialize(self):
        """握手，并记下服务器给的会话 id。【返回】(status, result 或 None)"""
        status, headers, obj = self._post(rpc_request(self._id + 1, "initialize", {
            "protocolVersion": PROTOCOL_VERSION, "capabilities": {}, "clientInfo": CLIENT_INFO}))
        if status == 200:
            self.session_id = headers.get("Mcp-Session-Id")
            self._post(rpc_notification("notifications/initialized"))
        return status, (obj or {}).get("result")

    def list_tools(self):
        """tools/list。【返回】(status, 工具列表或 None)"""
        status, _, obj = self._post(rpc_request(self._id + 1, "tools/list"))
        return status, ((obj or {}).get("result") or {}).get("tools")

    def call_tool(self, name, arguments=None):
        """tools/call。【返回】(status, 完整的 JSON-RPC 响应)"""
        status, _, obj = self._post(rpc_request(self._id + 1, "tools/call",
                                                {"name": name, "arguments": arguments or {}}))
        return status, obj

    def delete_session(self):
        """显式结束会话（规范建议客户端不再需要会话时这样做）。【返回】HTTP 状态码"""
        req = urllib.request.Request(self.base, method="DELETE",
                                     headers={"Mcp-Session-Id": self.session_id or ""})
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                return resp.status
        except urllib.error.HTTPError as e:
            return e.code


if __name__ == "__main__":
    if "--stdio-server" in sys.argv:
        serve_stdio()
        sys.exit(0)
    here = os.path.dirname(os.path.abspath(__file__))
    print("=== stdio 传输 ===")
    with StdioClient([sys.executable, os.path.join(here, "mcp_lab.py"), "--stdio-server"]) as c:
        print("serverInfo:", c.server_info, "| capabilities:", c.capabilities)
        print("工具：", [t["name"] for t in c.list_tools()])
        print("调用 office_days(e-05):", c.call_tool("office_days", {"employee_id": "e-05"})["result"]["content"][0]["text"])
        print("调用一个不存在的工具 :", c.call_tool("send_email", {})["error"])
    print("\n=== Streamable HTTP 传输 ===")
    with serve_http() as base:
        h = HttpMCPClient(base)
        status, res = h.initialize()
        print(f"initialize → HTTP {status}，会话 id {h.session_id}，协议版本 {res['protocolVersion']}")
        print("工具：", [t["name"] for t in h.list_tools()[1]])
        print("调用 list_employees(legal):", h.call_tool("list_employees", {"team": "legal"})[1]["result"]["content"][0]["text"])
        print("DELETE 会话 →", h.delete_session())