#!/usr/bin/env python3
"""
Supervisor regression tests.
"""

from __future__ import annotations

import argparse
import importlib.util
import io
import queue
import signal
import sys
import tempfile
import textwrap
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

_SUPERVISOR_PATH = Path(__file__).parent.parent / "sandai" / "operator" / "supervisor.py"
_SPEC = importlib.util.spec_from_file_location("sandai.operator.supervisor", _SUPERVISOR_PATH)
if _SPEC is None or _SPEC.loader is None:
    raise RuntimeError(f"unable to load supervisor module from {_SUPERVISOR_PATH}")
_MODULE = importlib.util.module_from_spec(_SPEC)
sys.modules[_SPEC.name] = _MODULE
_SPEC.loader.exec_module(_MODULE)

RestartPolicy = _MODULE.RestartPolicy
Supervisor = _MODULE.Supervisor
main = _MODULE.main
run_supervisor_in_thread = _MODULE.run_supervisor_in_thread


def _write_worker_script(path: Path) -> None:
    path.write_text(
        textwrap.dedent(
            """\
import os
import signal
import sys
import time
from pathlib import Path

mode = sys.argv[1]
state_path = Path(sys.argv[2])
state_path.parent.mkdir(parents=True, exist_ok=True)
count = 0
if state_path.exists():
    count = int(state_path.read_text(encoding="utf-8") or "0")
count += 1
state_path.write_text(str(count), encoding="utf-8")

print(f"BOOT {count}", flush=True)

if mode == "fail_once":
    if count == 1:
        sys.exit(3)
    print("RECOVERED", flush=True)
    sys.exit(0)

if mode == "always_fail":
    sys.exit(7)

if mode == "sleep_fail":
    time.sleep(float(sys.argv[3]))
    sys.exit(7)

if mode == "exit_zero":
    print("DONE", flush=True)
    sys.exit(0)

if mode == "print_env_and_exit":
    print(f"ENV SDRUN_MODE={os.getenv('SDRUN_MODE')}", flush=True)
    print(f"ENV SDRUN_WORLD_SIZE={os.getenv('SDRUN_WORLD_SIZE')}", flush=True)
    print(f"ENV SDRUN_RANK={os.getenv('SDRUN_RANK')}", flush=True)
    print(f"ENV SDRUN_LOCAL_RANK={os.getenv('SDRUN_LOCAL_RANK')}", flush=True)
    sys.exit(0)

if mode == "stamp_and_wait":
    stamp_path = Path(sys.argv[3])
    with stamp_path.open("a", encoding="utf-8") as fh:
        fh.write(f"{os.getenv('SDRUN_RANK')} {time.time()}\\n")
    print("STAMPED", flush=True)
    while True:
        time.sleep(0.1)

if mode == "wait_signal":
    def handle(signum, _frame):
        print(f"SIGNAL {signal.Signals(signum).name}", flush=True)
        sys.exit(128 + signum)

    signal.signal(signal.SIGTERM, handle)
    signal.signal(signal.SIGINT, handle)
    print("READY", flush=True)
    while True:
        time.sleep(0.1)

raise SystemExit(f"unknown mode: {mode}")
"""
        ),
        encoding="utf-8",
    )


def _result_from_queue(result_queue: "queue.Queue[int | BaseException]", timeout: float = 10.0) -> int:
    result = result_queue.get(timeout=timeout)
    if isinstance(result, BaseException):
        raise result
    return result


def test_restart_policy_parse_and_validation() -> None:
    assert RestartPolicy.parse("never") == RestartPolicy(mode="never", max_restarts=0)
    assert RestartPolicy.parse("always") == RestartPolicy(mode="always", max_restarts=None)
    assert RestartPolicy.parse("2") == RestartPolicy(mode="count", max_restarts=2)

    try:
        RestartPolicy.parse("-1")
    except argparse.ArgumentTypeError:
        pass
    else:
        raise AssertionError("negative restart count should be rejected")


def test_supervisor_restarts_failed_worker_once_and_keeps_logs() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=1,
            restart_policy=RestartPolicy.parse("1"),
            command=[sys.executable, str(worker_script), "fail_once", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        exit_code = supervisor.run()

        assert exit_code == 0
        assert state_path.read_text(encoding="utf-8") == "2"
        assert "BOOT 1" in stdout.getvalue()
        assert "BOOT 2" in stdout.getvalue()
        assert "RECOVERED" in stdout.getvalue()
        assert "restarting (1/1)" in stderr.getvalue()


def test_supervisor_does_not_restart_clean_exit_even_with_always_policy() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=1,
            restart_policy=RestartPolicy.parse("always"),
            command=[sys.executable, str(worker_script), "exit_zero", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        exit_code = supervisor.run()

        assert exit_code == 0
        assert state_path.read_text(encoding="utf-8") == "1"
        assert "DONE" in stdout.getvalue()
        assert "restarting" not in stderr.getvalue()


def test_supervisor_does_not_restart_clean_exit_even_with_count_policy() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=1,
            restart_policy=RestartPolicy.parse("3"),
            command=[sys.executable, str(worker_script), "exit_zero", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        exit_code = supervisor.run()

        assert exit_code == 0
        assert state_path.read_text(encoding="utf-8") == "1"
        assert "restarting" not in stderr.getvalue()


def test_supervisor_exits_after_all_workers_finish_and_sums_exit_codes() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        second_state_path = tmp / "state-2.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "always_fail", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        original_start_worker = supervisor._start_worker

        def _start_worker(index: int, *, restart_count: int):
            worker_state_path = state_path if index == 0 else second_state_path
            supervisor.command = [sys.executable, str(worker_script), "always_fail", str(worker_state_path)]
            return original_start_worker(index, restart_count=restart_count)

        supervisor._start_worker = _start_worker  # type: ignore[method-assign]

        exit_code = supervisor.run()

        assert exit_code == 14
        assert stderr.getvalue().count("failure policy is ignore, waiting for remaining workers") == 2


def test_supervisor_failure_exit_shutdown_stops_other_workers_and_returns_first_exit_code() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        wait_state_path = tmp / "wait_state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "wait_signal", str(wait_state_path)],
            failure_exit_policy="shutdown",
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        original_start_worker = supervisor._start_worker

        def _start_worker(index: int, *, restart_count: int):
            if index == 0:
                supervisor.command = [sys.executable, str(worker_script), "wait_signal", str(wait_state_path)]
            else:
                supervisor.command = [sys.executable, str(worker_script), "sleep_fail", str(state_path), "0.3"]
            return original_start_worker(index, restart_count=restart_count)

        supervisor._start_worker = _start_worker  # type: ignore[method-assign]

        exit_code = supervisor.run()

        assert exit_code == 7
        assert "failure policy is shutdown, stopping remaining workers" in stderr.getvalue()
        assert "SIGNAL SIGTERM" in stdout.getvalue()


def test_supervisor_restarts_before_applying_failure_exit_policy() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        wait_state_path = tmp / "wait_state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("1"),
            command=[sys.executable, str(worker_script), "wait_signal", str(wait_state_path)],
            success_exit_policy="shutdown",
            failure_exit_policy="shutdown",
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        original_start_worker = supervisor._start_worker

        def _start_worker(index: int, *, restart_count: int):
            if index == 0:
                supervisor.command = [sys.executable, str(worker_script), "fail_once", str(state_path)]
            else:
                supervisor.command = [sys.executable, str(worker_script), "wait_signal", str(wait_state_path)]
            return original_start_worker(index, restart_count=restart_count)

        supervisor._start_worker = _start_worker  # type: ignore[method-assign]

        exit_code = supervisor.run()

        assert exit_code == 0
        assert "restarting (1/1)" in stderr.getvalue()
        assert "failure policy is shutdown" not in stderr.getvalue()
        assert "success policy is shutdown" in stderr.getvalue()


def test_supervisor_ignores_clean_exit_by_default_when_other_workers_are_running() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        wait_state_path = tmp / "wait_state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "wait_signal", str(wait_state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        original_start_worker = supervisor._start_worker
        started_workers: dict[int, object] = {}

        def _start_worker(index: int, *, restart_count: int):
            if index == 0:
                worker = original_start_worker(
                    index,
                    restart_count=restart_count,
                )
                started_workers[index] = worker
                return worker

            supervisor.command = [sys.executable, str(worker_script), "exit_zero", str(state_path)]
            try:
                worker = original_start_worker(index, restart_count=restart_count)
                started_workers[index] = worker
                return worker
            finally:
                supervisor.command = [sys.executable, str(worker_script), "wait_signal", str(wait_state_path)]

        supervisor._start_worker = _start_worker  # type: ignore[method-assign]

        thread, result_queue = run_supervisor_in_thread(supervisor)
        deadline = time.time() + 5.0
        while stdout.getvalue().count("READY") < 1 or "DONE" not in stdout.getvalue():
            if time.time() > deadline:
                raise AssertionError(f"workers did not reach expected state; stdout={stdout.getvalue()} stderr={stderr.getvalue()}")
            time.sleep(0.05)

        running_worker = started_workers[0]
        assert running_worker.process.poll() is None
        assert "SIGNAL SIGTERM" not in stdout.getvalue()

        supervisor._handle_signal(signal.SIGTERM, None)
        exit_code = _result_from_queue(result_queue, timeout=5.0)
        thread.join(timeout=1.0)

        assert exit_code == 128 + signal.SIGTERM
        assert "SIGNAL SIGTERM" in stdout.getvalue()


def test_supervisor_can_shutdown_remaining_workers_on_clean_exit() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        wait_state_path = tmp / "wait_state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "wait_signal", str(wait_state_path)],
            success_exit_policy="shutdown",
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        original_start_worker = supervisor._start_worker

        def _start_worker(index: int, *, restart_count: int):
            if index == 0:
                return original_start_worker(index, restart_count=restart_count)

            supervisor.command = [sys.executable, str(worker_script), "exit_zero", str(state_path)]
            try:
                return original_start_worker(index, restart_count=restart_count)
            finally:
                supervisor.command = [sys.executable, str(worker_script), "wait_signal", str(wait_state_path)]

        supervisor._start_worker = _start_worker  # type: ignore[method-assign]

        exit_code = supervisor.run()

        assert exit_code == 0
        assert "success policy is shutdown" in stderr.getvalue()
        assert "SIGNAL SIGTERM" in stdout.getvalue()


def test_supervisor_injects_sdrun_environment_variables() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "print_env_and_exit", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        exit_code = supervisor.run()

        assert exit_code == 0
        output = stdout.getvalue()
        assert "ENV SDRUN_MODE=true" in output
        assert "ENV SDRUN_WORLD_SIZE=2" in output
        assert "ENV SDRUN_RANK=0" in output or "ENV SDRUN_RANK=1" in output
        assert "ENV SDRUN_LOCAL_RANK=0" in output or "ENV SDRUN_LOCAL_RANK=1" in output


def test_supervisor_startup_stagger_delays_later_workers() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        stamp_path = tmp / "stamps.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=2,
            restart_policy=RestartPolicy.parse("never"),
            command=[
                sys.executable,
                str(worker_script),
                "stamp_and_wait",
                str(state_path),
                str(stamp_path),
            ],
            startup_stagger=0.3,
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        thread, result_queue = run_supervisor_in_thread(supervisor)
        deadline = time.time() + 5.0
        while stdout.getvalue().count("STAMPED") < 2:
            if time.time() > deadline:
                raise AssertionError(f"workers did not stamp in time; stdout={stdout.getvalue()} stderr={stderr.getvalue()}")
            time.sleep(0.05)

        rows = stamp_path.read_text(encoding="utf-8").strip().splitlines()
        assert len(rows) == 2
        first_ts = float(rows[0].split()[1])
        second_ts = float(rows[1].split()[1])
        assert second_ts - first_ts >= 0.25
        assert "delaying start of worker-1 by 0.300s" in stderr.getvalue()

        supervisor._handle_signal(signal.SIGTERM, None)
        exit_code = _result_from_queue(result_queue, timeout=5.0)
        thread.join(timeout=1.0)

        assert exit_code == 128 + signal.SIGTERM


def test_supervisor_forwards_sigterm_to_worker() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        worker_script = tmp / "worker.py"
        state_path = tmp / "state.txt"
        _write_worker_script(worker_script)

        stdout = io.StringIO()
        stderr = io.StringIO()
        supervisor = Supervisor(
            worker_count=1,
            restart_policy=RestartPolicy.parse("never"),
            command=[sys.executable, str(worker_script), "wait_signal", str(state_path)],
            stdout=stdout,
            stderr=stderr,
            install_signal_handlers=False,
            shutdown_timeout=2.0,
        )

        thread, result_queue = run_supervisor_in_thread(supervisor)
        deadline = time.time() + 5.0
        while "READY" not in stdout.getvalue():
            if time.time() > deadline:
                raise AssertionError(f"worker never became ready; stdout={stdout.getvalue()} stderr={stderr.getvalue()}")
            time.sleep(0.05)

        supervisor._handle_signal(signal.SIGTERM, None)
        exit_code = _result_from_queue(result_queue, timeout=5.0)
        thread.join(timeout=1.0)

        assert exit_code == 128 + signal.SIGTERM
        assert "SIGNAL SIGTERM" in stdout.getvalue()
        assert "received signal SIGTERM" in stderr.getvalue()


def test_main_requires_worker_command() -> None:
    stderr = io.StringIO()
    old_stderr = sys.stderr
    sys.stderr = stderr
    try:
        try:
            main(["-w", "2"])
        except SystemExit as exc:
            assert exc.code == 2
        else:
            raise AssertionError("main should exit when worker command is missing")
    finally:
        sys.stderr = old_stderr

    assert "worker command is required after '--'" in stderr.getvalue()


def test_main_rejects_unknown_restart_policy() -> None:
    stderr = io.StringIO()
    old_stderr = sys.stderr
    sys.stderr = stderr
    try:
        try:
            main(["--restart", "failed", "--", sys.executable, "-c", "print('x')"])
        except SystemExit as exc:
            assert exc.code == 2
        else:
            raise AssertionError("main should reject unknown restart policy")
    finally:
        sys.stderr = old_stderr

    assert "restart must be 'never', 'always', or a non-negative integer" in stderr.getvalue()


def main_entry() -> int:
    tests = [
        ("restart_policy_parse", test_restart_policy_parse_and_validation),
        ("restart_once", test_supervisor_restarts_failed_worker_once_and_keeps_logs),
        ("clean_exit_no_restart", test_supervisor_does_not_restart_clean_exit_even_with_always_policy),
        ("clean_exit_no_restart_count", test_supervisor_does_not_restart_clean_exit_even_with_count_policy),
        ("sum_exit_codes", test_supervisor_exits_after_all_workers_finish_and_sums_exit_codes),
        ("failure_exit_shutdown", test_supervisor_failure_exit_shutdown_stops_other_workers_and_returns_first_exit_code),
        ("restart_before_failure_exit", test_supervisor_restarts_before_applying_failure_exit_policy),
        ("clean_exit_ignore", test_supervisor_ignores_clean_exit_by_default_when_other_workers_are_running),
        ("clean_exit_shutdown", test_supervisor_can_shutdown_remaining_workers_on_clean_exit),
        ("sdrun_env", test_supervisor_injects_sdrun_environment_variables),
        ("startup_stagger", test_supervisor_startup_stagger_delays_later_workers),
        ("sigterm_forwarding", test_supervisor_forwards_sigterm_to_worker),
        ("main_requires_command", test_main_requires_worker_command),
        ("main_rejects_restart_failed", test_main_rejects_unknown_restart_policy),
    ]

    passed = 0
    failed = 0

    print("Supervisor regression tests")
    print("=" * 60)
    for name, test_func in tests:
        try:
            test_func()
            print(f"PASS {name}")
            passed += 1
        except Exception as exc:
            print(f"FAIL {name}: {exc}")
            import traceback

            traceback.print_exc()
            failed += 1
    print("=" * 60)
    print(f"passed={passed} failed={failed}")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main_entry())
