"""G3 audit regressions that need two real ranks (CPU, gloo): U01-3.

Reuses the 2-rank harness pieces of test_trainer_dist (`_build_trainer`,
`_free_port`); only the child entry point is local, because spawn resolves the
body by module and `test_trainer_dist._BODIES` cannot see bodies defined here.
"""

from __future__ import annotations

import json
import os
import sys
import traceback

import torch.multiprocessing as mp
from test_trainer_dist import JOIN_TIMEOUT_S, WS, _build_trainer, _free_port


def _body_pin_on_rank1_only(rank: int, run_dir: str) -> dict:
    """The R5 pin flag is set on RANK 1 ONLY (a rank-local tripwire fired
    there), then both ranks train through saves at micro 4 and 8 with keep-N=1."""
    trainer, _model = _build_trainer(run_dir, init_seed=0)
    trainer.cfg.train.keep_checkpoints = 1
    if rank == 1:
        trainer._pin_next_save = True
    trainer.train()
    root = os.path.join(run_dir, "checkpoints")
    return {
        "dirs": sorted(d for d in os.listdir(root) if not d.endswith(".partial")),
        "pinned": sorted(trainer.pinned_checkpoints),
        "flag_left": bool(trainer._pin_next_save),
    }


_BODIES = {"pin_rank1": _body_pin_on_rank1_only}


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) -> 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_S)
    hung = [p for p in procs if p.is_alive()]
    for p in hung:
        p.kill()
        p.join()
    errs = []
    for r in range(WS):
        ef = os.path.join(run_dir, f"{mode}_{r}.err")
        if os.path.isfile(ef):
            errs.append(f"rank {r}:\n{open(ef).read()}")
    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_u01_3_pin_set_on_a_non_zero_rank_protects_the_checkpoint_on_rank_0(
    tmp_workdir, monkeypatch
):
    """U01-3. `_pin_next_save` is raised by a RANK-LOCAL tripwire verdict, but
    pruning and resume.pt are rank 0's. A pin that exists on rank 1 only left
    rank 0's pinned set empty: keep-N deleted the very checkpoint meant as the
    rewind point and a resume forgot the pin. The pin must be collective.
    Mutant: drop the all_reduce in save_checkpoint -> dirs == [micro_00000008]."""
    rows = _run("pin_rank1", tmp_workdir, monkeypatch)
    for r in rows:
        assert r["pinned"] == [4], rows
        assert r["flag_left"] is False, rows
    assert rows[0]["dirs"] == ["micro_00000004", "micro_00000008"], rows
    import torch

    state = torch.load(
        os.path.join(tmp_workdir, "checkpoints", "micro_00000004", "resume.pt"),
        weights_only=True,
        map_location="cpu",
    )
    assert state["pinned_checkpoints"] == [4]