#!/usr/bin/env python3
"""A small MCP client: JSON-RPC over stdio or Streamable HTTP, with no dependencies.

Purpose: connect an MCP server to the Part 24 agent, and show what the transport
    actually is. The stdio binding is newline-delimited JSON-RPC over a subprocess's
    standard streams, and the Streamable HTTP binding is one POST per message to a
    single endpoint, so both fit in a file you can read in ten minutes. The bridge
    lists the server's tools, converts them into the OpenAI function schemas the agent
    already speaks, prefixes their names with a server label so two servers cannot
    collide, and calls them. It speaks the 2026-07-28 revision, which carries its
    metadata per request, and falls back to the older initialize handshake exactly as
    the specification's backward-compatibility section describes.
Platform: all (pure Python; stdio needs a POSIX-like process model, so use WSL2 on Windows)
Minimum memory: 8 GB on the machine serving the model; this module needs almost none
Assumes: Python 3.9 or later and no third-party packages. A server to talk to: this
    part's mcp-server.py, or any other MCP server you have reviewed. Nothing here
    validates that a server is trustworthy; the lesson's checklist is that job.

Usage: imported by minimal-agent.py --mcp and by mcp-client-check.py:
           from mcpbridge import McpBridge
           bridge = McpBridge("python3 mcp-server.py")
       run directly to list a server's tools and print its identity:
           python3 mcpbridge.py --server "python3 mcp-server.py"
           python3 mcpbridge.py --url http://127.0.0.1:8000/mcp
"""

from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
import urllib.error
import urllib.request
from typing import Any, Dict, List, Optional

PROTOCOL_VERSION = "2026-07-28"
LEGACY_VERSION = "2025-06-18"
CLIENT_INFO = {"name": "course-minimal-agent", "version": "1.0.0"}


class McpError(Exception):
    """A JSON-RPC error, or a transport that gave up."""


def _meta() -> Dict[str, Any]:
    """The per-request metadata the current revision carries in every params object."""
    return {
        "io.modelcontextprotocol/protocolVersion": PROTOCOL_VERSION,
        "io.modelcontextprotocol/clientInfo": CLIENT_INFO,
        "io.modelcontextprotocol/clientCapabilities": {},
    }


def _label_from(source: str) -> str:
    """A short, name-safe label for one server, used to prefix its tool names.

    The specification warns that a client aggregating several servers may meet two
    tools with the same name and should disambiguate, and that the server's own name
    is not unique enough to rely on. A label taken from how you started it is.
    """
    words = [w for w in shlex.split(source) if not w.startswith("-")]
    stem = (words[-1] if words else source).replace("\\", "/").split("/")[-1]
    stem = stem.split(".")[0] or "mcp"
    return "".join(c if c.isalnum() else "_" for c in stem).strip("_")[:24] or "mcp"


class McpBridge:
    """One connection to one MCP server, exposed as OpenAI-shaped tools."""

    def __init__(self, server: Optional[str] = None, url: Optional[str] = None,
                 timeout: int = 60, label: Optional[str] = None,
                 stderr_to: Optional[Any] = None) -> None:
        if bool(server) == bool(url):
            raise McpError("give exactly one of a server command or a URL")
        self.url = url
        self.timeout = timeout
        self.label = label or _label_from(server or url or "mcp")
        self._next_id = 0
        self._process: Optional[subprocess.Popen] = None
        self.legacy = False

        if server:
            # stderr is the server's logging channel per the specification, so it is
            # never parsed; sending it to the terminal is what makes a broken server
            # debuggable rather than merely silent.
            self._process = subprocess.Popen(  # noqa: S603 - command supplied by the operator
                shlex.split(server), stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=stderr_to if stderr_to is not None else sys.stderr,
                text=True, bufsize=1)
        self.server_info = self._handshake()
        self.tools = self._list_tools()
        self.tool_names = {self.label + "_" + t["name"] for t in self.tools}

    # ------------------------------------------------------------- transport

    def request(self, method: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """One JSON-RPC request, with the per-request metadata this revision needs."""
        return self._request(method, self._params(params or {}))

    def _request(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
        self._next_id += 1
        message = {"jsonrpc": "2.0", "id": self._next_id, "method": method, "params": params}
        reply = self._over_http(message) if self.url else self._over_stdio(message)
        if "error" in reply:
            raise McpError("%s: %s" % (method, reply["error"].get("message", reply["error"])))
        return reply.get("result", {})

    def _over_stdio(self, message: Dict[str, Any]) -> Dict[str, Any]:
        assert self._process is not None and self._process.stdin and self._process.stdout
        self._process.stdin.write(json.dumps(message) + "\n")
        self._process.stdin.flush()
        while True:
            line = self._process.stdout.readline()
            if not line:
                raise McpError("the server closed its output stream; check its stderr above")
            try:
                reply = json.loads(line)
            except ValueError:
                # A server that prints to stdout corrupts the stream. Say so plainly:
                # this is the most common stdio bug and it looks like a protocol fault.
                raise McpError("non-JSON on the server's stdout: %r. A server must write "
                               "nothing but MCP messages to stdout; send logs to stderr."
                               % line[:120]) from None
            if reply.get("id") == message["id"]:
                return reply  # anything else is a notification for another request

    def _over_http(self, message: Dict[str, Any]) -> Dict[str, Any]:
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
            "MCP-Protocol-Version": PROTOCOL_VERSION,
            "Mcp-Method": message["method"],
        }
        name = message["params"].get("name") or message["params"].get("uri")
        if name:
            headers["Mcp-Name"] = name
        request = urllib.request.Request(self.url, data=json.dumps(message).encode("utf-8"),
                                         headers=headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                body = response.read().decode("utf-8")
                if "text/event-stream" in (response.headers.get("Content-Type") or ""):
                    return self._last_sse_message(body, message["id"])
                return json.loads(body)
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", "replace")[:300]
            raise McpError("HTTP %s from %s: %s" % (exc.code, self.url, detail)) from None
        except urllib.error.URLError as exc:
            raise McpError("could not reach %s: %s" % (self.url, exc.reason)) from None

    @staticmethod
    def _last_sse_message(body: str, want_id: int) -> Dict[str, Any]:
        """Pick the JSON-RPC response out of an SSE stream, ignoring notifications."""
        for line in body.splitlines():
            if not line.startswith("data:"):
                continue
            try:
                payload = json.loads(line[5:].strip())
            except ValueError:
                continue
            if payload.get("id") == want_id:
                return payload
        raise McpError("the response stream ended without a response to request %d" % want_id)

    # ------------------------------------------------------------- handshake

    def _handshake(self) -> Dict[str, Any]:
        """Probe with server/discover, and fall back to initialize for older servers.

        The specification recommends this probe rather than keying the fallback to one
        error code, because a server speaking an older revision may answer an unknown
        pre-initialize request with any implementation-defined error, or with nothing.
        """
        try:
            return self._request("server/discover", {"_meta": _meta()})
        except McpError:
            self.legacy = True
        result = self._request("initialize", {
            "protocolVersion": LEGACY_VERSION, "capabilities": {}, "clientInfo": CLIENT_INFO})
        self._notify("notifications/initialized", {})
        return result

    def _notify(self, method: str, params: Dict[str, Any]) -> None:
        if self.url or self._process is None or self._process.stdin is None:
            return
        self._process.stdin.write(json.dumps(
            {"jsonrpc": "2.0", "method": method, "params": params}) + "\n")
        self._process.stdin.flush()

    def _params(self, extra: Dict[str, Any]) -> Dict[str, Any]:
        return dict(extra) if self.legacy else dict(extra, _meta=_meta())

    # ------------------------------------------------------------------ use

    def _list_tools(self) -> List[Dict[str, Any]]:
        tools: List[Dict[str, Any]] = []
        cursor = None
        while True:
            params = self._params({"cursor": cursor} if cursor else {})
            result = self._request("tools/list", params)
            tools.extend(result.get("tools", []))
            cursor = result.get("nextCursor")
            if not cursor:
                return tools

    def schemas(self) -> List[Dict[str, Any]]:
        """The server's tools, in the shape a chat-completions request wants.

        The description is copied through unchanged, which is the point at which a
        poisoned server's instructions would enter your prompt. Print this before you
        trust a server you did not write.
        """
        out = []
        for tool in self.tools:
            out.append({"type": "function", "function": {
                "name": self.label + "_" + tool["name"],
                "description": tool.get("description", ""),
                "parameters": tool.get("inputSchema", {"type": "object"}),
            }})
        return out

    def call(self, prefixed_name: str, arguments: Dict[str, Any]) -> str:
        name = prefixed_name[len(self.label) + 1:]
        result = self._request("tools/call", self._params(
            {"name": name, "arguments": arguments}))
        parts = [block.get("text", "") for block in result.get("content", [])
                 if block.get("type") == "text"]
        if not parts and "structuredContent" in result:
            parts = [json.dumps(result["structuredContent"])]
        text = "\n".join(p for p in parts if p) or "(the tool returned no text)"
        return ("error: " + text) if result.get("isError") else text

    def close(self) -> None:
        """Shut the server down the way the specification describes: close stdin, wait, kill."""
        if self._process is None:
            return
        try:
            if self._process.stdin:
                self._process.stdin.close()
            self._process.wait(timeout=5)
        except (subprocess.TimeoutExpired, OSError):
            self._process.kill()
        finally:
            self._process = None


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--server", default=None, help="command that starts a stdio server")
    parser.add_argument("--url", default=None, help="a Streamable HTTP MCP endpoint")
    parser.add_argument("--json", action="store_true", help="print the tool schemas as JSON")
    args = parser.parse_args()

    bridge = McpBridge(server=args.server, url=args.url)
    try:
        if args.json:
            print(json.dumps(bridge.schemas(), indent=2))
            return
        print("connected over %s, %s protocol" % (
            "Streamable HTTP" if args.url else "stdio",
            "legacy initialize" if bridge.legacy else PROTOCOL_VERSION))
        print("server said: %s" % json.dumps(bridge.server_info)[:300])
        print("\n%d tool(s), exposed to the agent with the prefix %r:" % (
            len(bridge.tools), bridge.label + "_"))
        for tool in bridge.tools:
            print("\n  %s%s" % (bridge.label + "_", tool["name"]))
            print("    description: %s" % (tool.get("description", "") or "(none)"))
            print("    input schema: %s" % json.dumps(tool.get("inputSchema", {})))
        print("\nRead every description above in full before trusting this server: the model "
              "sees all of it, and a host's summary view does not.")
    finally:
        bridge.close()


if __name__ == "__main__":
    try:
        main()
    except McpError as exc:
        raise SystemExit("mcp: %s" % exc)