"""Regression tests for the 2026-10-14 audit findings on the DPO stage driver
(U03-1, U03-6, U11-1, U11-3, U11-4, U11-5).

Each test pins the behaviour the audit reproduced, not the source text:
  * U03-1 / U11-3  arm_result.json: rank 0 only, atomic, real checkpoint path
  * U03-6          the rerun verdict is drawn from the union of every rank's alarms
  * U11-1          a resumed arm's probe is relative to the resume point
  * U11-4          stale-checkpoint rotation happens once, on rank 0
  * U11-5          dotted overrides are applied or refused, never ignored

The multi-rank tests are 4-process CPU gloo (CUDA hidden), like
test_trainer_dist.py.
"""

from __future__ import annotations

import ast
import importlib.util
import json
import os
import socket
import traceback
from pathlib import Path

import pytest
import torch
import torch.multiprocessing as mp
from test_trainer import _tiny_trainer, tiny  # noqa: F401  (fixture re-export)

ROOT = Path(__file__).resolve().parents[1]
CONFIG = str(ROOT / "configs" / "dpo_steppairs.toml")
WS = 4
JOIN_TIMEOUT_S = 180


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


# ---------------------------------------------------------------------------
# child bodies (module level: spawn must import them)
# ---------------------------------------------------------------------------


def _body_arm_result(rank: int, work: str) -> dict:
    from align.train.dpo_stage import write_arm_result

    # Rank-dependent content, as in the real run: checkpoint is "" off rank 0,
    # and the lists differ in length (the tail of a longer torn document is
    # what made the recorded manifest unparseable).
    manifest = {
        "arm": "arm_b",
        "checkpoint": "/real/ckpt/micro_00000300" if rank == 0 else "",
        "tripwire_alarms": [{"alarm": "x" * (50 + 400 * (WS - rank))}] * (3 + rank),
    }
    path = write_arm_result(os.path.join(work, "arm_result.json"), manifest)
    return {"path": path}


def _body_rotate(rank: int, work: str) -> dict:
    from align.train.dpo_stage import rotate_stale_checkpoints

    ckpt = os.path.join(work, "run", "checkpoints")
    dst = rotate_stale_checkpoints(ckpt)
    # After the collective, EVERY rank must see the same directory state.
    return {"dst": dst, "ckpt_exists": os.path.isdir(ckpt)}


def _body_alarms(rank: int, work: str) -> dict:
    from align.train.dpo_stage import (
        apply_tripwire_response,
        load_stage_config,
        reduce_tripwire_alarms,
    )

    local = [{"alarm": "logps_chosen", "step": 7}] if rank == 2 else []
    alarms = reduce_tripwire_alarms(local)
    cfg, record = apply_tripwire_response(load_stage_config(CONFIG), alarms)
    return {"rerun": record["rerun_required"], "variant": cfg.variant, "alarms": alarms}


_BODIES = {"arm_result": _body_arm_result, "rotate": _body_rotate, "alarms": _body_alarms}


def _child(rank: int, port: int, mode: str, work: 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, work)
        with open(os.path.join(work, f"{mode}_{rank}.json"), "w") as f:
            json.dump(result, f)
    except BaseException:
        with open(os.path.join(work, f"{mode}_{rank}.err"), "w") as f:
            f.write(traceback.format_exc())
        code = 1
    finally:
        try:
            if dist.is_initialized():
                dist.destroy_process_group()
        except Exception:
            pass
    os._exit(code)


def _run(mode: str, work: 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, work)) 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 = [
        f"rank {r}:\n{open(os.path.join(work, f'{mode}_{r}.err')).read()}"
        for r in range(WS)
        if os.path.isfile(os.path.join(work, 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(work, f"{mode}_{r}.json"))) for r in range(WS)]


# ---------------------------------------------------------------------------
# U03-1 / U11-3: arm_result.json
# ---------------------------------------------------------------------------


def test_arm_result_written_only_by_rank0_in_a_4_rank_launch(tmp_workdir, monkeypatch):
    _run("arm_result", tmp_workdir, monkeypatch)
    path = os.path.join(tmp_workdir, "arm_result.json")
    got = json.load(open(path))  # unparseable (torn) on the recorded Arm B run
    assert got["checkpoint"] == "/real/ckpt/micro_00000300"
    assert len(got["tripwire_alarms"]) == 3  # rank 0's content, not a later rank's
    assert not [f for f in os.listdir(tmp_workdir) if ".tmp" in f], "tmp file left behind"


def test_non_rank0_never_writes_the_arm_result(tmp_workdir, monkeypatch):
    from align.train.dpo_stage import write_arm_result

    monkeypatch.setenv("RANK", "1")
    monkeypatch.setenv("WORLD_SIZE", "1")
    path = os.path.join(tmp_workdir, "arm_result.json")
    assert write_arm_result(path, {"checkpoint": ""}) == path
    assert not os.path.exists(path)


def test_arm_result_write_is_atomic_a_failed_write_keeps_the_previous_file(tmp_workdir, monkeypatch):
    from align.train import dpo_stage

    monkeypatch.setenv("RANK", "0")
    monkeypatch.setenv("WORLD_SIZE", "1")
    path = os.path.join(tmp_workdir, "arm_result.json")
    dpo_stage.write_arm_result(path, {"checkpoint": "good"})

    def boom(*a, **k):
        raise RuntimeError("disk died mid-dump")

    monkeypatch.setattr(dpo_stage.json, "dump", boom)
    with pytest.raises(RuntimeError, match="mid-dump"):
        dpo_stage.write_arm_result(path, {"checkpoint": "bad"})
    assert json.load(open(path)) == {"checkpoint": "good"}
    assert os.listdir(tmp_workdir) == ["arm_result.json"]


# ---------------------------------------------------------------------------
# U11-4: fresh-arm rotation
# ---------------------------------------------------------------------------


def test_stale_checkpoint_rotation_happens_once_and_no_rank_raises(tmp_workdir, monkeypatch):
    ckpt = os.path.join(tmp_workdir, "run", "checkpoints")
    os.makedirs(os.path.join(ckpt, "micro_00000100"))
    res = _run("rotate", tmp_workdir, monkeypatch)  # any rank raising fails here
    run = os.path.join(tmp_workdir, "run")
    stale = [d for d in os.listdir(run) if d.startswith("checkpoints_stale_")]
    assert len(stale) == 1
    assert os.path.isdir(os.path.join(run, stale[0], "micro_00000100"))
    assert [r["dst"] is not None for r in res] == [True, False, False, False]
    assert not any(r["ckpt_exists"] for r in res)


def test_rotation_leaves_an_empty_or_missing_dir_alone(tmp_workdir, monkeypatch):
    from align.train.dpo_stage import rotate_stale_checkpoints

    monkeypatch.setenv("RANK", "0")
    monkeypatch.setenv("WORLD_SIZE", "1")
    ckpt = os.path.join(tmp_workdir, "checkpoints")
    assert rotate_stale_checkpoints(ckpt) is None
    os.makedirs(ckpt)
    assert rotate_stale_checkpoints(ckpt) is None
    assert os.path.isdir(ckpt)


def _dpo_main_ast():
    return ast.parse((ROOT / "scripts" / "dpo.py").read_text())


def test_dpo_script_has_no_unguarded_per_rank_rename():
    renames = [
        n
        for n in ast.walk(_dpo_main_ast())
        if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "rename"
    ]
    assert not renames, "rotation must go through dpo_stage.rotate_stale_checkpoints (rank 0 + barrier)"


# ---------------------------------------------------------------------------
# U03-6: collective rerun verdict
# ---------------------------------------------------------------------------


def test_rerun_verdict_is_identical_on_every_rank_when_one_rank_alarms(tmp_workdir, monkeypatch):
    res = _run("alarms", tmp_workdir, monkeypatch)
    assert all(r["rerun"] for r in res), res  # rank-local filtering: only rank 2 would rerun
    assert len({r["variant"] for r in res}) == 1
    assert res[0]["alarms"] == res[3]["alarms"] and res[0]["alarms"][0]["rank"] == 2


def test_dpo_script_reduces_alarms_before_the_rerun_decision():
    tree = _dpo_main_ast()
    first = {}
    for n in ast.walk(tree):
        if isinstance(n, ast.Call) and isinstance(n.func, ast.Name):
            first[n.func.id] = min(first.get(n.func.id, 10**9), n.lineno)
    assert first["reduce_tripwire_alarms"] < first["apply_tripwire_response"]
    for n in ast.walk(tree):
        if isinstance(n, ast.Call) and getattr(n.func, "id", "") == "apply_tripwire_response":
            assert "summary" not in ast.dump(n.args[1]), "must be fed the reduced list, not the rank-local one"


# ---------------------------------------------------------------------------
# U11-1: resume + probe
# ---------------------------------------------------------------------------


def test_resumed_arm_probe_is_measured_after_the_resume_point(tiny, tmp_workdir):
    from align.train.dpo_stage import probe_micro_cap

    accum, probe_steps = 4, 8
    first, _ = _tiny_trainer(
        tiny, tmp_workdir, n_batches=40, grad_accum=accum, save_every_micro_steps=8, max_micro_steps=16
    )
    first.train()
    resumed, _ = _tiny_trainer(tiny, tmp_workdir, n_batches=40, grad_accum=accum, save_every_micro_steps=8)
    resumed.resume_from_latest()
    assert resumed.micro_step >= probe_steps  # the audit's precondition

    resumed.cfg.train.probe_steps = probe_steps
    resumed.cfg.train.max_micro_steps = probe_micro_cap(resumed.micro_step, probe_steps)
    resumed.train()
    assert len(resumed.step_times) == probe_steps
    probe = resumed.throughput_probe(tokens_per_step=1000)  # raised ValueError before the fix
    assert probe["n_steps"] == probe_steps // accum and probe["mean_step_time_s"] > 0


def test_probe_cap_is_relative_to_the_resume_position():
    from align.train.dpo_stage import probe_micro_cap

    assert probe_micro_cap(0, 20) == 20  # fresh run: unchanged
    assert probe_micro_cap(300, 20) == 320


def test_dpo_script_caps_the_probe_with_probe_micro_cap():
    src = (ROOT / "scripts" / "dpo.py").read_text()
    assert "max_micro_steps = probe_micro_cap(trainer.micro_step" in src
    assert "max_micro_steps = run_cfg.train.probe_steps" not in src


# ---------------------------------------------------------------------------
# U11-5: overrides
# ---------------------------------------------------------------------------


def _dpo_script():
    spec = importlib.util.spec_from_file_location("dpo_script_under_test", ROOT / "scripts" / "dpo.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


@pytest.mark.parametrize("spelling", ["--lr=1e-6", "lr=1e-6"])
def test_both_override_spellings_reach_the_stage_config(spelling):
    from align.train.dpo_stage import load_stage_config

    args = _dpo_script()._parse_args([CONFIG, spelling, "--resume"])
    assert args.overrides == ["lr=1e-6"] and args.resume is True
    assert load_stage_config(CONFIG, args.overrides).lr == 1e-6
    assert load_stage_config(CONFIG).lr == 5e-7  # no override -> the TOML value


def test_override_typo_is_refused_at_launch_by_struct_mode():
    from align.config import ConfigError
    from align.train.dpo_stage import load_stage_config

    args = _dpo_script()._parse_args([CONFIG, "--learning_rate=1e-6"])
    with pytest.raises(ConfigError):
        load_stage_config(CONFIG, args.overrides)


@pytest.mark.parametrize("bad", [["--lr", "1e-6"], ["--bogus-flag"], ["=3"]])
def test_non_override_leftovers_are_rejected_not_ignored(bad):
    with pytest.raises(SystemExit):
        _dpo_script()._parse_args([CONFIG, *bad])