"""Codex audit U01 (distributed training control plane) -- confirmed-open findings.

Each test is named after the incident it prevents and FAILED on the pre-fix code.

Run: CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=2 PYTHONPATH=$PWD .venv/bin/python -m pytest -q tests/test_audit_u01_ctrlplane.py
"""

from __future__ import annotations

import gc
import glob
import json
import os
import struct
import weakref
import zipfile

import pytest
import torch
from test_model_parity import tiny_config
from test_trainer import _tiny_trainer

import align.train.trainer as trainer_mod
from align.exit_codes import read_halt_marker
from align.train.probes import GradNormTripwire, LengthDriftTripwire, LogpsChosenTripwire
from align.train.trainer import HaltError


@pytest.fixture
def tiny(tiny_model_constants):
    return tiny_config(tiny_model_constants)


def _ckpts(run_dir):
    return sorted(glob.glob(os.path.join(run_dir, "checkpoints", "micro_*")))


def _flip_payload_byte(path: str) -> None:
    """Same-size bit flip inside a stored tensor of a torch.save zip (the audit's r8)."""
    b = bytearray(open(path, "rb").read())
    with zipfile.ZipFile(path) as z:
        info = [i for i in z.infolist() if "/data/" in i.filename and i.file_size > 8][0]
    off = info.header_offset
    n, m = struct.unpack("<HH", bytes(b[off + 26 : off + 30]))
    b[off + 30 + n + m + 3] ^= 0x40
    open(path, "wb").write(bytes(b))


# ---------------------------------------------------------------------------
# U01-8
# ---------------------------------------------------------------------------


class _SpyGradNorm(GradNormTripwire):
    def __init__(self, **kw):
        super().__init__(**kw)
        self.seen: list[float] = []

    def update(self, grad_norm):
        self.seen.append(float(grad_norm))
        return super().update(grad_norm)


def test_grad_norm_halt_wire_sees_every_optimizer_step_not_every_logged_one(tiny, tmp_workdir):
    """THE 80%-BLIND HALT WIRE (U01-8). The grad_norm wire was fed only from `_log`; with
    log_every=10 and grad_accum=2 one optimizer step in five was inspected."""
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=10, grad_accum=2, armed=("grad_norm",),
        log_every_micro_steps=10, save_every_micro_steps=10,
    )
    spy = _SpyGradNorm()
    t.tripwires["grad_norm"] = spy
    t.train()
    assert t.optimizer_step == 5
    finite = [g for g in spy.seen if g == g]
    assert len(finite) == 5, f"wire inspected {len(finite)} of 5 optimizer steps: {spy.seen}"


def test_grad_norm_explosion_halts_at_the_boundary_even_between_log_steps(tiny, tmp_workdir):
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=10, grad_accum=2, armed=("grad_norm",),
        log_every_micro_steps=10, save_every_micro_steps=10,
    )
    t.tripwires["grad_norm"] = GradNormTripwire(absolute_ceiling=1e-12)
    with pytest.raises(HaltError, match="grad_norm") as ei:
        t.train()
    assert t.micro_step == 2, "the first optimizer boundary must halt, not the first log step"
    assert ei.value.exit_code == HaltError.exit_code


# ---------------------------------------------------------------------------
# U01-1
# ---------------------------------------------------------------------------


def _capture_step_grad_norms(t):
    norms = []
    orig = t.train_module.optim_step

    def spy():
        norms.append(
            sum(float(p.grad.float().norm()) ** 2 for p in t.model.parameters() if p.grad is not None) ** 0.5
        )
        return orig()

    t.train_module.optim_step = spy
    return norms


def test_tail_partial_accumulation_window_is_stepped_and_leaves_no_stale_grads(tiny, tmp_workdir):
    """THE DROPPED-TAIL DEFECT (U01-1): 3 batches at accum 2 planned ceil(3/2)=2 optimizer steps,
    took 1, and left 21 params holding the pending .grad of the unstepped micro-batch."""
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=3, grad_accum=2, save_every_micro_steps=2
    )
    s = t.train()
    assert t.optimizer_step == t.scheduler.total_steps == 2
    assert all(p.grad is None or float(p.grad.abs().sum()) == 0.0 for p in t.model.parameters())
    assert s["tail_window"]["micros"] == 1 and s["tail_window"]["action"] == "stepped"


def test_tail_window_gradient_is_rescaled_to_a_mean_over_the_micros_it_holds(tiny, tmp_workdir):
    """A 1-micro tail at accum 2 carries grad/2; stepping it as-is under-weights it. It must
    equal the gradient a grad_accum=1 trainer computes on the same batch."""
    a, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=1, grad_accum=2, save_every_micro_steps=2)
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=1, grad_accum=1, save_every_micro_steps=1)
    na, nb = _capture_step_grad_norms(a), _capture_step_grad_norms(b)
    a.train()
    b.train()
    assert len(na) == len(nb) == 1
    assert na[0] == pytest.approx(nb[0], rel=1e-4)


def test_a_capped_run_stopping_mid_window_zeroes_the_pending_grads(tiny, tmp_workdir):
    t, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "a"), n_batches=6, grad_accum=2, max_micro_steps=3,
        save_every_micro_steps=2,
    )
    s = t.train()
    assert all(p.grad is None or float(p.grad.abs().sum()) == 0.0 for p in t.model.parameters())
    assert s["tail_window"]["action"] == "dropped"


# ---------------------------------------------------------------------------
# U01-2
# ---------------------------------------------------------------------------


def test_resume_without_optimizer_state_reseeds_the_fp32_masters_from_the_loaded_weights(tiny, tmp_workdir):
    """THE REVERTING-FIRST-STEP DEFECT (U01-2): OffloadAdamW masters stay at launch weights when
    the checkpoint carries no optimizer state; the first step then overwrites params from them."""
    a, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    a.train()
    ck = _ckpts(os.path.join(tmp_workdir, "a"))[-1]
    st = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu")
    for k in ("optimizer_state", "optimizer_per_rank"):
        st.pop(k, None)
    torch.save(st, os.path.join(ck, "resume.pt"))
    os.remove(os.path.join(ck, "integrity.json")) if os.path.exists(os.path.join(ck, "integrity.json")) else None
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    b.resume_from(ck)
    o = b.train_module.optimizer
    d = max(
        float((p.detach().float() - m.detach().float()).abs().max()) for p, m in zip(o.params, o._master)
    )
    assert d == 0.0, f"masters diverge from the restored params by {d}"


# ---------------------------------------------------------------------------
# U01-3
# ---------------------------------------------------------------------------


def _saved(tiny, tmp_workdir, name="a", **kw):
    run = os.path.join(tmp_workdir, name)
    t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, save_every_micro_steps=4, **kw)
    t.train()
    return run, _ckpts(run)[-1]


def test_a_same_size_bit_flip_in_resume_pt_is_refused_on_load(tiny, tmp_workdir):
    """THE UNHASHED-PAYLOAD DEFECT (U01-3): resume.pt / optimizer files had no sha256; a flipped
    byte in a stored tensor loaded as garbage."""
    _, ck = _saved(tiny, tmp_workdir)
    _flip_payload_byte(os.path.join(ck, "resume.pt"))
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    with pytest.raises(Exception, match="sha256"):
        b.resume_from(ck)


def test_a_bit_flip_in_a_per_rank_optimizer_file_is_refused_on_load(tiny, tmp_workdir):
    def hook(cfg):
        cfg.train.per_rank_optimizer_state = True

    _, ck = _saved(tiny, tmp_workdir, cfg_hook=hook)
    _flip_payload_byte(os.path.join(ck, "optimizer_rank0.pt"))
    b, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4, cfg_hook=hook
    )
    with pytest.raises(Exception, match="sha256"):
        b.resume_from(ck)


def test_a_checkpoint_without_the_hash_manifest_still_loads_with_a_warning(tiny, tmp_workdir, caplog):
    _, ck = _saved(tiny, tmp_workdir)
    assert os.path.isfile(os.path.join(ck, "integrity.json")), "new checkpoints must record hashes"
    os.remove(os.path.join(ck, "integrity.json"))
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    with caplog.at_level("WARNING"):
        b.resume_from(ck)
    assert any("integrity" in r.message for r in caplog.records)


# ---------------------------------------------------------------------------
# U01-5
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "kw,match",
    [
        (dict(seed=7), "seed"),
        (dict(grad_accum=4), "grad_accum"),
        (dict(n_batches=6), "data"),
    ],
)
def test_resume_refuses_a_changed_seed_accum_or_data_unless_overridden(tiny, tmp_workdir, kw, match):
    """THE SILENT-RECONFIGURATION RESUME (U01-5): seed / grad_accum / data length changed between
    save and resume, nothing noticed."""
    _, ck = _saved(tiny, tmp_workdir)
    args = dict(n_batches=4, grad_accum=2, save_every_micro_steps=4)
    seed = kw.pop("seed", None)
    args.update(kw)
    if seed is not None:
        args["seed"] = seed
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), **args)
    with pytest.raises(RuntimeError, match=match):
        b.resume_from(ck)
    b.allow_resume_mismatch = True
    b.resume_from(ck)  # explicit override passes


def test_resume_of_a_checkpoint_that_predates_the_fingerprint_warns_and_loads(tiny, tmp_workdir, caplog):
    _, ck = _saved(tiny, tmp_workdir)
    p = os.path.join(ck, "resume.pt")
    st = torch.load(p, weights_only=True, map_location="cpu")
    for k in ("seed", "grad_accum", "data_fingerprint"):
        st.pop(k, None)
    torch.save(st, p)
    os.remove(os.path.join(ck, "integrity.json"))
    b, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    with caplog.at_level("WARNING"):
        b.resume_from(ck)
    assert any("fingerprint" in r.message or "predates" in r.message for r in caplog.records)


# ---------------------------------------------------------------------------
# U01-6
# ---------------------------------------------------------------------------


def test_tripwire_state_roundtrips_through_resume_pt(tiny, tmp_workdir):
    """THE RESET-ON-RELAUNCH DEFECT (U01-6): baselines and latches died with the process."""
    _, ck = _saved(tiny, tmp_workdir, armed=("grad_norm",))
    st = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu")
    assert "tripwire_per_rank" in st
    b, _ = _tiny_trainer(
        tiny, os.path.join(tmp_workdir, "b"), n_batches=4, grad_accum=2, save_every_micro_steps=4,
        armed=("grad_norm",),
    )
    assert b.tripwires["grad_norm"]._history == []
    b.resume_from(ck)
    assert len(b.tripwires["grad_norm"]._history) == 2


def test_wire_state_dicts_roundtrip_baselines_and_latches():
    w = LogpsChosenTripwire(window=3)
    for i in range(5):
        w.update(-10.0 - i, 10.0, 0.1 * i)
    w.fired, w._fired_at_row = True, 4
    w2 = LogpsChosenTripwire(window=3)
    w2.load_state_dict(json.loads(json.dumps(w.state_dict())))
    assert w2.fired and w2._baseline_nll == w._baseline_nll and w2._nll == w._nll and w2._rows == w._rows

    g = GradNormTripwire()
    g.update(1.0)
    g.fired = True
    g2 = GradNormTripwire()
    g2.load_state_dict(g.state_dict())
    assert g2.fired and g2._history == [1.0]

    d = LengthDriftTripwire(baseline=0.5, min_pairs=10)
    d.update(0.5, n_pairs=7)
    d2 = LengthDriftTripwire(baseline=0.5, min_pairs=10)
    d2.load_state_dict(d.state_dict())
    assert d2._pairs_seen == 7


# ---------------------------------------------------------------------------
# U01-9
# ---------------------------------------------------------------------------


def _golden_report(**cfg_extra):
    return {
        "per_set": {"math500": {"n": 500, "pass@1": 0.3}, "gsm8k": {"n": 500, "pass@1": 0.55}},
        "config": {"mde": 0.05, "n_problems": 1000, "eval_set_hash": "g", **cfg_extra},
    }


def test_golden_report_produced_by_a_different_eval_source_is_not_recertified(tmp_workdir):
    """THE STALE-REPORT RE-CERTIFICATION (U01-9): the ledger drops the golden pass when
    align/eval changes, and `record_golden_from_report` then re-recorded it from the OLD report."""
    from align.train.stage_gates import GateLedger, eval_source_hash, record_golden_from_report

    cur = eval_source_hash()
    ledger = GateLedger(os.path.join(tmp_workdir, "a"), eval_source_hash=cur)
    stale = os.path.join(tmp_workdir, "stale.json")
    rep = _golden_report()
    rep["eval_source_hash"] = "0" * 16
    json.dump(rep, open(stale, "w"))
    out = record_golden_from_report(ledger, report_path=stale)
    assert out["recorded"] is False and out["passed"] is None and "eval" in out["reason"]
    assert not ledger.is_passed("golden_regression")

    ok = os.path.join(tmp_workdir, "ok.json")
    rep["eval_source_hash"] = cur
    json.dump(rep, open(ok, "w"))
    assert record_golden_from_report(ledger, report_path=ok)["passed"] is True


def test_an_unstamped_legacy_golden_report_still_records_with_a_warning(tmp_workdir, caplog):
    from align.train.stage_gates import GateLedger, record_golden_from_report

    ledger = GateLedger(os.path.join(tmp_workdir, "a"))
    p = os.path.join(tmp_workdir, "legacy.json")
    json.dump(_golden_report(), open(p, "w"))
    with caplog.at_level("WARNING"):
        assert record_golden_from_report(ledger, report_path=p)["passed"] is True
    assert any("eval_source_hash" in r.message for r in caplog.records)


def test_stamp_golden_report_writes_the_producing_hash(tmp_workdir):
    from align.train.stage_gates import eval_source_hash, stamp_golden_report

    p = os.path.join(tmp_workdir, "r.json")
    json.dump(_golden_report(), open(p, "w"))
    stamp_golden_report(p)
    assert json.load(open(p))["eval_source_hash"] == eval_source_hash()


# ---------------------------------------------------------------------------
# U01-11
# ---------------------------------------------------------------------------


def test_model_is_collectable_after_the_post_train_handoff(tiny, tmp_workdir):
    """THE HANDOFF LEAK (U01-11): `_probe_target` / `trunk_probe._model` kept the training replica
    alive after scripts/sft.py set `trainer.model = None`, starving the post-train eval build."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    t.train()
    w = weakref.ref(t.model)
    t.model = None
    t.optimizer = None
    gc.collect()
    assert w() is None


def test_trainer_release_drops_every_model_reference(tiny, tmp_workdir):
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), n_batches=4, grad_accum=2, save_every_micro_steps=4)
    t.train()
    w = weakref.ref(t.model)
    t.release()
    gc.collect()
    assert w() is None and t._probe_target is None and t.trunk_probe._model is None


# ---------------------------------------------------------------------------
# U01-21
# ---------------------------------------------------------------------------


def test_a_failed_arm_time_self_test_is_a_halt_with_a_marker(tiny, tmp_workdir):
    """THE RETRIED-DETERMINISTIC-FAILURE DEFECT (U01-21): a rank-0 self-test/probe failure was a
    plain RuntimeError, so no HALT.json was written and the supervisor relaunched it."""
    run = os.path.join(tmp_workdir, "a")
    t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, armed=("logps_chosen",), save_every_micro_steps=4)
    with pytest.raises(HaltError, match="SELF-TEST FAILED"):
        t.train()
    assert read_halt_marker(run) is not None


def test_a_failed_trunk_probe_is_a_halt_with_a_marker(tiny, tmp_workdir, monkeypatch):
    run = os.path.join(tmp_workdir, "a")
    t, _ = _tiny_trainer(tiny, run, n_batches=4, grad_accum=2, save_every_micro_steps=4)

    def boom(step, **kw):
        raise RuntimeError("TRUNK MOVEMENT PROBE FAILED at step 2")

    monkeypatch.setattr(t.trunk_probe, "check", boom)
    with pytest.raises(HaltError, match="TRUNK MOVEMENT PROBE FAILED"):
        t.train()
    assert read_halt_marker(run) is not None


def test_a_missing_tripwire_input_is_still_the_original_keyerror(tiny, tmp_workdir):
    """Scope guard: only probe / self-test failures are converted; the KeyError contract of the
    per-step tripwire check (test_control_plane) is unchanged."""
    t, _ = _tiny_trainer(tiny, os.path.join(tmp_workdir, "a"), armed=("logps_chosen",))
    t.tripwires["logps_chosen"].INPUT_KEYS = ()
    with pytest.raises(KeyError):
        t.train()
    assert trainer_mod.HaltError is HaltError