"""Regression tests for the 2026-10-14 audit, group G7 (p2 non-frozen files).

U12-1  --redo must delete the stage's artefacts (and every downstream one), not just the markers.
U12-5  the dev split must be decontaminated against test/MOBench/pool by INFORMAL text too.
U12-7  `align.model.weight_average` main() must take the heavy lock / preflight / thread cap.
U06-1  Qwen3 `_attn_bias` returns None (flash path) for all-ones and right-padded masks.
U06-2  weight_average fsyncs shards/index/marker/dirs before the rename, and a re-run verifies shards.

CPU only, nothing takes the real heavy lock.
Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_audit_g7_p2_nonfrozen_fixes.py -q
"""

from __future__ import annotations

import importlib.util
import json
import os
from pathlib import Path

import pytest
import torch

from test_expert_iteration import EI, FakeRunner, World, stages_done
from test_weight_average import _tensors, _write_ckpt

ROOT = Path(__file__).resolve().parents[1]


# ======================================================================== U12-1


def _round_artefacts(run: Path, r: int = 1):
    rd = run / f"round_{r:02d}"
    return {
        "gen": rd / "gen" / "gen_r0.jsonl", "ver": rd / "gen" / "ver_r0.jsonl", "rundir": rd / "gen" / "run_r0",
        "pass_rates": rd / "pass_rates.json", "data": rd / "sft_data", "sft": rd / "sft", "export": rd / "export",
        "avg": rd / "avg", "eval": rd / "eval", "report": rd / "report", "manifest": rd / "round_manifest.json",
        "select": rd / "select",
    }


def _plant_stale(art: dict):
    """Artefacts the fake children do not write but a real run leaves behind (pipeline run dirs, a stale eval file
    under a different tag)."""
    art["rundir"].mkdir(parents=True, exist_ok=True)
    (art["rundir"] / "pipeline.jsonl").write_text("{}")
    (art["eval"] / "gen_n4_c999.jsonl").write_text("{}")


def test_redo_train_removes_the_checkpoint_and_every_downstream_artefact(world):
    assert EI.main(world.argv(), FakeRunner(world)) == 0
    art = _round_artefacts(world.d / "run")
    _plant_stale(art)
    cfg = EI.parse_args(world.argv(redo="1:train"))
    EI.apply_redo(cfg, cfg.redo)
    for k in ("sft", "export", "avg", "eval", "report", "manifest"):
        assert not art[k].exists(), k
    assert not (art["export"].parent / "export_check").exists()
    for k in ("gen", "ver", "pass_rates", "data", "select"):  # upstream of train: untouched
        assert art[k].exists(), k
    assert stages_done(world.d / "run") == {"select", "gen_verify", "measure", "sft_data"}


def test_redo_gen_verify_removes_samples_and_everything_after_it(world):
    assert EI.main(world.argv(), FakeRunner(world)) == 0
    art = _round_artefacts(world.d / "run")
    _plant_stale(art)
    cfg = EI.parse_args(world.argv(redo="1:gen_verify"))
    EI.apply_redo(cfg, cfg.redo)
    for k in ("gen", "ver", "rundir", "pass_rates", "data", "sft", "export", "avg", "eval", "report", "manifest"):
        assert not art[k].exists(), k
    assert art["select"].exists()


def test_redo_reruns_train_from_scratch_and_regenerates_eval(world):
    """End to end: the children must not find the previous model's checkpoint or eval samples."""
    assert EI.main(world.argv(), FakeRunner(world)) == 0
    seen = {}

    class Watching(FakeRunner):
        def run(self, argv, env, log):
            k = self.key(argv)
            L = EI.Layout(cfg, 1)
            if k == "train":
                seen["ckpt_at_train"] = EI.latest_checkpoint(L.sft)
            if k == "eval":
                seen["eval_files_at_eval"] = sorted(p.name for p in L.eval.glob("*")) if L.eval.exists() else []
            return super().run(argv, env, log)

    cfg = EI.parse_args(world.argv(redo="1:train"))
    assert EI.main(world.argv(redo="1:train"), Watching(world)) == 0
    assert seen["ckpt_at_train"] is None
    assert seen["eval_files_at_eval"] == []


def test_redo_of_a_model_stage_also_wipes_later_rounds_but_a_cosmetic_redo_does_not(world):
    w = World(world.d)
    assert EI.main(w.argv(rounds=2), FakeRunner(w)) == 0
    r2 = _round_artefacts(w.d / "run", 2)
    assert r2["sft"].exists() and r2["eval"].exists()
    cfg = EI.parse_args(w.argv(rounds=2, redo="1:report"))
    EI.apply_redo(cfg, cfg.redo)
    assert r2["sft"].exists() and r2["eval"].exists()  # round 2 consumed round 1's model, which a report redo leaves intact
    cfg = EI.parse_args(w.argv(rounds=2, redo="1:train"))
    EI.apply_redo(cfg, cfg.redo)
    for k in ("gen", "ver", "data", "sft", "export", "avg", "eval", "manifest"):
        assert not r2[k].exists(), k


@pytest.fixture
def world(tmp_workdir):
    return World(Path(tmp_workdir))


# ======================================================================== U12-5


def _load_bv():
    spec = importlib.util.spec_from_file_location("build_valid_g7", ROOT / "scripts" / "p2" / "build_valid.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


INFORMAL = ("Let n be a positive integer such that the sum of the first n odd numbers equals the square of the "
            "number of its divisors plus seventeen hundred and forty two")


def _prob(pid, split, informal, sig):
    C = EI.C
    stmt = f"import Mathlib\n\ntheorem t_{pid.replace('-', '_')} {sig} := by sorry"
    return {"problem_id": pid, "source": "s", "split": split, "formal_statement": stmt,
            "informal_prefix": f"/-- {informal}-/\n" if informal else None, "statement_sha": C.statement_sha(stmt)}


def test_build_valid_drops_valid_problems_whose_informal_text_overlaps_the_pool(tmp_workdir):
    bv = _load_bv()
    d = Path(tmp_workdir)
    C = bv.C
    pd = d / "problems"
    pd.mkdir()
    C.write_jsonl_atomic(pd / "minif2f.jsonl", [_prob("mf-test", "test", "unrelated words about triangles and circles", "(x : Nat) : x = x")])
    C.write_jsonl_atomic(pd / "mobench.jsonl", [_prob("mo", "test", None, "(y : Nat) : y = y")])
    # same informal problem, DIFFERENT formalisation (different signature/name): no sha/alpha/name layer sees it
    C.write_jsonl_atomic(pd / "train_pool.jsonl", [_prob("Goedel-Pset-1", "train", INFORMAL, "(a : Nat) (h : a = 1) : a + 1 = 2")])
    src = d / "valid.jsonl"
    rows = [
        {"id": "valid_leak", "formal_statement": "theorem valid_leak (n : Nat) (h : n = 5) : n * 2 = 10 := sorry", "informal_stmt": INFORMAL},
        {"id": "valid_ok", "formal_statement": "theorem valid_ok (m : Nat) (h : m = 6) : m * 2 = 12 := sorry",
         "informal_stmt": "Show that twice six is twelve when m is six"},
    ]
    C.write_jsonl_atomic(src, rows)
    out = d / "valid_out.jsonl"
    assert bv.main(["--src", str(src), "--problems-dir", str(pd), "--out", str(out), "--no-elab"]) == 0
    kept = [r["problem_id"] for r in C.iter_jsonl(out)]
    assert kept == ["valid_ok"]
    man = json.loads(Path(str(out) + ".manifest.json").read_text())
    assert [x["id"] for x in man["dropped"]] == ["valid_leak"] and "informal" in man["dropped"][0]["reason"]


def test_driver_refuses_a_pool_whose_informal_text_overlaps_the_eval_problems(world):
    from test_p2_support import write_jsonl

    ev = json.loads((world.d / "minif2f.jsonl").read_text().splitlines()[0])
    ev["informal_prefix"] = f"/-- {INFORMAL}-/\n"
    ev_rows = [ev] + [json.loads(x) for x in (world.d / "minif2f.jsonl").read_text().splitlines()[1:]]
    write_jsonl(world.d / "minif2f.jsonl", ev_rows)
    cfg = EI.parse_args(world.argv())
    assert EI.split_problems(cfg) is None  # the pool is clean
    pool = [json.loads(x) for x in (world.d / "pool.jsonl").read_text().splitlines()]
    pool[0]["informal_prefix"] = f"/-- {INFORMAL}-/\n"
    write_jsonl(world.d / "pool.jsonl", pool)
    err = EI.split_problems(cfg)
    assert err and "informal" in err and pool[0]["problem_id"] in err


# ======================================================================== U12-7


def test_weight_average_main_takes_the_heavy_guards_and_caps_threads(tmp_workdir, monkeypatch):
    from align.model import weight_average as wa
    from align.utils import guard

    calls = []
    monkeypatch.setattr(guard, "apply_resource_guards", lambda *a, **k: calls.append((a, k)))
    threads = []
    monkeypatch.setattr(torch, "set_num_threads", lambda n: threads.append(n))
    A, B = _tensors(1, torch.float32), _tensors(2, torch.float32)
    _write_ckpt(f"{tmp_workdir}/a", A)
    _write_ckpt(f"{tmp_workdir}/b", B)
    argv = ["--base", f"{tmp_workdir}/a", "--ft", f"{tmp_workdir}/b", "--ft-weight", "0.5", "--out", f"{tmp_workdir}/o"]
    assert wa.main(argv) == 0
    assert len(calls) == 1 and calls[0][1].get("heavy") is True and calls[0][1].get("label") == "weight-average"
    assert threads and max(threads) <= calls[0][1]["cpu_threads"] <= 8


def test_weight_average_guards_run_before_any_checkpoint_is_read(tmp_workdir, monkeypatch):
    """A refused preflight must leave the output untouched (no .partial, no reads)."""
    from align.model import weight_average as wa
    from align.utils import guard

    def refuse(*a, **k):
        raise guard.PreflightRefused("host busy")

    monkeypatch.setattr(guard, "apply_resource_guards", refuse)
    A, B = _tensors(1, torch.float32), _tensors(2, torch.float32)
    _write_ckpt(f"{tmp_workdir}/a", A)
    _write_ckpt(f"{tmp_workdir}/b", B)
    with pytest.raises(guard.PreflightRefused):
        wa.main(["--base", f"{tmp_workdir}/a", "--ft", f"{tmp_workdir}/b", "--ft-weight", "0.5", "--out", f"{tmp_workdir}/o"])
    assert not os.path.exists(f"{tmp_workdir}/o") and not os.path.exists(f"{tmp_workdir}/o.partial")


# ======================================================================== U06-2


def test_weight_average_fsyncs_everything_before_the_rename(tmp_workdir, monkeypatch):
    from align.model import weight_average as wa

    A, B = _tensors(1, torch.bfloat16), _tensors(2, torch.bfloat16)
    _write_ckpt(f"{tmp_workdir}/a", A)
    _write_ckpt(f"{tmp_workdir}/b", B)
    events = []
    real_fsync, real_rename = os.fsync, os.rename

    def fsync(fd):
        try:
            events.append(("fsync", os.readlink(f"/proc/self/fd/{fd}")))
        except OSError:
            events.append(("fsync", "?"))
        return real_fsync(fd)

    def rename(src, dst, *a, **k):
        events.append(("rename", str(src), str(dst)))
        return real_rename(src, dst, *a, **k)

    monkeypatch.setattr(os, "fsync", fsync)
    monkeypatch.setattr(os, "rename", rename)
    out = f"{tmp_workdir}/o"
    rep = wa.average_checkpoints(f"{tmp_workdir}/a", f"{tmp_workdir}/b", out, 0.5)
    ri = next(i for i, e in enumerate(events) if e[0] == "rename" and e[2] == out)
    before = {e[1] for e in events[:ri] if e[0] == "fsync"}
    after = {e[1] for e in events[ri + 1:] if e[0] == "fsync"}
    part = out + ".partial"
    shards = [f"{part}/model-{i:05d}-of-{rep['shards']:05d}.safetensors" for i in range(1, rep["shards"] + 1)]
    for p in shards + [f"{part}/{wa.INDEX_NAME}", f"{part}/{wa.MARKER}", part]:
        assert p in before, f"{p} not fsynced before the rename"
    assert os.path.dirname(out) in after, "parent directory not fsynced after the rename"


def _finished(tmp_workdir):
    from align.model import weight_average as wa

    A, B = _tensors(1, torch.bfloat16), _tensors(2, torch.bfloat16)
    _write_ckpt(f"{tmp_workdir}/a", A)
    _write_ckpt(f"{tmp_workdir}/b", B)
    out = f"{tmp_workdir}/o"
    wa.average_checkpoints(f"{tmp_workdir}/a", f"{tmp_workdir}/b", out, 0.5)
    return wa, out


def test_rerun_on_a_finished_dir_is_a_noop_when_intact(tmp_workdir):
    wa, out = _finished(tmp_workdir)
    assert wa.average_checkpoints(f"{tmp_workdir}/a", f"{tmp_workdir}/b", out, 0.5)["skipped"] is True


@pytest.mark.parametrize("damage", ["truncate", "zero_length", "missing"])
def test_rerun_refuses_a_marked_dir_whose_shard_is_damaged(tmp_workdir, damage):
    wa, out = _finished(tmp_workdir)
    shard = sorted(f for f in os.listdir(out) if f.endswith(".safetensors"))[0]
    p = os.path.join(out, shard)
    if damage == "missing":
        os.remove(p)
    else:
        size = os.path.getsize(p)
        with open(p, "r+b") as f:
            f.truncate(0 if damage == "zero_length" else size - 7)
    with pytest.raises(wa.AverageError, match="incomplete|damaged|missing|size"):
        wa.average_checkpoints(f"{tmp_workdir}/a", f"{tmp_workdir}/b", out, 0.5)


# ======================================================================== U06-1


def _q3():
    from align.model import qwen3 as q3

    return q3


def test_attn_bias_is_none_for_all_ones_and_right_padded_masks():
    q3 = _q3()
    bias = q3.Qwen3Model._attn_bias
    t = 6
    ones = torch.ones(2, t, dtype=torch.long)
    right = ones.clone()
    right[1, 4:] = 0
    right[0, 5:] = 0
    assert bias(ones, t, torch.float32, "cpu") is None
    assert bias(right, t, torch.float32, "cpu") is None
    assert bias(ones.bool(), t, torch.float32, "cpu") is None


def test_attn_bias_is_kept_for_left_padding_and_interior_holes():
    q3 = _q3()
    bias = q3.Qwen3Model._attn_bias
    t = 6
    left = torch.ones(2, t, dtype=torch.long)
    left[1, :2] = 0
    hole = torch.ones(1, t, dtype=torch.long)
    hole[0, 3] = 0
    for m in (left, hole):
        b = bias(m, t, torch.float32, "cpu")
        assert b is not None and b.shape == (m.shape[0], 1, t, t)
    # a masked key is really masked for a real query
    assert bias(left, t, torch.float32, "cpu")[1, 0, 5, 0] < -1e30


def test_attn_bias_still_validates_the_mask_length():
    q3 = _q3()
    with pytest.raises(ValueError, match="attention_mask covers"):
        q3.Qwen3Model._attn_bias(torch.ones(1, 5, dtype=torch.long), 6, torch.float32, "cpu")