#!/usr/bin/env python3
"""Verifier for the exact-dataset-construction-replication task.

The experiment behind this task: given an exhaustive, code-derived natural-
language specification (and the raw source files, but no code and no results),
can an agent reconstruct the benchmark *exactly* -- down to each record's `id`?
So the grade is a plain exact match against the frozen reference release in
tests/key/, not a tiered/pool-based rubric.

For each of the 8 subtasks x 3 splits we canonicalize every record
(json.dumps(sort_keys=True), so object key order and file whitespace are free)
and compare as multisets:

  exact_match     a gold record is reproduced iff a byte-identical record --
                  INCLUDING its `id` -- appears in the agent's file for that
                  split. This is the headline "certificate" number: hitting it
                  requires reproducing every pseudo-random draw (the sampling,
                  the split assignment, the id numbering), not just the content.
  content_match   the same comparison with each record's `id` removed first.
                  This isolates whether the RIGHT instances (right fields, right
                  split) were built, independent of the RNG-driven numbering and
                  ordering -- so a run that reconstructs the data but seeds
                  differently still scores here.

reward = exact_match fraction over all reference records (graded in [0, 1]).
Writes /logs/verifier/reward.json (+ reward.txt) and a per-subtask report.json.
Pure stdlib.
"""
from __future__ import annotations

import json
import os
import sys
from collections import Counter

HERE = os.path.dirname(os.path.abspath(__file__))
KEY = os.path.join(HERE, "key")
WORKSPACE = os.getenv("WORKSPACE", "/workspace")
REWARD_DIR = os.getenv("REWARD_DIR", "/logs/verifier")

SUBTASKS = ["argument", "wortlaut", "systematik", "geschichte", "zweck",
            "konkretes_gesetz", "nicht_abstrakt", "nicht_selbst_aufgestellt"]
SPLITS = ["train", "validation", "test"]


def canon(rec, drop_id=False):
    if drop_id and isinstance(rec, dict):
        rec = {k: v for k, v in rec.items() if k != "id"}
    return json.dumps(rec, sort_keys=True, ensure_ascii=False)


def load_list(path):
    """Return a list of records, or None if the file is missing/not a JSON list."""
    if not os.path.exists(path):
        return None
    try:
        with open(path, encoding="utf-8") as fh:
            data = json.load(fh)
    except Exception:  # noqa: BLE001
        return None
    return data if isinstance(data, list) else None


def multiset_overlap(gold, agent, drop_id):
    """Count of gold records reproduced in agent (multiset intersection size)."""
    g = Counter(canon(r, drop_id) for r in gold)
    a = Counter(canon(r, drop_id) for r in agent)
    return sum((g & a).values())


def grade_split(subtask, split):
    gold = load_list(os.path.join(KEY, subtask, f"{split}.json"))
    if gold is None:
        raise SystemExit(f"key missing/invalid: {subtask}/{split}.json")
    agent = load_list(os.path.join(WORKSPACE, subtask, f"{split}.json"))
    present = agent is not None
    agent = agent or []
    return {
        "subtask": subtask,
        "split": split,
        "n_gold": len(gold),
        "n_agent": len(agent),
        "file_present": present,
        "n_exact": multiset_overlap(gold, agent, drop_id=False),
        "n_content": multiset_overlap(gold, agent, drop_id=True),
    }


def main() -> int:
    os.makedirs(REWARD_DIR, exist_ok=True)

    cells = [grade_split(s, sp) for s in SUBTASKS for sp in SPLITS]
    n_total = sum(c["n_gold"] for c in cells)
    n_exact = sum(c["n_exact"] for c in cells)
    n_content = sum(c["n_content"] for c in cells)

    exact_frac = round(n_exact / n_total, 6) if n_total else 0.0
    content_frac = round(n_content / n_total, 6) if n_total else 0.0

    # per-subtask rollup (for the human report; not part of the numeric reward)
    per_subtask = {}
    for s in SUBTASKS:
        sc = [c for c in cells if c["subtask"] == s]
        g = sum(c["n_gold"] for c in sc)
        per_subtask[s] = {
            "n_gold": g,
            "exact_frac": round(sum(c["n_exact"] for c in sc) / g, 6) if g else 0.0,
            "content_frac": round(sum(c["n_content"] for c in sc) / g, 6) if g else 0.0,
            "files_present": sum(1 for c in sc if c["file_present"]),
        }

    # reward.json is consumed as VerifierResult.rewards (flat dict of numbers).
    reward = {
        "reward": exact_frac,
        "exact_match_frac": exact_frac,
        "content_match_frac": content_frac,
        "n_records_total": n_total,
        "n_records_exact": n_exact,
        "n_records_content": n_content,
        "n_subtasks_fully_exact": sum(1 for v in per_subtask.values() if v["exact_frac"] == 1.0),
        "n_files_present": sum(1 for c in cells if c["file_present"]),
    }
    report = {"per_subtask": per_subtask, "per_split": cells}

    with open(os.path.join(REWARD_DIR, "reward.json"), "w", encoding="utf-8") as fh:
        json.dump(reward, fh, indent=2)
    with open(os.path.join(REWARD_DIR, "reward.txt"), "w", encoding="utf-8") as fh:
        fh.write(str(exact_frac))
    with open(os.path.join(REWARD_DIR, "report.json"), "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=2)

    # human-readable summary -> run log
    print(f"reward(exact)={exact_frac}  content={content_frac}  "
          f"exact={n_exact}/{n_total}  files={reward['n_files_present']}/24")
    for s in SUBTASKS:
        v = per_subtask[s]
        print(f"  {s:26} exact={v['exact_frac']:.4f}  content={v['content_frac']:.4f}  "
              f"files={v['files_present']}/3")
    return 0 if n_exact == n_total else 1


if __name__ == "__main__":
    sys.exit(main())