#!/usr/bin/env python
"""Verify P2 Generation JSONL with the strict Lean verifier -> resumable Verification JSONL.

    .venv/bin/python scripts/p2/verify_generations.py \\
        --generations gens_r0.jsonl --problems minif2f.jsonl --out ver_r0.jsonl --workers 8

Output: one line per generation, keyed by (problem_id, sample_idx, round, cap_tokens):

    {problem_id, sample_idx, round, cap_tokens, accepted, reasons, verify_s, budget_s,
     verifier: {pins, config, workers}, code_sha256, stage, diagnostics[, repl_errors, repl_code]}

Properties (each one is a defect the first version of such a script has):

  * RESUMABLE. The file is append-only and fsynced per record; a rerun skips every key that is
    already there. A line cut in half by SIGKILL is dropped on restart (`repair_torn_tail`); a
    malformed line elsewhere raises (appends only ever damage the tail).
  * NO STALE VERDICTS. Every record carries `code_sha256` of the code that was verified. A
    resumed run whose generation file was regenerated (same keys, different samples) finds a
    different hash and RAISES instead of attaching the old verdict to new text.
  * NO LEAN FOR NOTHING. `finish_reason == "length"` (a truncated proof is not a proof) and
    code-less samples are recorded as rejected with reasons ["no_code"] without calling Lean.
  * A DEAD VERIFIER IS NOT A ZERO. Results whose reason is an infrastructure death (crash,
    protocol, memory, output_too_large) are held back until their chunk is complete; if more
    than half of a chunk of >= 8 died that way the run aborts WITHOUT writing them, so a
    resume retries them. Otherwise they are written (a proof that kills its REPL scores 0).
  * HOST RULES. >2 workers => exclusive heavy-job flock (waited for, not skipped) and the
    host preflight; workers <= 8; scratch under /apps.
  * `--feedback` additionally stores, for samples rejected with `compile_error`, `repl_errors`
    (every Lean error: pos / endPos / data) and `repl_code` (the exact text the positions index:
    Goedel's `compile_view(lean_code)`), which `extras/vllm_rollout/lean_gen.py` needs to build
    correction prompts and which `StrictResult.diagnostics` (3 errors, 240 chars, no positions)
    cannot supply. Rejections that are not compile errors (banned construct, statement
    mismatch, timeout...) carry no Lean error list; the generator skips and counts them.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import logging
import os
import sys
import threading
import time
from dataclasses import asdict
from pathlib import Path
from typing import Callable, Optional

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


def _load_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


C = sys.modules.get("p2_common") or _load_common()
log = logging.getLogger("p2.verify")

INFRA_REASONS = ("crash", "protocol", "memory", "output_too_large")
MAX_WORKERS = 8


class StaleVerification(C.ContractError):
    """An existing verification record belongs to different code than the generation now holds."""


class VerifierDied(RuntimeError):
    """More than half of a chunk died for infrastructure reasons; nothing from it was written."""


class _NullWriter:
    def append(self, rec) -> None:
        pass

    def close(self) -> None:
        pass


class Result:
    """What the verify callback returns: the minimum the record needs."""

    __slots__ = ("accepted", "reasons", "elapsed_s", "stage", "diagnostics", "feedback")

    def __init__(self, accepted, reasons, elapsed_s, stage="", diagnostics=None, feedback=None):
        self.accepted = bool(accepted)
        self.reasons = list(reasons)
        self.elapsed_s = float(elapsed_s)
        self.stage = stage
        self.diagnostics = diagnostics or {}
        self.feedback = feedback

    @property
    def infra(self) -> bool:
        return any(r in INFRA_REASONS for r in self.reasons)


# =========================================================================== core (verifier-agnostic)


def load_existing(path: Path) -> dict[tuple, dict]:
    out: dict[tuple, dict] = {}
    if not path.exists():
        return out
    for rec in C.iter_jsonl(path):
        for k in (*C.KEY_FIELDS, "accepted", "reasons", "verify_s", "budget_s", "verifier", "code_sha256"):
            if k not in rec:
                raise C.ContractError(f"{path}: verification record lacks {k!r}: {str(rec)[:200]}")
        key = C.gen_key(rec)
        if key in out:
            raise C.ContractError(f"{path}: duplicate verification key {key} (the file is append-only)")
        out[key] = rec
    return out


def iter_generations(paths):
    seen: set[tuple] = set()
    for p in paths:
        for rec in C.iter_jsonl(p):
            C.validate_generation(rec)
            key = C.gen_key(rec)
            if key in seen:
                raise C.ContractError(f"{p}: duplicate generation key {key} across the input files")
            seen.add(key)
            yield rec


def make_record(gen: dict, res: Result, budget_s: float, verifier_meta: dict) -> dict:
    rec = {
        "problem_id": gen["problem_id"],
        "sample_idx": gen["sample_idx"],
        "round": gen["round"],
        "cap_tokens": gen["cap_tokens"],
        "accepted": res.accepted,
        "reasons": res.reasons,
        "verify_s": round(res.elapsed_s, 4),
        "budget_s": budget_s,
        "verifier": verifier_meta,
        "code_sha256": C.code_sha(gen["lean_code"]),
        "stage": res.stage,
        "diagnostics": res.diagnostics,
    }
    if res.feedback is not None:
        rec.update(res.feedback)  # repl_errors / repl_code (or repl_infra): top-level, as lean_gen.py reads them
    return rec


def no_code_cause(gen: dict) -> Optional[str]:
    if gen["finish_reason"] == "length":
        return "finish_reason=length"
    if gen["lean_code"] is None or not gen["lean_code"].strip():
        return "no_extracted_code"
    return None


def verify_all(
    gen_paths,
    problems: dict[str, dict],
    out_path: Path,
    batch_fn: Callable[[list[tuple[str, str]], Callable[[int, Result], None]], None],
    budget_s: float,
    verifier_meta: dict,
    chunk: int = 128,
    limit: Optional[int] = None,
    dry_run: bool = False,
) -> dict:
    """Verify every generation not yet in `out_path`. `batch_fn(items, on_result)` must call
    `on_result(index, Result)` once per item (from any thread). Returns run statistics."""
    out_path = Path(out_path)
    torn = C.repair_torn_tail(out_path)
    existing = load_existing(out_path)
    stats = {"generations": 0, "already_done": 0, "no_code": 0, "verified": 0, "accepted": 0, "torn_bytes_removed": torn}
    pending: list[dict] = []
    writer = _NullWriter() if dry_run else C.AppendJsonl(out_path)
    try:
        for gen in iter_generations(gen_paths):
            stats["generations"] += 1
            if gen["problem_id"] not in problems:
                raise C.ContractError(f"generation {C.gen_key(gen)}: problem_id not in the problems file(s)")
            key = C.gen_key(gen)
            if key in existing:
                if existing[key]["code_sha256"] != C.code_sha(gen["lean_code"]):
                    raise StaleVerification(
                        f"{key}: the stored verdict was computed for different code (sha "
                        f"{existing[key]['code_sha256']} vs {C.code_sha(gen['lean_code'])}); the generation file was "
                        "regenerated under the same keys. Use a fresh --out, never reuse verdicts."
                    )
                stats["already_done"] += 1
                continue
            cause = no_code_cause(gen)
            if cause is not None:
                rec = make_record(gen, Result(False, ["no_code"], 0.0, "skipped", {"cause": cause}), budget_s, verifier_meta)
                writer.append(rec)
                existing[key] = rec
                stats["no_code"] += 1
                continue
            pending.append(gen)
            if limit is not None and len(pending) >= limit:
                break
        log.info("%d generations: %d already verified, %d no_code written, %d to verify", stats["generations"],
                 stats["already_done"], stats["no_code"], len(pending))
        stats["pending"] = len(pending)
        t0 = time.monotonic()
        for lo in range(0, 0 if dry_run else len(pending), chunk):
            part = pending[lo : lo + chunk]
            items = [(problems[g["problem_id"]]["formal_statement"], g["lean_code"]) for g in part]
            held: list[tuple[int, Result]] = []
            lock = threading.Lock()
            results: dict[int, Result] = {}

            def on_result(i: int, res: Result, part=part) -> None:
                with lock:
                    results[i] = res
                    if res.infra:
                        held.append((i, res))
                        return
                writer.append(make_record(part[i], res, budget_s, verifier_meta))

            batch_fn(items, on_result)
            missing = [i for i in range(len(part)) if i not in results]
            if missing:
                raise RuntimeError(f"batch_fn returned without a result for items {missing[:5]}")
            if len(part) >= 8 and len(held) / len(part) > 0.5:
                raise VerifierDied(
                    f"{len(held)}/{len(part)} verifications of chunk {lo // chunk} died for infrastructure reasons "
                    f"({sorted({r for _, x in held for r in x.reasons})}); not writing them, so a resume retries. "
                    "Either the pool is dead or the policy is emitting REPL-killing proofs: look before rerunning."
                )
            for i, res in held:
                writer.append(make_record(part[i], res, budget_s, verifier_meta))
            for res in results.values():
                stats["verified"] += 1
                stats["accepted"] += int(res.accepted)
            log.info("verified %d/%d (%.0fs elapsed, %d accepted so far)", min(lo + chunk, len(pending)), len(pending),
                     time.monotonic() - t0, stats["accepted"])
    finally:
        writer.close()
    return stats


# =========================================================================== real Lean backend


def collect_feedback(worker, statement: str, code: str, timeout_s: float = 120.0) -> Optional[dict]:
    """Full Lean error list for `code`, in the form the correction-round prompt needs.

    Reproduces Goedel's compile step: `compile_view(code)` (import / set_option / open lines
    dropped) is run in an environment that already holds the statement's header, and the
    REPL's error dicts are kept as `repl_errors` with the exact text they index as `repl_code`
    (`extras/vllm_rollout/lean_gen.py` joins on these two keys). `StrictResult.diagnostics`
    cannot serve: it keeps 3 errors, 240 characters each, without positions.

    Returns None when the statement/code cannot be run at all (nothing to show a model)."""
    from align.verify.lean.repl_pool import ReplError
    from align.verify.lean.strict import lex, parse_statement

    ref = parse_statement(statement)
    if not ref.ok or lex(code).broken:
        return None
    view = C.compile_view(code)
    try:
        env = worker.header_env(ref.header_text, "")
        resp = worker.command(view, env, timeout_s)
    except ReplError as e:
        return {"repl_infra": e.code}

    def pos(p):
        return None if not p else {"line": int(p["line"]), "column": int(p["column"])}

    errors = [
        {"pos": pos(m.get("pos")), "endPos": pos(m.get("endPos")), "data": m.get("data", "")}
        for m in resp.get("messages", [])
        if isinstance(m, dict) and m.get("severity") == "error"
    ]
    return {"repl_errors": errors, "repl_code": view}


def build_real(workers: int, budget_s: float, feedback: bool):
    from align.verify.lean import ReplPool, StrictConfig, StrictVerifier, find_toolchain
    from align.verify.lean.repl_pool import ReplLimits

    tc = find_toolchain()
    cfg = StrictConfig(proof_timeout_s=budget_s)
    pool = ReplPool.from_toolchain(
        tc, n_workers=workers, limits=ReplLimits(), preflight=True,
        scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "p2_verify_pool"),
    )
    verifier = StrictVerifier(pool, cfg)
    cfgd = asdict(cfg)
    cfgd["axiom_whitelist"] = sorted(cfgd["axiom_whitelist"])
    meta = {"pins": tc.pins, "config": cfgd, "workers": pool.n_workers}

    def batch_fn(items, on_result):
        def one(w, ix_item):
            ix, (stmt, code) = ix_item
            sr = verifier._verify_on(w, (stmt, code))
            fb = None
            if feedback and not sr.accepted and sr.reasons and sr.reasons[0] == "compile_error":
                fb = collect_feedback(w, stmt, code)
            on_result(ix, Result(sr.accepted, sr.reasons, sr.elapsed_s, sr.stage, _jsonable(sr.diagnostics), fb))
            return None

        out = pool.map(one, list(enumerate(items)))
        for o in out:
            if isinstance(o, BaseException):
                raise o

    return batch_fn, meta, pool


def _jsonable(d: dict) -> dict:
    return json.loads(json.dumps(d, default=str))


# =========================================================================== CLI


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--generations", nargs="+", required=True)
    ap.add_argument("--problems", nargs="+", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--inherit-lock", type=int, default=None, metavar="PARENT_PID",
                    help="the heavy-job lock is already held by this (parent) pid; verify that and do not take it again")
    ap.add_argument("--budget-s", type=float, default=180.0, help="per-proof wall budget (compile + audit); recorded in every row")
    ap.add_argument("--feedback", action="store_true", help="store full Lean errors for compile_error rejections")
    ap.add_argument("--chunk", type=int, default=128)
    ap.add_argument("--limit", type=int, default=None, help="verify at most this many pending samples (smoke tests)")
    ap.add_argument("--dry-run", action="store_true", help="count what would be verified; start no Lean")
    a = ap.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    if not 1 <= a.workers <= MAX_WORKERS:
        raise SystemExit(f"--workers must be in [1, {MAX_WORKERS}]")
    out = Path(a.out).resolve()
    if not str(out).startswith("/apps/"):
        raise SystemExit(f"--out {out} must be on /apps (never / or /tmp)")

    problems: dict[str, dict] = {}
    for p in a.problems:
        for pid, rec in C.load_problems(p).items():
            if pid in problems and problems[pid]["formal_statement"] != rec["formal_statement"]:
                raise C.ContractError(f"problem {pid!r} appears in two problem files with different statements")
            problems[pid] = rec

    if a.dry_run:
        print(json.dumps(verify_all(a.generations, problems, out, None, a.budget_s, {}, a.chunk, a.limit, dry_run=True)))
        return 0

    from align.utils import guard

    if a.inherit_lock is not None:  # a supervisor holds the lock for the whole job (scripts/p2/launch_baseline.py)
        guard.inherit_heavy_job_lock(a.inherit_lock)
    elif a.workers > 2:
        log.info("waiting for the heavy-job lock (Lean pool with %d workers)", a.workers)
        guard.acquire_heavy_job_lock("p2-verify-generations", blocking=True)
    guard.preflight_host_check(strict=True)
    batch_fn, meta, pool = build_real(a.workers, a.budget_s, a.feedback)
    stats = verify_all(a.generations, problems, out, batch_fn, a.budget_s, meta, a.chunk, a.limit)
    stats["pool"] = asdict(pool.stats)
    print(json.dumps(stats))
    return 0


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