#!/usr/bin/env python3
"""Протокольный зонд MCP-сервера (stdio, JSON-RPC 2.0).

Запускает сервер как настоящий MCP-клиент и проверяет, что он отвечает по протоколу,
а не «выглядит работающим». Использование:

  python3 scripts/mcp_probe.py --cmd ./memo-mcp/build/install/memo-mcp/bin/memo-mcp
  python3 scripts/mcp_probe.py --cmd ./... --call memo_search \
      --args '{"path":"/root/WORK/memo-e2e/infra","query":"домены","k":3}'

Выход 0 — протокол в порядке; 1 — есть ошибки (печатаются).
"""

import argparse
import json
import shlex
import subprocess
import sys


def rpc(proc, method, params=None, notify=False, rid=None):
    msg = {"jsonrpc": "2.0", "method": method}
    if params is not None:
        msg["params"] = params
    if not notify:
        msg["id"] = rid if rid is not None else 1
    proc.stdin.write(json.dumps(msg) + "\n")
    proc.stdin.flush()
    if notify:
        return None
    line = proc.stdout.readline()
    if not line:
        return None
    try:
        return json.loads(line)
    except json.JSONDecodeError:
        return {"_raw": line.strip()}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cmd", required=True, help="команда запуска MCP-сервера (stdio)")
    ap.add_argument("--call", help="имя инструмента для tools/call")
    ap.add_argument("--args", default="{}", help="JSON-аргументы инструмента")
    ap.add_argument("--timeout", type=float, default=60.0)
    a = ap.parse_args()

    errors: list[str] = []
    proc = subprocess.Popen(
        shlex.split(a.cmd),
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, bufsize=1,
    )
    try:
        init = rpc(proc, "initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "mcp_probe", "version": "1"},
        }, rid=1)
        if not init or "result" not in init:
            errors.append(f"initialize: нет result (ответ {init})")
        else:
            info = init["result"].get("serverInfo", {})
            print(f"OK  initialize: {info.get('name')} {info.get('version')}")

        rpc(proc, "notifications/initialized", {}, notify=True)

        tools = rpc(proc, "tools/list", {}, rid=2)
        names: list[str] = []
        if not tools or "result" not in tools:
            errors.append(f"tools/list: нет result (ответ {tools})")
        else:
            names = [t["name"] for t in tools["result"].get("tools", [])]
            if not names:
                errors.append("tools/list: пустой список инструментов")
            else:
                print(f"OK  tools/list: {len(names)} -> {', '.join(names)}")

        if a.call:
            if names and a.call not in names:
                errors.append(f"tools/call: инструмента {a.call} нет в tools/list")
            else:
                res = rpc(proc, "tools/call",
                          {"name": a.call, "arguments": json.loads(a.args)}, rid=3)
                if not res or "result" not in res:
                    errors.append(f"tools/call: нет result (ответ {res})")
                elif res["result"].get("isError"):
                    errors.append(f"tools/call: isError=true ({res['result']})")
                else:
                    content = res["result"].get("content") or []
                    text = content[0].get("text", "") if content else ""
                    if not text.strip():
                        errors.append("tools/call: пустой content[0].text")
                    else:
                        print(f"OK  tools/call {a.call}: {len(text)} символов")
                        print("---- первые 600 символов ----")
                        print(text[:600])
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
        err = proc.stderr.read()
        if err.strip():
            print("---- stderr сервера ----")
            print(err[:1000])

    if errors:
        print("\nОШИБКИ:")
        for e in errors:
            print(" -", e)
        return 1
    print("\nПРОТОКОЛ В ПОРЯДКЕ")
    return 0


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