#!/usr/bin/env python3
"""Drive the fluxgit-mcp-sidecar over stdio (newline-delimited JSON-RPC 2.0).

Usage: mcp-drive.py <sidecar-binary> <repo-path> <outdir>
Calls initialize, tools/list, then repo.brief and (if available) repo.scope
and repo.conflictPreflight; writes each raw response JSON to <outdir>.
"""
import json, subprocess, sys, os

BIN, REPO, OUTDIR = sys.argv[1], sys.argv[2], sys.argv[3]
PROTOCOL_VERSION = "2024-11-05"
os.makedirs(OUTDIR, exist_ok=True)
BENCH_ROOT = os.path.realpath(os.path.dirname(REPO))

def normalize(text):
    """Remove machine-specific paths before artifacts are written or counted."""
    candidates = {os.path.dirname(REPO), BENCH_ROOT}
    for candidate in sorted(candidates, key=len, reverse=True):
        text = text.replace(candidate, "$BENCH_ROOT")
    return text

proc = subprocess.Popen([BIN], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE)

def send(obj):
    proc.stdin.write((json.dumps(obj) + "\n").encode())
    proc.stdin.flush()

def read_frame():
    """Read current newline-delimited output or a legacy Content-Length frame."""
    first = proc.stdout.readline()
    if not first:
        raise RuntimeError("sidecar closed stdout; stderr: " +
                           proc.stderr.read().decode()[:2000])
    if not first.lower().startswith(b"content-length:"):
        return first.rstrip(b"\r\n").decode()

    length = int(first.split(b":", 1)[1].strip())
    while True:
        header = proc.stdout.readline()
        if not header:
            raise RuntimeError("truncated legacy frame headers")
        if header in (b"\n", b"\r\n"):
            break
    body = proc.stdout.read(length)
    if len(body) != length:
        raise RuntimeError("truncated legacy frame body")
    return body.decode()

def recv(expect_id):
    while True:
        line = read_frame()
        msg = json.loads(line)
        if msg.get("id") == expect_id:
            return msg, line

rid = 0
def call(method, params):
    global rid
    rid += 1
    send({"jsonrpc": "2.0", "id": rid, "method": method, "params": params})
    return recv(rid)

# 1. initialize
init, _ = call("initialize", {
    "protocolVersion": PROTOCOL_VERSION,
    "capabilities": {},
    "clientInfo": {"name": "token-bench", "version": "1.0"},
})
send({"jsonrpc": "2.0", "method": "notifications/initialized"})
print("initialized:", init["result"]["serverInfo"], file=sys.stderr)

# 2. tools/list (for schemas; not counted in the benchmark totals)
tl, raw = call("tools/list", {})
with open(os.path.join(OUTDIR, "tools-list.json"), "w") as f:
    f.write(normalize(raw))
names = [t["name"] for t in tl["result"]["tools"]]
print("tools:", len(names), file=sys.stderr)

def tool_call(name, args, outfile):
    msg, raw = call("tools/call", {"name": name, "arguments": args})
    with open(os.path.join(OUTDIR, outfile), "w") as f:
        f.write(normalize(raw))
    ok = "error" not in msg and not msg.get("result", {}).get("isError", False)
    print(f"{name}: {'OK' if ok else 'ERROR'} raw_bytes={len(raw.encode())}", file=sys.stderr)
    return msg, raw

tool_call("repo.brief", {"repoPath": REPO}, "repo-brief.json")
if "repo.scope" in names:
    tool_call("repo.scope", {"repoPath": REPO, "path": "src/billing", "churnDays": 365}, "repo-scope.json")
if "repo.conflictPreflight" in names:
    tool_call("repo.conflictPreflight",
              {"repoPath": REPO, "currentRef": "feature/usage-metering", "targetRef": "main"},
              "conflict-preflight.json")

proc.stdin.close()
proc.wait(timeout=10)