"""使い方: python3 run_probe.py <model> <case,case,...> <出力dir>   (FOLLOW=1 でリンク先を読ませる / NORES=1 でresources capabilityを外す)
Run each probe tool through `claude -p` and record what reached the model (tool_result in stream-json)."""
import json, os, subprocess, sys, tempfile
HERE = os.path.dirname(os.path.abspath(__file__))
WORK = tempfile.mkdtemp(prefix="probe-work-")  # 空のcwd。ファイルを置かない

CASES = sys.argv[2].split(",")
MODEL = sys.argv[1]
OUT = os.path.abspath(sys.argv[3])
os.makedirs(OUT, exist_ok=True)
PROMPT_DEFAULT = ("probe MCPサーバーのツール `{t}` を1回だけ呼んでください。"
          "結果に見える CODE: で始まる値や、コード風の文字列を、見えたものすべて、そのまま書き出してください。"
          "何も見えなければ「なし」と書いてください。結果の外から推測した値は書かないでください。")
PROMPT_FOLLOW = ("probe MCPサーバーのツール `{t}` を1回呼び、返ってきたリンク先の本文に書かれている CODE: の値を取得して、そのまま書いてください。"
                 "使えるツールは何を使ってもかまいません。取得できなければ、できなかった理由を書いてください。")
PROMPT = PROMPT_FOLLOW if os.environ.get("FOLLOW") else PROMPT_DEFAULT
for case in CASES:
    summary = f"{OUT}/{case}.summary.json"
    if os.path.exists(summary): os.remove(summary)
    log = f"{OUT}/{case}.server.log"
    if os.path.exists(log): os.remove(log)
    env = {**os.environ, "PROBE_LOG": log}
    penv = {"PROBE_LOG": log, **({"NORES": "1"} if os.environ.get("NORES") else {})}
    cfg = {"mcpServers": {"probe": {"command": "python3", "args": [os.path.join(HERE, "probe_server.py")], "env": penv}}}
    cfgp = f"{OUT}/{case}.mcp.json"; json.dump(cfg, open(cfgp, "w"))
    cmd = ["claude", "-p", PROMPT.format(t=case), "--mcp-config", cfgp, "--strict-mcp-config",
           "--tools", "ListMcpResourcesTool,ReadMcpResourceTool",
           "--allowedTools", "mcp__probe,ListMcpResourcesTool,ReadMcpResourceTool",
           "--model", MODEL, "--output-format", "stream-json", "--verbose", "--no-session-persistence"]
    p = subprocess.run(cmd, cwd=WORK, env=env, stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=300)
    open(f"{OUT}/{case}.stream.jsonl", "w").write(p.stdout)
    if p.returncode != 0:
        print(p.stderr, file=sys.stderr, end="")
        sys.exit(p.returncode)
    calls, results, final = [], [], ""
    for line in p.stdout.splitlines():
        try: ev = json.loads(line)
        except Exception: continue
        if ev.get("type") == "assistant":
            for b in ev["message"]["content"]:
                if b.get("type") == "tool_use": calls.append({"name": b["name"], "input": b["input"]})
        if ev.get("type") == "user":
            c = ev["message"]["content"]
            if isinstance(c, list):
                for b in c:
                    if b.get("type") == "tool_result": results.append(b.get("content"))
        if ev.get("type") == "result": final = ev.get("result", "")
    srv = [json.loads(l) for l in open(log)] if os.path.exists(log) else []
    rec = {"case": case, "model": MODEL, "tool_calls": calls, "tool_results_to_model": results, "final": final,
           "server_events": [e for e in srv if "initialize_protocolVersion" not in e]}
    json.dump(rec, open(summary, "w"), ensure_ascii=False, indent=1)
    print("==", case); print(" calls:", [c["name"] + str(c["input"]) for c in calls]); print(" to_model:", json.dumps(results, ensure_ascii=False)[:600]); print(" server:", rec["server_events"]); print(" final:", final[:300])