#!/usr/bin/env python3
"""AI Quant Researcher Benchmark v1 - MCP call helper (sealed).
Usage: python3 /tmp/claude-1000/bench/mcp_bench.py <RUN_ID> <tool_name> '<json args>'   |  <RUN_ID> tools/list
Only the research tape is allowed; every call is logged; budget 80 calls."""
import json, subprocess, sys, time, os
APP = "/home/chris/.local/share/strategyarena-candidates/20261007-exact-cuda-230/apps/strategyarena-tauri"
ALLOWED = {"ohlcv-sha256-9f8ce49cc04885b6"}
BUDGET = 80
run, name = sys.argv[1], sys.argv[2]
args = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
logp = f"/tmp/claude-1000/bench/log_{run}.jsonl"
n = sum(1 for _ in open(logp)) if os.path.exists(logp) else 0
def log(rec):
    with open(logp, "a") as f: f.write(json.dumps(rec) + "\n")
if n >= BUDGET:
    print(json.dumps({"ok": False, "error": "BENCHMARK_CALL_BUDGET_EXHAUSTED", "budget": BUDGET})); sys.exit(0)
blob = json.dumps(args)
import re
for ds in set(re.findall(r"ohlcv-sha256-[0-9a-f]{16}", blob)):
    if ds not in ALLOWED:
        log({"t": time.time(), "tool": name, "blocked": ds})
        print(json.dumps({"ok": False, "error": "BENCHMARK_SEALED_DATASET", "message": f"{ds} is sealed for evaluation. Only ohlcv-sha256-9f8ce49cc04885b6 (BTC 4h 2020-10..2023-09) is allowed."})); sys.exit(0)
msgs = [{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": f"bench-{run}", "version": "1"}}}]
msgs.append({"jsonrpc": "2.0", "id": 2, "method": "tools/list"} if name == "tools/list" else {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": name, "arguments": args}})
p = subprocess.run(["node", "scripts/strategyarena-agent-mcp-stdio.mjs"], cwd=APP, input="\n".join(map(json.dumps, msgs)) + "\n", capture_output=True, text=True, timeout=900)
lines = [l for l in p.stdout.strip().splitlines() if l.strip()]
d = json.loads(lines[-1]) if lines else {"error": p.stderr[-1000:]}
res = d.get("result", d)
if name == "tools/list":
    out = {"tools": [{"name": t["name"], "description": t["description"]} for t in res.get("tools", [])]}
    if len(sys.argv) > 3: want = set(json.loads(sys.argv[3])); out = {"tools": [t for t in res.get("tools", []) if t["name"] in want]}
else:
    out = res.get("structuredContent", res)
    if name == "list_sandbox_datasets" and isinstance(out, dict) and "datasets" in out:
        out["datasets"] = [json.load(open("/tmp/claude-1000/bench/allowed_dataset.json"))]; out["total"] = 1; out["nextOffset"] = None
ok = isinstance(out, dict) and out.get("ok", True) is not False and "error" not in out
log({"t": time.time(), "tool": name, "ok": ok, "err": (out.get("error") if isinstance(out, dict) else None)})
s = json.dumps(out, indent=1)
outp = f"/tmp/claude-1000/bench/{run}/last_{name.replace('/','_')}.json"; os.makedirs(os.path.dirname(outp), exist_ok=True); open(outp, "w").write(s)
print(s[:40000] + ("\n...[truncated; full output in " + outp + "]" if len(s) > 40000 else ""))