from __future__ import annotations

import json
from pathlib import Path

from argus.ray.discovery import build_head_record, write_head_record_atomic
from argus.ray.worker_watchdog import Watchdog


def test_watchdog_restarts_on_first_seen_and_on_change(tmp_path: Path):
    head_file = tmp_path / "head.json"
    calls: list[list[str]] = []

    def runner(argv: list[str]) -> int:
        calls.append(argv)
        return 0

    wd = Watchdog(
        head_ip_file=str(head_file),
        node_ip="10.0.0.2",
        resources_json=json.dumps({"worker_node": 100}),
        poll_s=1,
        runner=runner,
        ray_start_block=False,
    )

    write_head_record_atomic(str(head_file), build_head_record(cluster_name="c", head_ip="1.1.1.1"))
    assert wd.tick_once() is not None
    assert any("ray" in c[0] and c[1] == "start" for c in calls)

    calls.clear()
    # Same address -> no restart
    write_head_record_atomic(str(head_file), build_head_record(cluster_name="c", head_ip="1.1.1.1"))
    wd.tick_once()
    assert calls == []

    # Address change -> restart
    write_head_record_atomic(str(head_file), build_head_record(cluster_name="c", head_ip="2.2.2.2"))
    wd.tick_once()
    assert any(c[1] == "stop" for c in calls)
    assert any(c[1] == "start" for c in calls)


def test_watchdog_block_mode_restarts_when_ray_exits(monkeypatch, tmp_path: Path):
    import argus.ray.worker_watchdog as mod

    head_file = tmp_path / "head.json"
    write_head_record_atomic(str(head_file), build_head_record(cluster_name="c", head_ip="1.1.1.1"))

    calls: list[list[str]] = []
    popen_argvs: list[list[str]] = []

    class DummyProc:
        def __init__(self, argv: list[str]):
            self.argv = argv
            self._poll: int | None = None
            self.terminated = False
            self.killed = False

        def poll(self):
            return self._poll

        def terminate(self):
            self.terminated = True

        def wait(self, timeout: float | None = None):
            # Force kill path to be exercised.
            raise TimeoutError("timeout")

        def kill(self):
            self.killed = True

    proc_holder: dict[str, DummyProc] = {}

    def fake_popen(argv, stdout=None, stderr=None):
        popen_argvs.append(list(argv))
        p = DummyProc(list(argv))
        proc_holder["p"] = p
        return p

    def runner(argv: list[str]) -> int:
        calls.append(argv)
        return 0

    monkeypatch.setattr(mod.subprocess, "Popen", fake_popen)

    wd = Watchdog(
        head_ip_file=str(head_file),
        node_ip="10.0.0.2",
        resources_json=json.dumps({"worker_node": 100}),
        poll_s=1,
        runner=runner,
        ray_start_block=True,
    )

    # First tick: should start ray with --block.
    assert wd.tick_once() is not None
    assert any(c[1] == "stop" for c in calls)
    assert len(popen_argvs) == 1
    assert "--block" in popen_argvs[0]

    # Simulate ray process exit; next tick should restart (stop + new popen).
    proc_holder["p"]._poll = 1
    calls.clear()
    wd.tick_once()
    assert any(c[1] == "stop" for c in calls)
    assert len(popen_argvs) == 2


def test_watchdog_main_once_invokes_runner(monkeypatch, tmp_path: Path):
    from argus.ray import worker_watchdog as mod

    head_file = tmp_path / "head.json"
    write_head_record_atomic(str(head_file), build_head_record(cluster_name="c", head_ip="1.1.1.1"))

    calls: list[list[str]] = []

    def runner(argv: list[str]) -> int:
        calls.append(argv)
        return 0

    monkeypatch.setattr(mod, "_default_runner", runner)
    rc = mod.main(["--head-ip-file", str(head_file), "--node-ip", "10.0.0.2", "--resources-kv", "worker_node=100", "--once"])
    assert rc == 0
    assert any(c[1] == "start" for c in calls)
