"""Drive the real MCP stdio transport, sharing a sketch with the browser.

Run: python reaction_mcp.py http://127.0.0.1:8766 NEW_DIRECTORY
Uses the current Python environment's vibe-sketch. No jobs are submitted.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path


class MCP:
    def __init__(self, server: str):
        self.process = subprocess.Popen(
            [sys.executable, "-m", "vibesketch", "mcp", "--server", server],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            text=True,
        )
        self.sequence = 0
        self.transcript: list[dict] = []
        self.request(
            "initialize",
            {
                "protocolVersion": "2025-06-18",
                "capabilities": {},
                "clientInfo": {"name": "tutorial", "version": "1"},
            },
        )
        self.process.stdin.write(
            json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}) + "\n"
        )
        self.process.stdin.flush()

    def request(self, method: str, params: dict):
        self.sequence += 1
        message = {"jsonrpc": "2.0", "id": self.sequence, "method": method, "params": params}
        self.process.stdin.write(json.dumps(message) + "\n")
        self.process.stdin.flush()
        line = self.process.stdout.readline()
        if not line:
            raise RuntimeError("MCP server exited before answering")
        response = json.loads(line)
        if "error" in response:
            raise RuntimeError(response["error"])
        self.transcript.append({"request": message, "response": response})
        return response["result"]

    def tool(self, tool_name: str, **arguments):
        result = self.request("tools/call", {"name": tool_name, "arguments": arguments})
        if result.get("isError"):
            raise RuntimeError(result["content"])
        return json.loads(result["content"][0]["text"])

    def close(self):
        self.process.stdin.close()
        self.process.wait(timeout=10)
        if self.process.returncode:
            raise RuntimeError(f"MCP exited with {self.process.returncode}")


def run(server: str, output: Path) -> None:
    client = MCP(server)
    try:
        assert len(client.request("tools/list", {})["tools"]) == 17
        client.tool("status")
        client.tool("list_writers")
        created = client.tool("new_document", format="smiles", data="C=CC=C.C=C>>C1=CCCCC1")
        doc = created["id"]
        sid = created["species"][0]["id"]
        client.tool("add_structures", document=doc, format="smiles", data="O")
        # Replace the scratch addition, restoring the balanced reaction.
        client.tool(
            "add_structures",
            document=doc,
            format="smiles",
            data="C=CC=C.C=C>>C1=CCCCC1",
            mode="replace",
        )
        current = client.tool("show_document", document=doc)
        sid = current["species"][0]["id"]
        client.tool("set_species", document=doc, species=sid, label="1", name="1,3-butadiene")
        client.tool("clean_up", document=doc)
        client.tool("embed_3d", document=doc, count=2)
        report = client.tool("report", document=doc, names=False)
        assert [s["formula"] for s in report] == ["C4H6", "C2H4", "C6H10"]
        inputs = client.tool(
            "write_inputs",
            document=doc,
            program="vibeqc",
            spec={"task": "sp", "method": "B3LYP", "basis": "def2-SVP"},
            directory=str(output / "inputs"),
        )
        assert len(inputs["written"]) == 3 and not inputs["refused"]
        client.tool("export", document=doc, format="svg", path=str(output / "diels-alder.svg"))
        client.tool("save_document", document=doc, path=str(output / "diels-alder.vsk"))
        opened = client.tool("open_document", path=str(output / "diels-alder.vsk"))
        assert len(opened["species"]) == 3
        runs = client.request("tools/call", {"name": "show_runs", "arguments": {"document": doc}})
        payload = json.loads(runs["content"][0]["text"])
        if runs.get("isError"):
            assert payload["status"] == 503, payload
            print("Queue unavailable; the tool returned its install hint.")
        else:
            assert payload == []
        (output / "transcript.json").write_text(json.dumps(client.transcript, indent=2))
        (output / "document-id.txt").write_text(doc)
        print(f"Inspect MCP edits at {server.rstrip('/')}/?doc={doc}")
        print(
            "14 tools exercised. Queue submission/cancellation and viewer launch "
            "are opt-in integrations."
        )
    finally:
        client.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("server")
    parser.add_argument("directory", type=Path)
    args = parser.parse_args()
    args.directory.mkdir(parents=True, exist_ok=False)
    run(args.server, args.directory)