"""同じ候補値を、見せる情報だけ変えて何度も採点する。使い方: python3 -I run_judge.py <model> [並列数]"""
import json, subprocess, sys, itertools, os, concurrent.futures as cf

CUT = 1000
EMPTY_DIR = os.environ.get("EMPTY_DIR", "/tmp/llmj/run")  # CLAUDE.md等が無い空フォルダ
SYSTEM = ("あなたはAIエージェントのツール呼び出しを検証する判定器です。"
          "与えられた情報だけを根拠に、提案された引数の値が正しいかを判定します。"
          "1行目には ACCEPT（正しい）、REJECT（誤り）、UNSURE（与えられた情報では判断できない）のうち、いずれか1語だけを書きます。"
          "1行目にそれ以外の文字を入れてはいけません。2行目に理由を1文で書いてください。")

def prompt(case, placement, ref, cand):
    ret = case["ret_visible"] if placement == "visible" else case["ret_buried"]
    shown = ret[:CUT]
    if placement == "buried_notice":
        shown += "\n…（長いため、ここから先は省略されています）"
    value = case[cand]
    parts = [f"## ユーザーの依頼\n{case['request']}",
             f"## これまでのやり取り\nエージェントの呼び出し: {case['search_call']}\nツールの返り値:\n{shown}",
             f"## 関数\n{case['fn']}({case['param']}: string) — {case['pdesc']}"]
    if ref != "none":
        # ref_wrong: 参照が誤値を指す。候補が誤値なら候補と同じ値、正解候補なら別の誤値（先頭エントリ）を指す
        rv = case["correct"] if ref == "ref_correct" else case["wrong_hidden" if cand == "wrong_hidden" else "wrong"]
        parts.append(f"## 参照（期待値）\n{case['param']} = {rv}\n根拠: 依頼で指定された名称に対応する、検索結果のエントリのID。")
    parts.append(f"## 提案された呼び出し\n{case['fn']}({case['param']}=\"{value}\")\n\nこの {case['param']} の値は正しいですか？")
    return "\n\n".join(parts)

def call(model, text):
    env = dict(os.environ, CLAUDE_CODE_DISABLE_AUTO_MEMORY="1")
    cmd = ["claude", "-p", "--model", model, "--tools", "", "--setting-sources", "", "--strict-mcp-config",
           "--disable-slash-commands", "--no-session-persistence", "--system-prompt", SYSTEM, "--output-format", "json"]
    for _ in range(3):
        r = subprocess.run(cmd, input=text, capture_output=True, text=True, env=env, cwd=EMPTY_DIR)
        try:
            j = json.loads(r.stdout)
            return j["result"], list(j.get("modelUsage", {}))
        except Exception:
            continue
    return "ERROR", []

def main():
    model = sys.argv[1]; workers = int(sys.argv[2]) if len(sys.argv) > 2 else 8
    cases = json.load(open("cases.json"))
    jobs = [(c["case_id"], pl, rf, cd) for c in cases
            for pl, rf, cd in itertools.product(["visible", "buried", "buried_notice"], ["none", "ref_correct", "ref_wrong"], ["correct", "wrong", "wrong_hidden"])
            if not (cd == "wrong_hidden" and pl == "visible")]
    byid = {c["case_id"]: c for c in cases}
    out = f"results_{model}.jsonl"
    done = set()
    if os.path.exists(out):
        for l in open(out):
            r = json.loads(l); done.add((r["case_id"], r["placement"], r["ref"], r["cand"]))
    todo = [j for j in jobs if j not in done]
    def work(j):
        cid, pl, rf, cd = j
        text, used = call(model, prompt(byid[cid], pl, rf, cd))
        return dict(case_id=cid, placement=pl, ref=rf, cand=cd, model_used=used, output=text)
    with cf.ThreadPoolExecutor(workers) as ex, open(out, "a") as f:
        for r in ex.map(work, todo):
            f.write(json.dumps(r, ensure_ascii=False) + "\n"); f.flush()
    print("done", len(todo))

if __name__ == "__main__":
    main()
