#!/usr/bin/env python
"""P2 integrator: assert the seams between the independently built pieces, on real files.

    .venv/bin/python scripts/p2/check_seams.py keys  --generations G.jsonl --verifications V.jsonl
    .venv/bin/python scripts/p2/check_seams.py sft   --generations G0 G1 --verifications V0 V1 --shards D/shards

Each check exists because the pieces were written by different hands against a prose contract,
and each failure below is silent when only one side is looked at:

  keys  generation -> verification. A verification file whose key set differs from the
        generation file's (a sample never verified, a verdict for a sample that is not on disk, a
        verdict at another cap or round) scores a different experiment than the one generated.
        Also checks `code_sha256` against the generation on disk (a stale verdict).
  sft   verification -> trainer input. The tokens the trainer takes loss on must be EXACTLY the
        `completion_ids` vLLM sampled for a sample the verifier ACCEPTED, and the context must be
        exactly its `prompt_ids`. This re-derives the labels through the same `ShardDataset` +
        `collate_sft` the trainer uses and does not trust `build_sft_data.py`'s own read-back
        (an independent implementation of the same property is the only thing that catches a
        shared misunderstanding).

Exit status 0 only if every check passes; a report JSON is printed either way.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))


def _common():
    if "p2_common" in sys.modules:
        return sys.modules["p2_common"]
    spec = importlib.util.spec_from_file_location("p2_common", Path(__file__).with_name("p2_common.py"))
    mod = importlib.util.module_from_spec(spec)
    sys.modules["p2_common"] = mod
    spec.loader.exec_module(mod)
    return mod


class SeamError(AssertionError):
    pass


def _load(paths, key_fields):
    out: dict[tuple, dict] = {}
    for p in paths:
        for rec in _common().iter_jsonl(p):
            k = tuple(rec[f] for f in key_fields)
            if k in out:
                raise SeamError(f"duplicate key {k} in {p}")
            out[k] = rec
    return out


def check_keys(gen_paths, ver_paths) -> dict:
    """Generation key set == verification key set, and every verdict is for the code on disk."""
    C = _common()
    gens = _load(gen_paths, C.KEY_FIELDS)
    vers = _load(ver_paths, C.KEY_FIELDS)
    missing = sorted(set(gens) - set(vers))
    extra = sorted(set(vers) - set(gens))
    if missing or extra:
        raise SeamError(f"key sets differ: {len(missing)} generated but unverified (e.g. {missing[:2]}), "
                        f"{len(extra)} verified but not generated (e.g. {extra[:2]})")
    stale = [k for k, g in gens.items() if vers[k]["code_sha256"] != C.code_sha(g["lean_code"])]
    if stale:
        raise SeamError(f"{len(stale)} verdicts were computed for different code than the generation holds, e.g. {stale[:2]}")
    # a length-capped sample is never a proof, whatever the verdict file says
    bad = [k for k, g in gens.items() if g["finish_reason"] == "length" and vers[k]["accepted"]]
    if bad:
        raise SeamError(f"{len(bad)} length-truncated samples are marked accepted, e.g. {bad[:2]}")
    rounds: dict = {}
    for k, v in vers.items():
        r = rounds.setdefault(k[2], [0, 0])
        r[0] += 1
        r[1] += int(bool(v["accepted"]))
    return {"keys": len(gens), "by_round": {str(r): {"samples": a, "accepted": b} for r, (a, b) in sorted(rounds.items())}}


def check_sft(gen_paths, ver_paths, shards_dir) -> dict:
    """Every stored row's loss tokens == completion_ids of an ACCEPTED generation, context == prompt_ids."""
    import numpy as np

    from align.data.loader import collate_sft
    from align.data.store import ShardDataset

    C = _common()
    gens = _load(gen_paths, C.KEY_FIELDS)
    vers = _load(ver_paths, C.KEY_FIELDS)
    rows = 0
    kinds: dict = {}
    seen: set = set()
    for split in ("train", "held"):
        d = Path(shards_dir) / "sft" / split
        if not d.is_dir():
            continue
        ds = ShardDataset(d)
        for i in range(len(ds)):
            row = ds[i]
            ex = row.extras
            key = (row.problem_id, int(ex["sample_idx"]), int(ex["round"]), int(ex["cap_tokens"]))
            if key not in gens:
                raise SeamError(f"{split}[{i}] {key}: no such generation")
            if not (key in vers and vers[key]["accepted"]):
                raise SeamError(f"{split}[{i}] {key}: trained on a sample the verifier did not accept")
            if key in seen:
                raise SeamError(f"{key}: stored twice")
            seen.add(key)
            g = gens[key]
            p = int(ex["prompt_len"])
            ids = np.asarray(row.input_ids)
            if p != len(g["prompt_ids"]) or not np.array_equal(ids[:p], np.asarray(g["prompt_ids"])):
                raise SeamError(f"{key}: stored context is not the sampled prompt_ids")
            lab = collate_sft([row], pad_token_id=0).labels[0].numpy()
            trained = lab[lab != -100]
            if not np.array_equal(trained, np.asarray(g["completion_ids"], dtype=trained.dtype)):
                raise SeamError(f"{key}: tokens with loss ({len(trained)}) != completion_ids ({len(g['completion_ids'])})")
            if not np.all(lab[:p] == -100):
                raise SeamError(f"{key}: loss leaks into the prompt")
            kinds[str(ex["kind"])] = kinds.get(str(ex["kind"]), 0) + 1
            rows += 1
    if rows == 0:
        raise SeamError("no SFT rows were found; a vacuous pass")
    accepted = sum(1 for v in vers.values() if v["accepted"])
    return {"rows": rows, "by_kind": kinds, "accepted_samples": accepted, "accepted_not_used": accepted - rows}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name in ("keys", "sft"):
        s = sub.add_parser(name)
        s.add_argument("--generations", nargs="+", required=True)
        s.add_argument("--verifications", nargs="+", required=True)
        if name == "sft":
            s.add_argument("--shards", required=True)
    a = ap.parse_args(argv)
    try:
        rep = check_keys(a.generations, a.verifications) if a.cmd == "keys" else check_sft(a.generations, a.verifications, a.shards)
        rep["passed"] = True
    except SeamError as e:
        rep = {"passed": False, "error": str(e)}
    print(json.dumps(rep, indent=2))
    return 0 if rep["passed"] else 1


if __name__ == "__main__":
    raise SystemExit(main())