#!/usr/bin/env python3
"""Opus 5.5 zero-shot judge (the no-training rival) on the same evidence a probe gets: the policy, the customer's
messages, the database changes and the relevant START records. No agent messages, no labels, no task text.
Effort high, adaptive thinking, one sample. Resumable (skips sim_ids already in --out).

  python judge.py --ids data/judge_sample.txt --out runs/judge.jsonl [--workers 16] [--digest FILE --tag digest]
"""
import argparse, concurrent.futures as cf, gzip, json, os, re, sys, threading
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import llm  # noqa: E402
import render  # noqa: E402
from state import policy, start_db  # noqa: E402

SYSTEM = """You audit the work of an AI customer-service agent after a chat has ended. The agent had to follow the
domain policy and do what the customer validly asked for, using tools that read and write the database. You see
the policy, the customer's messages (not the agent's), every database change the agent made, and the relevant
database records as they were before the chat. Decide whether the agent's work is WRONG: the end state violates
the policy, misses or gets wrong something the customer validly requested, or contains a change the customer did
not request or confirm. If the work is consistent with the policy and the customer's requests, it is OK."""

LOCK = threading.Lock()


def prompt(it, digest=None):
    d = it["domain"]
    s = [f"## Policy ({d})\n```\n{policy(d)}\n```"]
    if digest:
        s.append(f"## Notes on common failures in this domain (compiled from past graded chats)\n{digest}")
    s += [f"## Customer messages\n```\n" + "\n---\n".join(it["user_messages"]) + "\n```",
          f"## Database changes made during the chat (START -> END)\n```\n{render.changes_text(it['changes'])}\n```",
          "## Relevant records before the chat (START)\n```\n" +
          render.relevant_start(d, start_db(d), "\n".join(it["user_messages"]), it["changes"]) + "\n```",
          'Answer with a JSON object on the last line: {"verdict": "WRONG" or "OK", "reason": "<one sentence>"}']
    return "\n\n".join(s)


def parse(txt):
    m = re.findall(r'"verdict"\s*:\s*"(WRONG|OK)"', txt or "", re.I)
    return m[-1].upper() if m else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ids", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--digest", default=None, help="json {domain: digest text}")
    ap.add_argument("--tag", default="zeroshot")
    a = ap.parse_args()
    ids = [l.strip() for l in open(a.ids) if l.strip()]
    want = set(ids)
    items = {}
    for l in gzip.open(os.path.join(HERE, "data", "items.jsonl.gz"), "rt"):
        r = json.loads(l)
        if r["sim_id"] in want:
            items[r["sim_id"]] = r
    dig = json.load(open(a.digest)) if a.digest else {}
    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 ids if s not in done]
    print("todo", len(todo), "done", len(done), flush=True)
    calls = os.path.join(HERE, "runs", f"judge_calls_{a.tag}.jsonl")

    def one(sid):
        it = items[sid]
        try:
            txt = llm.call(prompt(it, dig.get(it["domain"])), system=SYSTEM, log_path=calls, purpose=f"judge_{a.tag}",
                           meta={"sim_id": sid}, max_tokens=16000)
            v, err = parse(txt), None
        except Exception as e:  # noqa: BLE001
            txt, v, err = "", None, str(e)[:300]
        row = {"sim_id": sid, "tag": a.tag, "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 r in ex.map(one, todo):
            n += 1
            if n % 50 == 0:
                print(n, flush=True)
    print("done", n)


if __name__ == "__main__":
    main()