from __future__ import annotations

import contextlib
import json
import multiprocessing as mp
import os
import signal
import subprocess
import sys
import time
from pathlib import Path
from types import SimpleNamespace

import pytest

from examples import _launch

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]


def _worker(dp_size: int, dp_rank: int) -> None:
    directory = Path(os.environ["LAUNCH_TEST_DIR"])
    mode = os.environ["LAUNCH_TEST_MODE"]
    failed_rank = int(os.environ["LAUNCH_TEST_FAILED_RANK"])
    if mode == "ignore_sigterm" and dp_rank != failed_rank:
        signal.signal(signal.SIGTERM, signal.SIG_IGN)
    (directory / f"rank-{dp_rank}.pid").write_text(str(os.getpid()))
    if mode == "success":
        return
    if mode != "interrupt" and dp_rank == failed_rank:
        # Ensure the sibling has started (and installed its signal handler).
        deadline = time.monotonic() + 10
        while len(list(directory.glob("rank-*.pid"))) < dp_size // int(os.environ["WORLD_SIZE"]):
            if time.monotonic() >= deadline:
                raise RuntimeError("sibling failed to start")
            time.sleep(0.01)
        raise RuntimeError("intentional worker failure")
    # Stand in for a rank stuck in a collective after another rank has failed.
    time.sleep(60)


def _run_case(
    tmp_path: Path,
    *,
    mode: str,
    failed_rank: int = 1,
    node_rank: int = 0,
    num_nodes: int = 1,
    dp_size: int = 2,
) -> subprocess.CompletedProcess[str]:
    environment = dict(
        os.environ,
        PYTHONPATH=str(Path(__file__).resolve().parents[2]),
        RANK=str(node_rank),
        WORLD_SIZE=str(num_nodes),
        LAUNCH_TEST_DIR=str(tmp_path),
        LAUNCH_TEST_MODE=mode,
        LAUNCH_TEST_FAILED_RANK=str(failed_rank),
        LAUNCH_TEST_DP_SIZE=str(dp_size),
    )
    process = subprocess.Popen(
        [sys.executable, str(Path(__file__).resolve())],
        env=environment,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        start_new_session=True,
    )
    try:
        if mode == "interrupt":
            deadline = time.monotonic() + 10
            while len(list(tmp_path.glob("rank-*.pid"))) < dp_size // num_nodes:
                if process.poll() is not None or time.monotonic() >= deadline:
                    pytest.fail("workers failed to start before interruption")
                time.sleep(0.01)
            process.send_signal(signal.SIGINT)
        stdout, stderr = process.communicate(timeout=15)
    except subprocess.TimeoutExpired:
        pytest.fail("launcher did not exit after a worker failed")
    finally:
        # Also bounds the regression test when run against the broken launcher.
        with contextlib.suppress(ProcessLookupError):
            os.killpg(process.pid, signal.SIGKILL)
        process.wait()
    assert json.loads((tmp_path / "remaining_children.json").read_text()) == []
    return subprocess.CompletedProcess(process.args, process.returncode, stdout, stderr)


def test_successful_workers_are_joined(tmp_path: Path) -> None:
    result = _run_case(tmp_path, mode="success")
    assert result.returncode == 0, result.stderr
    assert len(list(tmp_path.glob("rank-*.pid"))) == 2


@pytest.mark.parametrize(
    ("node_rank", "num_nodes", "dp_size", "failed_rank"),
    [(0, 1, 2, 1), (0, 1, 2, 0), (1, 2, 4, 3)],
)
def test_worker_failure_reaps_siblings(tmp_path: Path, node_rank: int, num_nodes: int, dp_size: int, failed_rank: int) -> None:
    result = _run_case(tmp_path, mode="failure", node_rank=node_rank, num_nodes=num_nodes, dp_size=dp_size, failed_rank=failed_rank)
    assert result.returncode != 0
    assert "intentional worker failure" in result.stderr
    assert f"Some processes failed: [{failed_rank}]" in result.stderr


def test_unresponsive_sibling_is_killed(tmp_path: Path) -> None:
    result = _run_case(tmp_path, mode="ignore_sigterm")
    assert result.returncode != 0
    assert "Some processes failed: [1]" in result.stderr


def test_interruption_reaps_workers(tmp_path: Path) -> None:
    result = _run_case(tmp_path, mode="interrupt")
    assert result.returncode != 0
    assert "KeyboardInterrupt" in result.stderr


def test_start_failure_reaps_started_worker(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("RANK", "0")
    monkeypatch.setenv("WORLD_SIZE", "1")
    context = mp.get_context("spawn")
    processes = []

    def fail_start() -> None:
        raise RuntimeError("could not start worker")

    def make_process(**kwargs: object) -> mp.Process:
        process = context.Process(target=time.sleep, args=(60,))
        if processes:
            monkeypatch.setattr(process, "start", fail_start)
        processes.append(process)
        return process

    monkeypatch.setattr(_launch, "get_context", lambda _: SimpleNamespace(Process=make_process))
    try:
        with pytest.raises(RuntimeError, match="could not start worker"):
            _launch.run_dp(_worker, 2)
        assert processes[0].exitcode is not None
        assert processes[1].pid is None
    finally:
        for process in processes:
            if process.pid is not None:
                if process.is_alive():
                    process.kill()
                process.join()


def test_single_rank_runs_in_parent(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("RANK", "0")
    monkeypatch.setenv("WORLD_SIZE", "1")
    calls = []
    _launch.run_dp(lambda size, rank: calls.append((size, rank, os.getpid())), 1)
    assert calls == [(1, 0, os.getpid())]


if __name__ == "__main__":
    _launch._PROCESS_TERMINATION_TIMEOUT_S = 0.2
    try:
        _launch.run_dp(_worker, int(os.environ["LAUNCH_TEST_DP_SIZE"]))
    finally:
        remaining = [child.pid for child in mp.active_children()]
        (Path(os.environ["LAUNCH_TEST_DIR"]) / "remaining_children.json").write_text(json.dumps(remaining))