#!/usr/bin/env python3
"""Bounded MCP smoke test for a server declared in an .mcp.json (ersatztv#777).

Written because the caller's original check only asked `[ -x command ]`, which a
DIRECTORY satisfies (`[ -x /bin ]` is true), so it could report a pass for a server
that cannot run. The same weakness recurs one level in: accepting any response
carrying the right id passes a server that was not the configured one — or that
answered with a malformed body.
Hence the shape and identity checks below: "it answered" is not "it answered
correctly", and a smoke test that cannot tell them apart is decoration.

Usage:
  mcp_smoke.py <.mcp.json> <server> [timeout] [--expect-server NAME]
                                              [--expect-tool NAME]...

Exit 0 only when the server answered `initialize` and `tools/list` with
well-formed bodies, matched `--expect-server` if given, and exposed every
`--expect-tool`. Failures exit non-zero with a diagnostic naming the stage.
Codes group failures by STAGE (config=3-5, command=6, project=7, spawn=8,
protocol=9-11, malformed=12, identity=13, tools=14) — several distinct causes
deliberately share a stage code, so read the message, not the number.

ACCEPTED LIMIT: a server that writes a gigabyte with no newline can still exhaust
memory before the timeout fires. Guarding that needs the very frame-capping reader
whose caps are themselves the defect described below, and the input here is our OWN
configured server on a developer machine — not an adversary. Stated rather than
silently unhandled.

Deliberately NOT bounded by message/line caps — adding them IS the defect: an
over-long line has its suffix re-parsed as a fresh message (a false green), and
a cap reached before the awaited reply reports
"server did not start". What IS bounded is the set of retained DECODED responses
— only a reply to the request in flight is kept, notifications are dropped as
they arrive — and wall-clock, via the caller's timeout. The raw line buffer is
explicitly NOT bounded; that is the accepted limit stated above.
"""

from __future__ import annotations

import contextlib
import json
import os
import secrets
import shutil
import signal
import subprocess
import sys
import threading
import time


def fail(msg: str, code: int) -> int:
    print(f"FAIL: {msg}")
    return code


def main() -> int:
    argv = sys.argv[1:]
    expect_server: str | None = None
    expect_tools: list[str] = []
    positional: list[str] = []
    i = 0
    while i < len(argv):
        if argv[i] == "--expect-server" and i + 1 < len(argv):
            expect_server = argv[i + 1]
            i += 2
        elif argv[i] == "--expect-tool" and i + 1 < len(argv):
            expect_tools.append(argv[i + 1])
            i += 2
        else:
            positional.append(argv[i])
            i += 1

    if len(positional) < 2:
        return fail(
            "usage: mcp_smoke.py <.mcp.json> <server> [timeout] [--expect-server NAME] [--expect-tool NAME]...", 2
        )
    cfg_path, server = positional[0], positional[1]
    if len(positional) > 2:
        try:
            budget = int(positional[2])
        except ValueError:
            return fail(f"timeout must be an integer, got {positional[2]!r}", 2)
        if budget <= 0:
            return fail(f"timeout must be positive, got {budget}", 2)
    else:
        budget = 180

    try:
        with open(cfg_path, encoding="utf-8") as fh:
            doc = json.load(fh)
    except FileNotFoundError:
        return fail(f"{cfg_path} does not exist", 3)
    except json.JSONDecodeError as exc:
        return fail(f"{cfg_path} is not valid JSON: {exc}", 4)
    except OSError as exc:
        return fail(f"{cfg_path} could not be read: {exc}", 4)

    servers = doc.get("mcpServers")
    if not isinstance(servers, dict):
        return fail(f"{cfg_path} has no 'mcpServers' object", 5)
    cfg = servers.get(server)
    if not isinstance(cfg, dict):
        return fail(f"{cfg_path} has no '{server}' server entry", 5)

    command = cfg.get("command")
    args = cfg.get("args") or []
    if not isinstance(command, str) or not command:
        return fail(f"'{server}' has no string 'command'", 5)
    if not isinstance(args, list) or not all(isinstance(a, str) for a in args):
        return fail(f"'{server}' has a non-string-list 'args'", 5)

    # A directory is executable, so `-x` alone is vacuous. A bare command name is
    # legitimate and resolves via PATH (e.g. "gitea-mcp-server"), so resolve first
    # and only then insist on a regular file.
    resolved = command if os.path.sep in command else shutil.which(command)
    if resolved is None:
        return fail(f"command not found on PATH: {command}", 6)
    if not os.path.isfile(resolved):
        return fail(f"command is not a regular file: {resolved}", 6)
    if not os.access(resolved, os.X_OK):
        return fail(f"command is not executable: {resolved}", 6)

    # The server runs with the config's directory as cwd, so a relative --project
    # must be validated against THAT, not against the caller's cwd.
    workdir = os.path.dirname(os.path.abspath(cfg_path)) or os.getcwd()
    for idx, a in enumerate(args):
        target = None
        if a == "--project" and idx + 1 < len(args):
            target = args[idx + 1]
        elif a.startswith("--project="):
            target = a.split("=", 1)[1]
        if target:
            probe = target if os.path.isabs(target) else os.path.join(workdir, target)
            if not os.path.exists(probe):
                return fail(f"--project path does not exist: {probe}", 7)

    env = dict(os.environ)
    extra = cfg.get("env") or {}
    if isinstance(extra, dict):
        env.update({k: v for k, v in extra.items() if isinstance(v, str)})

    try:
        proc = subprocess.Popen(
            [resolved, *args],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            env=env,
            cwd=workdir,
            start_new_session=True,  # own process group, so children die with us
        )
    except OSError as exc:
        return fail(f"could not start server: {exc}", 8)

    try:
        pgid = os.getpgid(proc.pid)
    except OSError:
        pgid = None

    # Retain ONLY the reply to the request currently in flight. Keeping every integer
    # id let a server pre-answer id 2 before it was asked, and `wait_for(2)` then
    # accepted an answer to a question never posed — a false green. It also let a
    # duplicate id overwrite an earlier reply, and let `responses` grow without bound.
    # One pending id at a time fixes all three.
    lock = threading.Lock()
    pending: int | None = None
    responses: dict[int, dict] = {}
    drained = threading.Event()

    def reader() -> None:
        try:
            for raw in proc.stdout:  # type: ignore[union-attr]
                line = raw.decode(errors="replace").strip()
                if not line.startswith("{"):
                    continue
                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    continue
                rid = msg.get("id")
                if not isinstance(rid, int):
                    continue  # a notification: nothing to retain
                with lock:
                    # unsolicited, or a second answer to an already-answered id
                    if rid != pending or rid in responses:
                        continue
                    responses[rid] = msg
        finally:
            drained.set()

    threading.Thread(target=reader, daemon=True).start()

    def send(payload: dict) -> bool:
        """False when the pipe is gone — an instantly-exiting server is the #777
        failure mode, so it must produce a diagnostic, not a BrokenPipeError."""
        try:
            proc.stdin.write((json.dumps(payload) + "\n").encode())  # type: ignore[union-attr]
            proc.stdin.flush()  # type: ignore[union-attr]
            return True
        except (BrokenPipeError, OSError, ValueError):
            return False

    def expect(req_id: int, payload: dict, deadline: float) -> dict | None:
        """Register the id BEFORE sending, so a reply cannot arrive unregistered."""
        nonlocal pending
        with lock:
            pending = req_id
        if not send(payload):
            return None
        return wait_for(req_id, deadline)

    def wait_for(req_id: int, deadline: float) -> dict | None:
        while time.time() < deadline:
            if req_id in responses:
                return responses[req_id]
            # Only conclude "no answer" once the process is gone AND stdout is fully
            # drained; otherwise a reply already in the pipe is reported as a no-show.
            if proc.poll() is not None and drained.wait(timeout=2):
                return responses.get(req_id)
            time.sleep(0.25)
        return responses.get(req_id)

    def cleanup() -> None:
        # `dotnet run` execs a CHILD (csharp-lsp-mcp), so the leader exiting on
        # SIGTERM says nothing about the descendant. Always follow up with SIGKILL to
        # the saved group: a stale server surviving a probe is exactly the litter
        # observed accumulating at start-up.
        if pgid is not None:
            for sig in (signal.SIGTERM, signal.SIGKILL):
                try:
                    os.killpg(pgid, sig)
                except OSError:
                    break  # no group members left
                with contextlib.suppress(subprocess.TimeoutExpired):
                    proc.wait(timeout=5)
                time.sleep(0.2)
        else:
            for sig in (signal.SIGTERM, signal.SIGKILL):
                try:
                    proc.send_signal(sig)
                    proc.wait(timeout=5)
                    break
                except (OSError, subprocess.TimeoutExpired):
                    continue

    deadline = time.time() + budget
    try:
        # Random ids close the residual pre-answer race: holding a lock across send()
        # cannot reject a frame the server emitted BEFORE the request, but a server
        # cannot pre-answer an id it cannot guess.
        id_init = secrets.randbelow(2**31 - 1000) + 1000
        id_tools = secrets.randbelow(2**31 - 1000) + 1000
        while id_tools == id_init:
            id_tools = secrets.randbelow(2**31 - 1000) + 1000
        init = expect(
            id_init,
            {
                "jsonrpc": "2.0",
                "id": id_init,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {"name": "mcp-smoke", "version": "0"},
                },
            },
            deadline,
        )
        if init is None:
            return fail(f"no 'initialize' response within {budget}s (server did not start)", 9)
        if "error" in init:
            return fail(f"initialize returned an error: {json.dumps(init['error'])[:300]}", 9)
        result = init.get("result")
        if not isinstance(result, dict):
            return fail("initialize response has no 'result' object (malformed)", 12)
        info = result.get("serverInfo")
        if not isinstance(info, dict) or not isinstance(info.get("name"), str):
            return fail("initialize result has no 'serverInfo.name' string (malformed)", 12)
        actual = info["name"]
        if expect_server is not None and actual != expect_server:
            return fail(f"wrong server: expected '{expect_server}', got '{actual}'", 13)

        send({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}})
        listed = expect(id_tools, {"jsonrpc": "2.0", "id": id_tools, "method": "tools/list", "params": {}}, deadline)
        if listed is None:
            return fail(f"no 'tools/list' response within {budget}s", 10)
        lresult = listed.get("result")
        if not isinstance(lresult, dict):
            return fail("tools/list response has no 'result' object (malformed)", 12)
        tools = lresult.get("tools")
        if not isinstance(tools, list):
            return fail("tools/list 'tools' is not a list (malformed)", 12)
        names = {t.get("name") for t in tools if isinstance(t, dict) and isinstance(t.get("name"), str)}
        if not names:
            return fail("server started but exposes zero well-formed tools", 11)
        missing = [t for t in expect_tools if t not in names]
        if missing:
            return fail(f"server '{actual}' is missing expected tool(s): {', '.join(missing)}", 14)
        print(f"OK: {server} -> {actual} {info.get('version', '')}, {len(names)} tools")
        return 0
    finally:
        cleanup()


if __name__ == "__main__":
    sys.exit(main())