#!/usr/bin/env python3
"""tau2-agreement H4 judges on the tau2 TEST-split judge sample (prereg amendment 3). Runs via a model proxy.
  --mode zeroshot : capital-tau2 judge.py prompt and system verbatim (policy, customer messages, changes, START records)
  --mode peer     : the same + the database changes of the d0 SAME-MODEL panel (h4_common.panels(..., "same")[0]),
                    labelled unverified (peer_judge.py wording)
Opus 5.5, adaptive thinking, effort high, one sample. Resumable.

  python h4_judges.py --mode zeroshot|peer --data DIR --ct2 CT2_DIR --state-home DIR --out OUT.jsonl [--workers 12]
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import json
import os
import re
import sys
import threading

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, ".."))
import h4_common as C  # noqa: E402
from peer_judge import PEER_NOTE  # noqa: E402

LOCK = threading.Lock()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", required=True, choices=["zeroshot", "peer"])
    ap.add_argument("--data", required=True)
    ap.add_argument("--ct2", required=True)
    ap.add_argument("--state-home", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()
    sys.path.insert(0, a.ct2)
    import judge as zj
    import llm
    import render
    import state
    state.HERE = a.state_home
    idx, items, sample, by_task = C.load(a.data)
    if a.limit:
        sample = sample[:a.limit]
    done = set()
    if os.path.exists(a.out):
        for l in open(a.out):
            r = json.loads(l)
            if r.get("verdict"):
                done.add(r["sim_id"])
    todo = [s for s in sample if s not in done]
    print("todo", len(todo), "done", len(done), flush=True)
    calls = a.out.replace(".jsonl", "_calls.jsonl")

    def one(sid):
        it = items[sid]
        panel = []
        if a.mode == "zeroshot":
            ptxt = zj.prompt(it)
        else:
            panel = C.panels(sid, C.pool(sid, idx, items, by_task, True), "same")[0]
            base = zj.prompt(it)
            head, tail = base.rsplit("\n\nAnswer with a JSON object", 1)
            peers = [PEER_NOTE.format(k=len(panel))]
            for i, p in enumerate(panel, 1):
                peers.append(f"### Attempt {i}: database changes (START -> END)\n```\n"
                             f"{render.changes_text(items[p]['changes'], max_chars=4000)}\n```")
            ptxt = head + "\n\n" + "\n\n".join(peers) + "\n\nAnswer with a JSON object" + tail
        try:
            txt = llm.call(ptxt, system=zj.SYSTEM, log_path=calls, purpose=f"h4_{a.mode}", meta={"sim_id": sid},
                           max_tokens=16000)
            v, err = zj.parse(txt), None
        except Exception as e:  # noqa: BLE001
            txt, v, err = "", None, str(e)[:300]
        row = {"sim_id": sid, "mode": a.mode, "panel": panel, "verdict": v, "flag": (v == "WRONG") if v else None,
               "error": err, "reason": (re.findall(r'"reason"\s*:\s*"([^"]*)"', txt) or [""])[-1][:400]}
        with LOCK, open(a.out, "a") as f:
            f.write(json.dumps(row) + "\n")
        return row

    n = 0
    with cf.ThreadPoolExecutor(a.workers) as ex:
        for _ in ex.map(one, todo):
            n += 1
            if n % 50 == 0:
                print(n, flush=True)
    print("done", n, flush=True)


if __name__ == "__main__":
    main()