#!/usr/bin/env python3
"""tau2-agreement H2b: the INFORMATION-MATCHED rival to consensus verification on new tau2 tasks.

Opus 5.5 judges the attempt exactly as capital-tau2's zero-shot judge does (same system prompt, same evidence: policy,
customer messages, database changes, relevant START records; adaptive thinking, effort high, one sample), PLUS the
database changes of the SAME K=4 peer attempts that CV-K4's first seeded draw uses (analyze_cv_tau2.py: draw 0 of
random.Random(f"20261009|{sim_id}")). Peers are labelled unverified. No label, no peer outcome, no agent messages.
Resumable (skips sim_ids already in --out). Runs through a model proxy.

  python peer_judge.py --data DIR --ct2 CAPITAL_TAU2_CODE_DIR --out peer.jsonl [--workers 12] [--limit N]
"""
from __future__ import annotations

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

SEED = 20261009
K = 4
LOCK = threading.Lock()
PEER_NOTE = """## Other attempts at the same customer request (unverified)
The same customer request was also handled in {k} other independent chats, by other agents. Below are the
database changes each of those attempts made. Their correctness is UNKNOWN: any or all of them may be wrong. They
are evidence about what the request involves, not ground truth. The customer's wording may differ between chats."""


def draw0_panel(sim_id, others):
    rnd = random.Random(f"{SEED}|{sim_id}")
    # analyze_cv_tau2.py draws K=2 first (50 draws), then K=3, K=4 ... from the same generator: replay exactly
    for kk in (2, 3):
        for _ in range(50):
            rnd.sample(others, min(kk, len(others)))
    return rnd.sample(others, min(K, len(others)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="dir with items.jsonl.gz, judge_zeroshot.jsonl, fold2_judge_zeroshot.jsonl")
    ap.add_argument("--ct2", required=True, help="companion/capital-tau2 code dir (render, state, llm, judge)")
    ap.add_argument("--state-home", required=True, help="dir containing data/start_{domain}.json.gz and data/policy_*.md")
    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  # capital-tau2 zero-shot judge: SYSTEM, prompt(), parse()
    import llm
    import render
    import state
    state.HERE = a.state_home
    items = {}
    by_task = {}
    for l in gzip.open(os.path.join(a.data, "items.jsonl.gz"), "rt"):
        r = json.loads(l)
        items[r["sim_id"]] = r
        by_task.setdefault(f"{r['domain']}:{r['task_id']}", []).append(r["sim_id"])
    judged = []
    for f in ("judge_zeroshot.jsonl", "fold2_judge_zeroshot.jsonl"):
        judged += [json.loads(l)["sim_id"] for l in open(os.path.join(a.data, f))]
    if a.limit:
        judged = judged[: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 judged 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]
        task = f"{it['domain']}:{it['task_id']}"
        others = [x for x in by_task[task] if x != sid]
        panel = draw0_panel(sid, others)
        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="judge_peer", 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, "tag": "peer4", "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)


if __name__ == "__main__":
    main()