"""Let a panel of judges argue about the questions they disagree on.

Majority vote treats a 2-1 split as settled. It isn't: a panel with one lenient
and one strict member lets the third member decide every split, and on genuine
equivalence questions the majority is confidently wrong (three judges can all
miss that two matrices differ by a unimodular factor). Deliberation at least
puts the reasoning on the table.

Each round every judge sees the current positions and may revise. Rounds stop
early on consensus. Positions are anonymised by default so a judge cannot defer
to a name it considers authoritative.

    python deliberate.py judged_A.json judged_B.json judged_C.json --dry-run
    python deliberate.py judged_A.json judged_B.json judged_C.json --rounds 5
    python deliberate.py ... --show-names        # reveal who said what

Writes deliberation.json: every round, every position, for every question.
Resumable -- questions already settled in the transcript are skipped.
"""

import argparse
import asyncio
import json
import os
import re
from pathlib import Path
from typing import Literal

from openai import AsyncOpenAI
from pydantic import BaseModel
from tqdm.asyncio import tqdm_asyncio

from text_only_share import hf_token, load_env

SYSTEM = """You are grading one answer against an answer key, together with other
graders. You will see the current position of each grader, including your own.

Reconsider from scratch. Change your verdict only if an argument you now see is
sound -- never because you are outnumbered. Holding a correct minority position
is the right outcome; agreeing to end the disagreement is not.

Pay particular attention to whether the two answers are equivalent under a
transformation the question itself permits: rotation, reflection, reordering,
algebraic identity, a unimodular factor, a change of units or notation. A
difference in form is not a difference in content.

Reply in exactly this form, nothing else:

reasoning: <one paragraph>
correct: yes

(use "correct: no" if the answer does not match the key)"""


# Verbatim from hle/hle_eval/run_judge_results.py. The opening verdicts were
# formed under this standard, so without it the graders are re-arguing to a
# rule they cannot see -- and the margin-of-error clause in particular decides
# cases like 0.85 against a key of 0.84, where one grader invoked it and then
# abandoned it with nothing to appeal to.
HLE_CRITERIA = """reasoning: Explain why the extracted_final_answer is correct or \
incorrect based on [correct_answer], focusing only on if there are meaningful \
differences between [correct_answer] and the extracted_final_answer. Do not comment \
on any background to the problem, do not attempt to solve the problem, do not argue \
for any answer different than [correct_answer], focus only on whether the answers match.

correct: Answer 'yes' if extracted_final_answer matches the [correct_answer] given \
above, or is within a small margin of error for numerical problems. Answer 'no' \
otherwise, i.e. if there if there is any inconsistency, ambiguity, non-equivalency, \
or if the extracted answer is incorrect."""


FAILS = []


class Turn(BaseModel):
    reasoning: str
    correct: Literal["yes", "no"]


def save(path, records):
    """Atomic write, so an interrupt lands either side of it and never inside.

    write_text truncates the file before it writes; Ctrl+C in that window leaves
    a partial JSON that will not load, destroying every completed question.
    """
    tmp = str(path) + ".tmp"
    with open(tmp, "w") as f:
        json.dump(records, f, indent=2)
    os.replace(tmp, path)


def name_of(path):
    m = re.search(r"__(.+)\.json$", Path(path).name)
    return m.group(1) if m else Path(path).stem


def solver_of(path):
    """Whose answers were graded. Recorded so a transcript is self-describing:
    the same question id exists under every solver, so without this a consumer
    cannot tell which run a verdict belongs to."""
    m = re.search(r"^judged_(.+?)__", Path(path).name)
    return m.group(1) if m else "?"


def vote(entry):
    return 1 if "yes" in entry["judge_response"]["correct"] else 0


YES = {"yes", "correct", "true", "match", "matches"}


YES = {"yes", "correct", "true", "match", "matches"}


def parse_prose(text):
    """Fallback for judges whose providers will not honour a schema.

    Models answer with 'correct: yes', '**Verdict:** Correct', '"correct": "no"'
    and similar; the verdict word matters, the decoration does not.
    """
    m = re.search(r'[*"#\s]*(?:verdict|correct)[*"]*\s*:\s*[*"\s]*'
                  r'(yes|no|correct|incorrect|true|false)\b', text, re.I)
    if not m:
        return None
    word = m.group(1).lower()
    r = re.search(r'[*"#\s]*reasoning[*"]*\s*[:\n]\s*[*"\s]*(.*?)'
                  r'(?=\n\s*[*"#\s]*(?:verdict|correct)[*"]*\s*:|\Z)',
                  text, re.S | re.I)
    body = r.group(1) if r else text
    # whatever we captured, a trailing verdict line is not reasoning
    body = re.sub(r'\n\s*[*"#\s]*(?:verdict|correct)[*"]*\s*:.*$', "", body,
                  flags=re.S | re.I)
    return {"correct": "yes" if word in YES else "no",
            "reasoning": body.strip()[:1200]}


def positions_block(positions, me, show_names, skip_me=False):
    lines = []
    for i, (who, pos) in enumerate(positions.items()):
        if skip_me and who == me:
            continue        # in a chain their own view is the assistant turn above
        label = who if show_names else (f"grader {chr(65+i)}" + (" (you)" if who == me else ""))
        lines.append(f"--- {label}: {'CORRECT' if pos['v'] else 'INCORRECT'}\n{pos['r']}")
    return "\n\n".join(lines)


def build(question, response, key, positions, me, show_names):
    """Snapshot prompt: no prior rounds, only where everyone stands now."""
    return [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content":
            f"[question]\n{question}\n\n[answer being graded]\n{response}\n\n"
            f"[answer key]\n{key}\n\n"
            f"[the standard these verdicts were formed under]\n{HLE_CRITERIA}\n\n"
            f"[current positions]\n"
            + positions_block(positions, me, show_names)},
    ]


def extend_chain(chain, ctx, positions, me, show_names, my_last):
    """Append one exchange to a grader's own running conversation.

    Its previous verdict goes back as an assistant turn rather than as another
    line of context, so the model is looking at something it said. That is the
    whole point of testing history: it is what creates consistency pressure, and
    what a snapshot prompt deliberately removes.
    """
    if not chain:
        chain.append({"role": "system", "content": SYSTEM})
        chain.append({"role": "user", "content":
                      f"[question]\n{ctx['q']}\n\n[answer being graded]\n{ctx['resp']}\n\n"
                      f"[answer key]\n{ctx['key']}\n\n"
                      f"[the standard these verdicts were formed under]\n{HLE_CRITERIA}\n\n"
                      f"[current positions]\n"
                      + positions_block(positions, me, show_names)})
        return chain
    if my_last is not None:
        chain.append({"role": "assistant", "content":
                      f"{'CORRECT' if my_last['v'] else 'INCORRECT'}\n{my_last['r']}"})
    others = positions_block(positions, me, show_names, skip_me=True)
    chain.append({"role": "user", "content":
                  "[the other graders have now said]\n" + others +
                  "\n\nReconsider if they are right, or restate your position and say why."})
    return chain


async def one_turn(client, model, msgs, sem, max_tok, structured):
    async with sem:
        try:
            if structured:
                r = await client.beta.chat.completions.parse(
                    model=model, max_completion_tokens=max_tok, messages=msgs,
                    response_format=Turn,
                    extra_body={"provider": {"require_parameters": True}})
                p = r.choices[0].message.parsed
                if p:
                    return {"v": 1 if p.correct == "yes" else 0, "r": p.reasoning}
            raise RuntimeError("no structured parse")
        except Exception:
            try:
                r = await client.chat.completions.create(
                    model=model, max_completion_tokens=max_tok, messages=msgs)
                raw = r.choices[0].message.content or ""
                got = parse_prose(raw)
                if got:
                    return {"v": 1 if got["correct"] == "yes" else 0, "r": got["reasoning"]}
                # call succeeded, no verdict in the text -- keep it to diagnose
                return {"v": None, "r": f"<no verdict found in: {raw.strip()[:200]!r}>"}
            except Exception as e:
                return {"v": None, "r": f"<api failed: {str(e)[:140]}>"}


def tally(positions, names):
    yes = sum(1 for n in names if positions[n]["v"] == 1)
    return f"{yes}-{len(names)-yes}"


async def deliberate(client, qid, ctx, panel, start, args, sem):
    """Run rounds for one question. Returns the transcript."""
    names = list(panel)
    positions = dict(start)
    if not args.quiet:
        print(f"{qid[:8]}  opening   {tally(positions, names)}", flush=True)
    transcript = [{"round": 0, "positions": {k: dict(v) for k, v in positions.items()}}]
    chains = {n: [] for n in names}          # only used with --history
    last_said = {n: None for n in names}
    for rnd in range(1, args.rounds + 1):
        if len({p["v"] for p in positions.values()}) == 1:
            break
        start_of_round = {n: dict(v) for n, v in positions.items()}
        if args.history:
            msgs = {n: extend_chain(chains[n], ctx, positions, n, args.show_names,
                                    last_said[n]) for n in panel}
        else:
            msgs = {n: build(ctx["q"], ctx["resp"], ctx["key"], positions, n,
                             args.show_names) for n in panel}
        results = await asyncio.gather(*[
            one_turn(client, panel[n], msgs[n], sem, args.max_tokens, args.structured.get(n, True))
            for n in panel])
        failed = [n for n, res in zip(panel, results) if res["v"] is None]
        if failed:
            # Freezing a position on failure looks identical to "held their
            # ground", which is the one thing this must never fake.
            FAILS.extend((qid, n, dict(zip(panel, results))[n]["r"]) for n in failed)
        for n, res in zip(panel, results):
            if res["v"] is not None:
                positions[n] = res
                # A failed turn must not enter the chain as if it were said.
                last_said[n] = res
        if len(failed) == len(panel):
            break                          # nothing can change; stop burning rounds
        transcript.append({"round": rnd, "positions": {k: dict(v) for k, v in positions.items()}})
    out = {"id": qid, "solver": args.solver, "rounds": transcript,
           "final": {n: positions[n]["v"] for n in panel},
           "consensus": len({p["v"] for p in positions.values()}) == 1}
    out["movement"] = classify(out)
    return out


def classify(d):
    """Which way did the question resolve, relative to where it started?

    A panel that only ever ratifies its opening majority is an expensive way to
    recompute a majority vote. The case that justifies deliberation is the other
    one: a lone dissenter with a sound argument turning the other two. Recording
    the direction is what tells those apart.
    """
    r0 = d["rounds"][0]["positions"]
    names = list(r0)
    opening = {n: r0[n]["v"] for n in names}
    yes = sum(opening.values())
    majority = 1 if yes * 2 > len(names) else 0
    dissenters = [n for n in names if opening[n] != majority]
    movers = [n for n in names if opening[n] != d["final"][n]]
    verdicts = set(d["final"].values())
    out = {"opening": f"{yes}-{len(names) - yes}",
           "opening_majority": majority,
           "opening_dissenters": dissenters,
           "movers": movers}
    if len(verdicts) != 1:
        out["direction"] = "no_consensus"
        return out
    outcome = verdicts.pop()
    out["outcome"] = outcome
    out["direction"] = "majority_held" if outcome == majority else "minority_won"
    return out


def show(d, meta, names, runs):
    """Print one question's deliberation so the arguments can be read."""
    qid = d["id"]
    row = meta.get(qid, {})
    start = d["rounds"][0]["positions"]
    print("\n" + "=" * 78)
    print(f"{qid}   {row.get('category','?')} · {row.get('answer_type','?')}"
          f"   {'CONSENSUS' if d['consensus'] else 'STILL SPLIT'}"
          f" after {len(d['rounds'])-1} round(s)")
    mv = d.get("movement") or classify(d)
    arrow = {"majority_held": "opening majority held",
             "minority_won": "DISSENTER TURNED THE MAJORITY",
             "no_consensus": "no consensus"}[mv["direction"]]
    print(f"opened {mv['opening']} -> {arrow}"
          f"   (moved: {', '.join(mv['movers']) or 'nobody'})")
    print("=" * 78)
    print(f"[question]\n{row.get('question','')}\n")
    print(f"[answer being graded]\n{runs[names[0]][qid].get('response','')}\n")
    print(f"[answer key]\n{meta[qid]['answer']}\n")
    for rnd in d["rounds"]:
        label = "opening positions" if rnd["round"] == 0 else f"round {rnd['round']}"
        moved = ""
        if rnd["round"]:
            prev = d["rounds"][rnd["round"] - 1]["positions"]
            flips = [n for n in names if prev[n]["v"] != rnd["positions"][n]["v"]]
            moved = f"   ({', '.join(flips)} changed)" if flips else "   (nobody moved)"
        print(f"--- {label}{moved}")
        for n in names:
            p_ = rnd["positions"][n]
            print(f"  [{n}: {'correct' if p_['v'] else 'incorrect'}]")
            print(f"    {p_['r']}")
        print(flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("judged", nargs="+", help="two or more judged_*.json")
    ap.add_argument("--models", nargs="+",
                    help="OpenRouter id per judged file, in order (default: infer from filename)")
    ap.add_argument("--rounds", type=int, default=5)
    ap.add_argument("--history", action="store_true",
                    help="give each grader its own running conversation, with its "
                         "previous verdicts as assistant turns, instead of a fresh "
                         "snapshot of current positions each round")
    ap.add_argument("--limit", type=int, default=0,
                    help="only the first N split questions, for an A/B on a subset")
    ap.add_argument("--batch-size", type=int, default=0,
                    help="stop after this many questions (0 = all). Cached results "
                         "are skipped, so repeated runs walk through the backlog.")
    ap.add_argument("--quiet", action="store_true",
                    help="do not print transcripts, only the summary")
    ap.add_argument("--out", default="deliberation.json")
    ap.add_argument("--max-tokens", type=int, default=1500)
    ap.add_argument("--num-workers", type=int, default=8)
    ap.add_argument("--show-names", action="store_true",
                    help="reveal which judge holds which position")
    ap.add_argument("--dataset", default="cais/hle")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    runs = {name_of(p): json.loads(Path(p).read_text()) for p in args.judged}
    names = list(runs)
    panel = dict(zip(names, args.models)) if args.models else {
        n: f"openai/{n}" if n.startswith("gpt") else n for n in names}
    args.structured = {}
    args.solver = solver_of(args.judged[0])

    common = sorted(set.intersection(*[set(r) for r in runs.values()]))
    common = [q for q in common if len({runs[n][q].get("response") for n in names}) == 1]
    split = [q for q in common if 0 < sum(vote(runs[n][q]) for n in names) < len(names)]
    if args.limit:
        # Applied before the estimate so --dry-run reflects it, and taken from
        # the front so both arms of an A/B get the identical question set.
        split = split[:args.limit]
    print(f"panel {names}\n{len(common)} comparable, {len(split)} split")
    print(f"up to {len(split)} x {len(names)} x {args.rounds} = "
          f"{len(split)*len(names)*args.rounds} calls if nothing converges early")
    if args.models is None:
        print("\nno --models given; inferred OpenRouter ids:")
        for n, m in panel.items():
            print(f"  {n:<28} -> {m}")
        print("pass --models explicitly if any of those are wrong")
    if args.dry_run:
        return

    from datasets import load_dataset
    load_env(".env")
    meta = {r["id"]: r for r in load_dataset(args.dataset, split="test", token=hf_token())}
    key = os.environ.get("OPENROUTER_API_KEY")
    if not key:
        raise SystemExit("no OPENROUTER_API_KEY")

    done = {}
    if Path(args.out).exists():
        done = {d["id"]: d for d in json.loads(Path(args.out).read_text())}
        print(f"{len(done)} questions already in {args.out} — skipping those")
    todo = [q for q in split if q not in done]
    if args.batch_size:
        todo = todo[:args.batch_size]
    print(f"{len(split)} split, {len(done)} already settled, {len(todo)} this batch")

    client = AsyncOpenAI(base_url="https://openrouter.ai/api/v1", api_key=key,
                         timeout=120.0, max_retries=1)
    sem = asyncio.Semaphore(args.num_workers)

    async def run():
        tasks = []
        for qid in todo:
            ctx = {"q": meta[qid]["question"],
                   "resp": runs[names[0]][qid].get("response", ""),
                   "key": meta[qid]["answer"]}
            start = {n: {"v": vote(runs[n][qid]),
                         "r": runs[n][qid]["judge_response"].get("reasoning", "")}
                     for n in names}
            tasks.append(deliberate(client, qid, ctx, panel, start, args, sem))
        out = []
        for fut in asyncio.as_completed(tasks):
            d = await fut
            out.append(d)
            if not args.quiet:
                show(d, meta, names, runs)
            save(args.out, list(done.values()) + out)
            print(f"[{len(out)}/{len(tasks)} done]\n", flush=True)
        return out

    try:
        fresh = asyncio.run(run())
    except KeyboardInterrupt:
        # Everything finished so far is already on disk from the per-question
        # save; report what survived instead of dying with a traceback.
        done_now = json.loads(Path(args.out).read_text()) if Path(args.out).exists() else []
        print(f"\ninterrupted -- {len(done_now)} questions saved in {args.out}")
        raise SystemExit(0)
    if FAILS:
        from collections import Counter
        print(f"\n{len(FAILS)} turns failed, so those positions could not update "
              f"{dict(Counter(n for _, n, _ in FAILS))}")
        for qid, who, why in FAILS[:3]:
            print(f"  {qid[:8]} {who}: {why[:170]}")

    allr = list(done.values()) + fresh
    save(args.out, allr)

    reached = sum(1 for d in allr if d["consensus"])
    used = [len(d["rounds"]) - 1 for d in allr]
    print(f"\nconsensus on {reached}/{len(allr)}   rounds used: "
          f"median {sorted(used)[len(used)//2]}, max {max(used) if used else 0}")
    from collections import Counter
    dirs = Counter((d.get("movement") or classify(d))["direction"] for d in allr)
    print(f"  opening majority held        {dirs['majority_held']}")
    print(f"  dissenter turned the majority {dirs['minority_won']}")
    print(f"  never converged              {dirs['no_consensus']}")
    who = Counter(n for d in allr for n in (d.get("movement") or classify(d))["movers"])
    if who:
        print("  who changed their mind: "
              + ", ".join(f"{n} x{c}" for n, c in who.most_common()))
    print("\nhow often each judge moved from its original verdict")
    for n in names:
        moved = sum(1 for d in allr
                    if d["rounds"][0]["positions"][n]["v"] != d["final"][n])
        print(f"  {n:<28} {moved:>4}/{len(allr)}")
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()