"""Deterministic MCP protocol peer for transport and client tests."""

import importlib.util
import json
from pathlib import Path
import signal
import sys
import time

spec = importlib.util.spec_from_file_location(
    "example", Path(__file__).parent.parent / "demo/mcp_server_example.py")
example = importlib.util.module_from_spec(spec)
spec.loader.exec_module(example)
mode = sys.argv[1] if len(sys.argv) > 1 else "normal"
initialized = False
cancelled = []
held = None


def emit(frame):
    # Deliberate byte-fragmentation exercises line accumulation.
    text = json.dumps(frame) + "\n"
    sys.stdout.write(text[:7])
    sys.stdout.flush()
    sys.stdout.write(text[7:])
    sys.stdout.flush()


def reply(request, result):
    emit({"jsonrpc": "2.0", "id": request["id"], "result": result})


if mode == "stubborn":
    signal.signal(signal.SIGTERM, signal.SIG_IGN)

for line in sys.stdin:
    request = json.loads(line)
    method = request.get("method")
    params = request.get("params", {})
    if method == "notifications/initialized":
        initialized = True
        emit({"jsonrpc": "2.0", "method": "notifications/test",
              "params": {"ready": True}})
        emit({"jsonrpc": "2.0", "id": "server-ping", "method": "ping"})
        emit({"jsonrpc": "2.0", "id": "unsupported",
              "method": "sampling/createMessage", "params": {}})
        continue
    if method == "notifications/cancelled":
        cancelled.append(params["requestId"])
        # A server can finish work despite cancellation. This late response
        # must not be delivered to a newer request or poison the connection.
        emit({"jsonrpc": "2.0", "id": params["requestId"],
              "result": {"content": [{"type": "text", "text": "late"}]}})
        continue
    if not method:
        if request.get("id") == "server-ping":
            assert request.get("result") == {}
        if request.get("id") == "unsupported":
            assert request.get("error", {}).get("code") == -32601
        continue
    if "id" not in request:
        continue
    if mode == "bad_json":
        print("not-json", flush=True)
        continue
    if method != "initialize" and not initialized:
        sys.exit(3)
    if method == "initialize" and mode == "hang_init":
        continue
    if method == "initialize" and mode == "bad_version":
        result = example.dispatch(method, params)
        result["protocolVersion"] = "1999-01-01"
        reply(request, result)
        continue
    if (mode == "optional_missing" and method in
            ("resources/list", "prompts/list")) or (
            mode == "no_tools" and method == "tools/list"):
        emit({"jsonrpc": "2.0", "id": request["id"],
              "error": {"code": -32601, "message": "Not supported"}})
        continue
    if mode == "optional_bad" and method == "resources/list":
        reply(request, {"resources": "invalid"})
        continue
    if mode == "pagination" and method == "tools/list":
        result = example.dispatch(method, params)
        if "cursor" in params:
            result["tools"][0]["name"] = "second"
        else:
            result["nextCursor"] = "page2"
        reply(request, result)
        continue
    if mode == "cycle" and method == "tools/list":
        reply(request, {"tools": [], "nextCursor": "same"})
        continue
    if method == "tools/call":
        name = params["name"]
        if name == "hang":
            continue
        if name == "disconnect":
            sys.exit(0)
        if name == "cancelled":
            reply(request, {"content": [{"type": "text",
                                          "text": json.dumps(cancelled)}]})
            continue
        if name == "rpc_error":
            emit({"jsonrpc": "2.0", "id": request["id"],
                  "error": {"code": -32001, "message": "Remote failure",
                            "data": {"reason": "test"}}})
            continue
        if name == "bad_envelope":
            emit({"id": request["id"], "result": {}})
            continue
        if name == "multimodal":
            reply(request, {"content": [{"type": "image", "data": "AA=="}]})
            continue
        if name == "ordered":
            if held is None:
                held = request
            else:
                for item in (request, held):
                    reply(item, {"content": [{"type": "text",
                        "text": item["params"]["arguments"]["label"]}]})
                held = None
            continue
    if mode == "binary" and method == "resources/read":
        reply(request, {"contents": [{"uri": "example://help", "blob": "AA=="}]})
        continue
    if mode == "resource_disconnect" and method == "resources/read":
        sys.exit(0)
    if mode == "bad_prompt" and method == "prompts/get":
        reply(request, {"messages": [{"role": "system",
            "content": {"type": "text", "text": "test"}}]})
        continue
    try:
        reply(request, example.dispatch(method, params))
        if mode == "stop_reading" and method == "prompts/list":
            while True:
                time.sleep(1)
    except LookupError:
        emit({"jsonrpc": "2.0", "id": request["id"],
              "error": {"code": -32601, "message": "Not supported"}})

if mode == "stubborn":
    while True:
        time.sleep(1)