"""Regressions from the adversarial review of the G0 control plane (2026-10-12).

Each test pins one defect found by reading `trainer.py` / `dist.py` /
`supervisor.py` for the shapes that wedged this box before: a rank acting alone
while a peer waits in a collective, a checkpoint trusted without being checked,
a retry loop that outlives its reason, and a handler that does more than set a
flag. Every test fails against the code it was written for (the mutant that
removes the fix is named in the docstring).

Run: .venv/bin/python -m pytest tests/test_control_plane_review.py -q -m "not gpu and not oracle" -p no:cacheprovider
"""

from __future__ import annotations

import json
import os
import signal
import socket
import sys
import threading
import time
import traceback

import pytest
import torch
import torch.multiprocessing as mp
from test_model_parity import tiny_config
from test_trainer import _tiny_trainer

import align.utils.supervisor as sv
from align.train.trainer import checkpoint_problem


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


def _truncate_half(path):
    size = os.path.getsize(path)
    with open(path, "r+b") as f:
        f.truncate(size // 2)


def _saved_run(tiny, run_dir):
    t, _ = _tiny_trainer(tiny, run_dir, save_every_micro_steps=4)
    t.train()
    root = os.path.join(run_dir, "checkpoints")
    assert sorted(os.listdir(root)) == ["micro_00000004", "micro_00000008"]
    return t, root


# ---------------------------------------------------------------------------
# 1. the idempotent re-save must not trust a broken directory
# ---------------------------------------------------------------------------


def test_resave_replaces_a_corrupt_existing_checkpoint_and_keeps_the_evidence(tiny, tmp_workdir):
    """THE TRUSTED-CORRUPTION DEFECT. micro_8 is cut; resume falls back to
    micro_4; the resumed run reaches 8 again and used to log 'already exists'
    and write nothing, leaving a broken newest directory (and, in
    `checkpoints_written`, a broken one counted against keep-N). Mutant: restore
    `os.path.isdir(final)` as the skip test."""
    run = os.path.join(tmp_workdir, "run")
    _saved_run(tiny, run)
    root = os.path.join(run, "checkpoints")
    bad = os.path.join(root, "micro_00000008")
    _truncate_half(os.path.join(bad, "model.safetensors"))
    assert checkpoint_problem(bad) is not None

    t2, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    assert t2.resume_from_latest().endswith("micro_00000004")
    t2.train()  # walks back up to 8

    assert checkpoint_problem(bad) is None, "micro_8 must have been re-saved"
    kept = [d for d in os.listdir(root) if d.startswith("corrupt_micro_00000008.")]
    assert len(kept) == 1, "the broken directory is moved aside, not deleted"
    assert checkpoint_problem(os.path.join(root, kept[0])) is not None
    # neither resume nor the supervisor's progress rule can see the quarantined dir
    assert t2.latest_checkpoint() == bad
    assert sv.latest_checkpoint_micro_step(run, checkpoint_problem) == 8


def test_resave_still_skips_a_complete_checkpoint(tiny, tmp_workdir):
    """The idempotence this fix must keep: a COMPLETE micro_N is neither
    rewritten nor moved (re-reaching an already-saved boundary is routine)."""
    run = os.path.join(tmp_workdir, "run")
    t, root = _saved_run(tiny, run)
    final = os.path.join(root, "micro_00000008")
    before = os.stat(os.path.join(final, "resume.pt")).st_mtime_ns
    assert t.save_checkpoint() == final
    assert os.stat(os.path.join(final, "resume.pt")).st_mtime_ns == before
    assert sorted(os.listdir(root)) == ["micro_00000004", "micro_00000008"]


# ---------------------------------------------------------------------------
# 2. a malformed manifest is a verdict, not a crash
# ---------------------------------------------------------------------------


def _rewrite_manifest(d, **kv):
    state = torch.load(os.path.join(d, "resume.pt"), weights_only=True, map_location="cpu")
    state.update(kv)
    torch.save(state, os.path.join(d, "resume.pt"))
    # The test models a LEGITIMATE checkpoint with a different manifest, so keep the integrity
    # sidecar (audit U01-3) truthful; a hand-edit without this is what the sidecar refuses.
    ip = os.path.join(d, "integrity.json")
    if os.path.isfile(ip):
        from align.train import sharded_ckpt

        rec = sharded_ckpt.read_integrity(d)
        rec["resume.pt"] = sharded_ckpt.file_sha256(os.path.join(d, "resume.pt"))
        sharded_ckpt.write_integrity(d, rec)


def test_malformed_per_rank_manifest_is_skipped_by_latest_checkpoint(tiny, tmp_workdir):
    """A `per_rank_files` manifest without `world_size` raised KeyError out of
    `checkpoint_problem`, so `latest_checkpoint` (and the supervisor's progress
    rule) died on the very directory they exist to step over. Mutant: call
    `_checkpoint_problem` directly from `checkpoint_problem`."""
    t, root = _saved_run(tiny, os.path.join(tmp_workdir, "run"))
    d = os.path.join(root, "micro_00000008")
    _rewrite_manifest(d, optimizer_state={"layout": "per_rank_files"})
    assert "malformed" in checkpoint_problem(d)
    assert t.latest_checkpoint() == os.path.join(root, "micro_00000004")


def test_safetensors_header_that_is_json_but_not_an_object_is_a_problem(tiny, tmp_workdir):
    t, root = _saved_run(tiny, os.path.join(tmp_workdir, "run"))
    d = os.path.join(root, "micro_00000008")
    body = b"[]"
    with open(os.path.join(d, "model.safetensors"), "wb") as f:
        f.write(len(body).to_bytes(8, "little") + body)
    assert checkpoint_problem(d) is not None
    assert t.latest_checkpoint() == os.path.join(root, "micro_00000004")


# ---------------------------------------------------------------------------
# 3. gather-layout resume at a different world size
# ---------------------------------------------------------------------------


def test_gather_layout_resume_refuses_a_different_world_size(tiny, tmp_workdir):
    """The gather layout recorded no world size: a 1-rank resume of a 4-rank
    save silently took shard 0 of 4. The list length IS the saved world size.
    Mutant: delete the length check."""
    run = os.path.join(tmp_workdir, "run")
    _t, root = _saved_run(tiny, run)
    d = os.path.join(root, "micro_00000008")
    state = torch.load(os.path.join(d, "resume.pt"), weights_only=True, map_location="cpu")
    _rewrite_manifest(d, optimizer_per_rank=state["optimizer_per_rank"] * 4)
    t2, _ = _tiny_trainer(tiny, run, save_every_micro_steps=4)
    with pytest.raises(RuntimeError, match=r"4 rank\(s\).*world_size=1"):
        t2.resume_from(d)


# ---------------------------------------------------------------------------
# 4. supervisor: a signal that lands before the child is visible
# ---------------------------------------------------------------------------


def test_signal_landing_before_the_child_is_registered_still_reaches_the_child(tmp_workdir):
    """THE LOST-SIGNAL WINDOW. `_on_signal` forwards SIGTERM to `self._proc`; a
    signal between `Popen` returning and the assignment found None, so the
    child (the job) never got the SIGTERM that makes it checkpoint -- it was
    SIGKILLed after grace_s (exit -9). Mutant: delete the forward after Popen."""
    script = os.path.join(tmp_workdir, "child.py")
    with open(script, "w") as f:
        f.write("import time\ntime.sleep(60)\n")
    run = os.path.join(tmp_workdir, "run")
    s = sv.Supervisor([sys.executable, script], run, max_retries=0, grace_s=4.0, preflight=lambda: None)
    real = sv.subprocess.Popen

    def popen(*a, **k):
        p = real(*a, **k)
        s._on_signal(signal.SIGTERM, None)  # lands before `self._proc = ...`
        return p

    sv.subprocess.Popen = popen
    try:
        t0 = time.monotonic()
        res = s.run(install_signal_handlers=False)
        took = time.monotonic() - t0
    finally:
        sv.subprocess.Popen = real
    assert res.reason == "supervisor_terminated"
    assert res.child_codes == [-signal.SIGTERM], res.child_codes  # TERM, not the grace-expiry KILL
    assert took < 3.5


def test_signal_handler_takes_no_lock_the_interrupted_code_may_hold(tmp_workdir):
    """THE HANDLER-DEADLOCK DEFECT. A Python signal handler runs on the main
    thread between bytecodes; `threading.Event.set()` takes the event's
    non-reentrant lock, which `Event.wait()` (the whole backoff) holds for a few
    bytecodes. Simulated by running the handler on a thread while the test
    holds every lock the stop flag owns. Mutant: `self._stop = threading.Event()`."""
    run = os.path.join(tmp_workdir, "run")
    s = sv.Supervisor([sys.executable, "-c", "pass"], run, preflight=lambda: None)
    held = [lk for lk in (getattr(s._stop, "_cond", None),) if lk is not None]
    done = threading.Event()

    def handler():
        s._on_signal(signal.SIGTERM, None)
        done.set()

    for lk in held:
        lk.acquire()
    try:
        th = threading.Thread(target=handler, daemon=True)
        th.start()
        assert done.wait(2.0), "the signal handler blocked on a lock: it can deadlock the supervisor"
    finally:
        for lk in held:
            lk.release()
    assert s._stop.is_set() and s._term_signal == signal.SIGTERM
    assert s._stop.wait(0.0) is True


def test_stop_flag_wait_times_out_and_wakes_on_set():
    flag = sv._SignalFlag()
    t0 = time.monotonic()
    assert flag.wait(0.25) is False and 0.2 <= time.monotonic() - t0 < 1.0
    threading.Timer(0.15, flag.set).start()
    t0 = time.monotonic()
    assert flag.wait(5.0) is True and time.monotonic() - t0 < 1.5


# ---------------------------------------------------------------------------
# 5. two real ranks: the gloo control group honours the armed timeout, and the
#    re-save verdict is one verdict
# ---------------------------------------------------------------------------

WS = 2


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def _body_gloo_timeout(rank: int, run_dir: str) -> dict:
    from align.utils import dist as u

    u.gloo_group()  # collective: every rank builds it
    if rank == 1:
        time.sleep(14)  # alive but silent: a wedged rank, not a dead one
        return {}
    t0 = time.monotonic()
    try:
        u.all_reduce_max([1])
        return {"raised": False, "secs": time.monotonic() - t0}
    except Exception as e:  # noqa: BLE001
        return {"raised": True, "secs": time.monotonic() - t0, "type": type(e).__name__}


def _body_resave_verdict(rank: int, run_dir: str) -> dict:
    """Rank 1's filesystem view says the directory is absent (stale attribute
    cache / node-local disk); rank 0 sees a complete one. One verdict, or rank 0
    skips the save while rank 1 enters the state-gather alone."""
    import align.train.trainer as tr

    ckpt_root = os.path.join(run_dir, "ckpts")
    os.makedirs(os.path.join(ckpt_root, "micro_00000004"), exist_ok=True)
    tr.checkpoint_problem = lambda p: None  # existence == complete, for this probe
    if rank == 1:
        real = os.path.isdir
        os.path.isdir = lambda p: False if p.endswith("micro_00000004") else real(p)

    class Stub:
        pass

    stub = Stub()
    stub.ckpt_root = ckpt_root
    return {"redundant": tr.Trainer._resave_is_redundant(stub, "micro_00000004")}


_BODIES = {"gloo_timeout": _body_gloo_timeout, "resave_verdict": _body_resave_verdict}


def _child(rank: int, port: int, mode: str, run_dir: str) -> None:
    import torch.distributed as dist

    os.environ.update(
        RANK=str(rank),
        WORLD_SIZE=str(WS),
        LOCAL_RANK=str(rank),
        MASTER_ADDR="127.0.0.1",
        MASTER_PORT=str(port),
    )
    from align.utils.dist import init_distributed

    init_distributed(backend="gloo")
    code = 0
    try:
        result = _BODIES[mode](rank, run_dir)
        with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f:
            json.dump(result, f)
    except BaseException:
        with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f:
            f.write(traceback.format_exc())
        code = 1
    finally:
        try:
            dist.destroy_process_group()
        except Exception:
            pass
    if code:
        sys.exit(code)


def _run(mode: str, run_dir: str, monkeypatch, join_timeout: float = 120.0) -> list[dict]:
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "")
    ctx = mp.get_context("spawn")
    port = _free_port()
    procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)]
    for p in procs:
        p.start()
    for p in procs:
        p.join(join_timeout)
    hung = [p for p in procs if p.is_alive()]
    for p in hung:
        p.kill()
        p.join()
    errs = [
        f"rank {r}:\n{open(os.path.join(run_dir, f'{mode}_{r}.err')).read()}"
        for r in range(WS)
        if os.path.isfile(os.path.join(run_dir, f"{mode}_{r}.err"))
    ]
    assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs)
    assert not errs, "\n".join(errs)
    return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)]


def test_gloo_control_group_honours_the_armed_process_group_timeout(tmp_workdir, monkeypatch):
    """THE UNARMED-WATCHDOG DEFECT. `dist.new_group(backend="gloo")` without a
    timeout gets torch's 30-minute gloo default, not the group's: every stop /
    HALT / skip-cap decision (the boundary all-reduce) ignored
    ALIGN_PG_TIMEOUT_S. Rank 1 stays alive but silent for 14 s; with a 2 s
    timeout rank 0 must give up in seconds, not when the peer finally exits.
    Mutant: `dist.new_group(backend="gloo")` with no timeout."""
    monkeypatch.setenv("ALIGN_PG_TIMEOUT_S", "2")
    r0, _r1 = _run("gloo_timeout", tmp_workdir, monkeypatch)
    assert r0["raised"] is True, r0
    assert r0["secs"] < 9, f"control collective waited {r0['secs']:.1f}s against a 2s timeout"


def test_resave_verdict_is_rank_zeros_and_identical_on_every_rank(tmp_workdir, monkeypatch):
    """A rank whose filesystem view differs must not decide alone. Mutant:
    decide per rank from `os.path.isdir`."""
    r0, r1 = _run("resave_verdict", tmp_workdir, monkeypatch)
    assert r0 == r1 == {"redundant": True}


def test_resave_verdict_on_one_process_is_local(tmp_workdir):
    class Stub:
        ckpt_root = os.path.join(tmp_workdir, "c")

    from align.train.trainer import Trainer

    os.makedirs(Stub.ckpt_root)
    assert Trainer._resave_is_redundant(Stub, "micro_00000004") is False  # nothing there
    os.makedirs(os.path.join(Stub.ckpt_root, "micro_00000004"))
    assert Trainer._resave_is_redundant(Stub, "micro_00000004") is False  # empty dir = incomplete
    assert any(d.startswith("corrupt_micro_00000004.") for d in os.listdir(Stub.ckpt_root))