#!/usr/bin/env python3
"""
validate.py — G6 validation/judge: a deterministic consistency pass over an extractor result.

NOT a gate. It never approves or rejects; it attaches confidence + flags so a HUMAN reviewer's
attention is routed to the parts most likely wrong. A clean paper -> all-green, skim; a flagged
paper -> the exact items to check, worst-first. Every value stays a *suggestion* a human confirms.

Checks (all deterministic, no GPU, ~free — run on every extraction):
  C1 marks-sum vs official max      — over-read (sum>max) = error; under (sum<max) = warn
  C2 part marks plausibility        — marks None / 0 / implausibly high
  C3 top-level question sequence    — gaps in 1..N (skipped when numbering was OCR-inferred '~')
  C4 sub-part contiguity            — within a question: a,b,c / .1,.2,.3 with no hole
  C5 coverage                       — missed parts vs ground truth (when the result carries it)

Usage:
  python validate.py results/genreport/edexcel1f/ocr_struct_filled.json
  python validate.py <structured.json> --out report.json
"""
import json, re, sys, argparse
from collections import defaultdict

IMPLAUSIBLE_PART_MARKS = 15        # a single sub-part above this is worth a human glance


def _qnum(q):
    """Numeric value of a top-level question id ('01'->1, '4'->4); None if inferred ('~3') / odd."""
    if q.startswith("~"):
        return None
    m = re.match(r"^0*(\d+)$", q)
    return int(m.group(1)) if m else None


def _subkey(label, q):
    """The part's own suffix within its question: '01.2'->'2', '4a'->'a', '1bi'->'bi'."""
    s = label[len(q):] if label.startswith(q) else label
    return s.lstrip(".").lstrip("~")


def validate(result):
    board = result.get("board")
    code = result.get("paper_code")
    flags, checks = [], []
    parts = [(p["label"], q["question"], p) for q in result.get("questions", []) for p in q["parts"]]
    conf = {}                                      # label -> high/medium/low
    low = set()                                    # labels a check has implicated

    def add(cid, severity, status, detail):
        checks.append({"id": cid, "severity": severity, "status": status, "detail": detail})
        if status != "ok":
            flags.append(f"[{severity}] {cid}: {detail}")

    # ---- C1: marks sum vs official maximum -------------------------------------------------
    mc = result.get("stats", {}).get("marks_check")
    exp = (mc or {}).get("expected_max") or result.get("front_matter", {}).get("max_marks")
    msum = (mc or {}).get("sum")
    if msum is None:
        msum = sum(p["marks"] for *_, p in parts if p.get("marks") is not None)
    if exp:
        if msum > exp:
            add("C1_marks_sum", "error", "over",
                f"marks sum {msum} EXCEEDS official max {exp} (+{msum-exp}) — an over-read; check the paper")
        elif msum < exp:
            add("C1_marks_sum", "warn", "under",
                f"marks sum {msum} below official max {exp} (-{exp-msum}) — missing parts or unread marks")
        else:
            add("C1_marks_sum", "info", "ok", f"marks sum {msum} == official max {exp}")
    else:
        add("C1_marks_sum", "info", "unknown", "no official max available to check the sum against")

    # ---- C2: per-part marks plausibility ---------------------------------------------------
    none_ct = zero_ct = 0
    for lab, q, p in parts:
        mk = p.get("marks")
        if mk is None:
            none_ct += 1; low.add(lab)
        elif mk == 0:
            zero_ct += 1; low.add(lab)
        elif mk > IMPLAUSIBLE_PART_MARKS:
            low.add(lab)
            add("C2_part_marks", "warn", "implausible",
                f"part {lab} has {mk} marks (> {IMPLAUSIBLE_PART_MARKS}) — verify it isn't a mis-read")
    if none_ct or zero_ct:
        add("C2_part_marks", "warn", "missing",
            f"{none_ct} part(s) with no mark, {zero_ct} with 0 marks — unread/garbled mark tokens")
    elif not any(c["id"] == "C2_part_marks" for c in checks):
        add("C2_part_marks", "info", "ok", "every part carries a plausible mark")

    # ---- C3: top-level question sequence + EXPECTED-question interpolation ------------------
    # If Q1, Q2 ... Q14 are recovered but 3-13 are not, the paper certainly HAS 3-13 — they were
    # just missed (e.g. a Docling page-collapse). We emit the full expected sequence with a per-Q
    # `recovered` flag so a live question-tree view can render the gaps as explicit "needs a second
    # pass" slots, and a targeted re-OCR knows exactly which questions to chase.
    qids = [q for q in dict.fromkeys(q for _, q, _ in parts)]
    nums = sorted({n for n in (_qnum(q) for q in qids) if n is not None})
    zero_pad = any(len(q) == 2 and q.startswith("0") for q in qids)   # AQA 'NN' vs Edexcel/OCR 'N'
    question_sequence = []
    if any(q.startswith("~") for q in qids):
        add("C3_question_seq", "info", "inferred",
            "question numbers were OCR-inferred ('~N') — sequence not checkable; treat labels as approximate")
    elif nums:
        # isolated high outliers (a content number mis-read as 'Q67' after Q1-10) are likely
        # spurious top-levels, not 50 missing questions — strip them off the top so the sequence
        # reflects the real paper, and flag them for review instead of flooding the tree with slots.
        core, suspect = nums[:], []
        while len(core) >= 2 and core[-1] - core[-2] > 4:
            suspect.insert(0, core.pop())
        hi = core[-1] if core else nums[-1]
        gaps = [n for n in range(nums[0], hi + 1) if n not in core]
        question_sequence = [{"n": n, "label": (f"{n:02d}" if zero_pad else str(n)),
                              "recovered": n in core} for n in range(nums[0], hi + 1)]
        if suspect:
            add("C3_question_seq", "warn", "spurious",
                f"isolated high question number(s) {suspect} after a {nums[0]}-{hi} run — likely a "
                f"content number mis-read as a top-level question; review/remove")
        if gaps:
            add("C3_question_seq", "warn", "gap",
                f"top-level questions {gaps} missing between {nums[0]}-{hi} — expected but "
                f"unrecovered; surface as second-pass slots in the question tree")
        elif not suspect:
            add("C3_question_seq", "info", "ok", f"questions {nums[0]}-{hi} contiguous")

    # ---- C4: sub-part contiguity within each question --------------------------------------
    def order(keys):
        """Map a question's child keys to an ordered scheme + report holes. Handles .N and a/b/c."""
        dig = sorted(int(k[0]) for k in keys if k[:1].isdigit())
        let = sorted(k[0] for k in keys if k[:1].isalpha())
        holes = []
        if dig:
            holes += [str(n) for n in range(dig[0], dig[-1] + 1) if n not in dig]
        if let:
            lo, hi = ord(let[0]), ord(let[-1])
            holes += [chr(c) for c in range(lo, hi + 1) if chr(c) not in let]
        return holes
    byq = defaultdict(list)
    for lab, q, p in parts:
        sk = _subkey(lab, q)
        if sk:
            byq[q].append(sk)
    seq_holes = {}
    for q, keys in byq.items():
        firsts = {k[0] for k in keys}            # immediate children only (a / 1 / etc.)
        h = order(firsts)
        if h:
            seq_holes[q] = h
    if seq_holes:
        add("C4_subpart_seq", "warn", "gap",
            "sub-part gaps: " + ", ".join(f"Q{q} missing {hs}" for q, hs in sorted(seq_holes.items())))
    else:
        add("C4_subpart_seq", "info", "ok", "sub-parts contiguous within every question")

    # ---- C5: coverage vs ground truth (when present) ---------------------------------------
    cov = result.get("coverage", {})
    if cov.get("coverage_pct") is not None:
        missed = cov.get("missed", [])
        if missed:
            add("C5_coverage", "warn", "missed",
                f"{cov['coverage_pct']}% vs GT ({cov['recovered']}/{cov['total']}); missed {missed[:10]}")
            low.update(missed)
        else:
            add("C5_coverage", "info", "ok", f"100% coverage vs GT ({cov['recovered']}/{cov['total']})")

    # ---- per-part confidence + paper summary -----------------------------------------------
    sum_mismatch = any(c["id"] == "C1_marks_sum" and c["status"] in ("over", "under") for c in checks)
    for lab, q, p in parts:
        if lab in low:
            conf[lab] = "low"
        elif sum_mismatch:
            conf[lab] = "medium"               # paper-level doubt taints every part a little
        else:
            conf[lab] = "high"
    severities = [c["severity"] for c in checks if c["status"] not in ("ok", "info", "unknown")]
    worst = "error" if "error" in severities else "warn" if "warn" in severities else "clean"

    return {
        "paper_code": code, "board": board,
        "summary": {
            "worst_severity": worst,
            "needs_priority_review": worst != "clean",
            "n_flags": len(flags),
            "marks_sum": msum, "official_max": exp,
            "parts_total": len(parts),
            "parts_low_conf": sum(1 for v in conf.values() if v == "low"),
            "questions_expected": len(question_sequence) or None,
            "questions_recovered": sum(1 for q in question_sequence if q["recovered"]) or None,
        },
        "flags": flags,
        "checks": checks,
        "part_confidence": conf,
        "question_sequence": question_sequence,   # full expected skeleton (recovered + missing slots)
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("structured")
    ap.add_argument("--out")
    a = ap.parse_args()
    rep = validate(json.load(open(a.structured)))
    s = rep["summary"]
    print(f"paper       : {rep['paper_code']}  ({rep['board']})")
    print(f"verdict     : {s['worst_severity'].upper()}  "
          f"{'-> PRIORITY REVIEW' if s['needs_priority_review'] else '-> all checks clean (still human-reviewable)'}")
    print(f"marks       : {s['marks_sum']}/{s['official_max']}  | parts {s['parts_total']} "
          f"({s['parts_low_conf']} low-confidence)")
    if s.get("questions_expected"):
        miss = [q["label"] for q in rep["question_sequence"] if not q["recovered"]]
        print(f"questions   : {s['questions_recovered']}/{s['questions_expected']} recovered"
              + (f"  | second-pass slots: {miss}" if miss else "  (complete sequence)"))
    if rep["flags"]:
        print("flags:")
        for f in rep["flags"]:
            print(f"  - {f}")
    else:
        print("flags       : none")
    if a.out:
        json.dump(rep, open(a.out, "w"), indent=2)
        print(f"-> wrote {a.out}")


if __name__ == "__main__":
    main()