"""Build a miniF2F VALIDATION problem file for model selection.

WHY. Round selection, the weight-averaging coefficient and the self-distillation stop rule are
all tuned by looking at scores. Tuning them on miniF2F-test makes the final test number partly
a tuning artefact (adaptive overfitting to 244 problems). miniF2F-valid (244 different problems,
same distribution) is the standard dev split; every development decision is taken on it and the
test split is touched only by the pre-registered final evaluation.

Source: cat-searcher/minif2f-lean4 `valid.jsonl`, pinned revision below. Statements are rewritten
into the SAME envelope as our test file (Goedel's header, `:= by sorry`), so prompts and the
verifier see identical structure. Decontamination against miniF2F-test, MOBench and the training
pool is by signature sha (statement_sha), renaming-invariant sha (alpha_sha), theorem name and
informal text (build_problems' 13-gram / 8-gram layers); any overlap is dropped and counted. Elaboration under the pinned Mathlib is checked with <= 2
Lean workers (not a heavy job), and failures are dropped and listed.
"""

from __future__ import annotations

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

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


def _load_sibling(name: str):
    """Load scripts/p2/<name>.py by path (the repo idiom; keeps the dependency firewall's import scan clean).
    A cached module is reused only if it IS this sibling file, and a failed import is not left cached."""
    path = Path(__file__).with_name(f"{name}.py").resolve()
    cached = sys.modules.get(name)
    if cached is not None and Path(getattr(cached, "__file__", "") or "").resolve() == path:
        return cached
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    try:
        spec.loader.exec_module(mod)
    except BaseException:
        sys.modules.pop(name, None)
        raise
    return mod


C = _load_sibling("p2_common")

SRC_REPO = "cat-searcher/minif2f-lean4"
SRC_REV = "70a1249ce240667f6bcdd1ccd62f847f0e065d57"
HEADER = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n"


def to_problem(row: dict) -> dict:
    stmt = row["formal_statement"].rstrip()
    if stmt.endswith(":= sorry"):
        stmt = stmt[: -len(":= sorry")] + ":= by sorry"
    elif not stmt.endswith(":= by sorry"):
        raise C.ContractError(f"{row['id']}: statement does not end in a sorry proof")
    informal = (row.get("informal_stmt") or "").strip()
    prefix = f"/-- {informal}-/\n" if informal else ""
    full = HEADER + prefix + stmt
    return {
        "problem_id": row["id"],
        "source": "minif2f",
        "split": "valid",
        "formal_statement": full,
        "informal_prefix": prefix or None,
        "statement_sha": C.statement_sha(full),
    }


def elaborate(pool, kept: list, timeout_s: float) -> tuple[list, list]:
    """Elaboration check with the grader ceiling-tested FIRST (known-good must pass, known-bad must
    fail), as build_problems.run_check does. Without it a broken toolchain or header would mark every
    statement as an elaboration failure and silently produce an empty dev split. `pool.map` returns
    exceptions in place; any exception stops the build instead of becoming a "failure" row."""
    bp = _load_sibling("build_problems")
    probe = pool.map(lambda w, s: bp.check_elaborates(w, s, timeout_s), [bp.SELF_TEST_GOOD, bp.SELF_TEST_BAD])
    for what, res, want in (("known-good", probe[0], True), ("known-bad", probe[1], False)):
        if isinstance(res, BaseException):
            raise res
        if res[0] is not want:
            raise RuntimeError(f"elaboration instrument failed its ceiling test: {what} statement gave {res}")
    res = pool.map(lambda w, rec: bp.check_elaborates(w, rec["formal_statement"], timeout_s), kept)
    good, bad = [], []
    for p, r in zip(kept, res):
        if isinstance(r, BaseException):
            raise r
        ok, why = r
        if ok:
            good.append(p)
        else:
            bad.append({"id": p["problem_id"], "reason": why})
    return good, bad


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--src", required=True, help="the pinned valid.jsonl")
    ap.add_argument("--problems-dir", required=True, help="dir holding minif2f.jsonl, mobench.jsonl, train_pool.jsonl")
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--elab-timeout", type=float, default=60.0)
    ap.add_argument("--no-elab", action="store_true", help="skip the Lean elaboration check (tests)")
    a = ap.parse_args(argv)
    if not 1 <= a.workers <= 2:
        raise SystemExit("--workers must be 1 or 2 (this is not a heavy job)")

    pdir = Path(a.problems_dir)
    shas, alphas, names, guarded = set(), set(), set(), []
    for f in ("minif2f.jsonl", "mobench.jsonl", "train_pool.jsonl"):
        for r in C.iter_jsonl(pdir / f):
            shas.add(r["statement_sha"])
            alphas.add(C.alpha_sha(r["formal_statement"]))
            names.add(r["problem_id"])
            guarded.append(r)
    # informal-text layers (13-gram / three 8-grams): "same informal problem, different formalisation" is the realistic leak, and
    # build_problems applies it to the pool against test+MOBench only, never against this dev split
    informal = _load_sibling("build_problems").Decontaminator(guarded)

    kept, dropped = [], []
    for row in C.iter_jsonl(a.src):
        try:
            p = to_problem(row)
        except C.ContractError as e:
            dropped.append({"id": row.get("id"), "reason": f"unsupported: {e}"})
            continue
        if p["statement_sha"] in shas or C.alpha_sha(p["formal_statement"]) in alphas or p["problem_id"] in names:
            dropped.append({"id": p["problem_id"], "reason": "overlaps test/MOBench/pool"})
            continue
        why = informal.check_informal(p["informal_prefix"])
        if why is not None:
            dropped.append({"id": p["problem_id"], "reason": f"informal text overlaps test/MOBench/pool ({why})"})
            continue
        kept.append(p)

    elab_bad = []
    if not a.no_elab and kept:
        from align.utils import guard
        from align.verify.lean import ReplPool, find_toolchain
        from align.verify.lean.repl_pool import ReplLimits

        guard.preflight_host_check(strict=True)
        # scratch on /apps (repo .tmp), never $TMPDIR: /apps/CLAUDE.md rule 5 (no writes to / or /tmp)
        pool = ReplPool.from_toolchain(find_toolchain(), n_workers=a.workers, limits=ReplLimits(), preflight=True,
                                       scratch_dir=str(REPO / ".tmp" / "p2_valid_pool"))
        try:
            kept, elab_bad = elaborate(pool, kept, a.elab_timeout)
        finally:
            pool.close()

    # split first, then its manifest, both atomic. The manifest records the split's sha256, so a kill
    # between the two writes (new split, old manifest) is DETECTABLE by any reader, never silent.
    C.write_jsonl_atomic(a.out, kept)
    manifest = {"source": SRC_REPO, "revision": SRC_REV, "kept": len(kept), "dropped": dropped,
                "elaboration_failures": elab_bad, "header": HEADER,
                "split_sha256": hashlib.sha256(Path(a.out).read_bytes()).hexdigest()}
    mpath = Path(a.out + ".manifest.json")
    mtmp = mpath.with_name(mpath.name + f".tmp{os.getpid()}")
    mtmp.write_text(json.dumps(manifest, indent=1))
    os.replace(mtmp, mpath)
    print(json.dumps({"kept": len(kept), "dropped": len(dropped), "elab_failed": len(elab_bad)}))
    return 0


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