"""Tests for the server lifecycle module (backends/servers.py)."""

import io
import os
import signal
import tempfile
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from unittest.mock import MagicMock, patch

from backends import ServerSpec, servers


class StartTests(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self._tmp.name)
        # A fake executable so Path(argv[0]).exists() passes.
        self.exe = self.dir / "fake_server"
        self.exe.write_bytes(b"#!/bin/sh\n")
        self.spec = ServerSpec("test", "http://127.0.0.1:9999",
                               [str(self.exe), "--port", "9999"])

    def tearDown(self):
        self._tmp.cleanup()

    def test_returns_false_when_executable_missing(self):
        spec = ServerSpec("nope", "http://127.0.0.1:1", ["/no/such/binary"])
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertFalse(servers.start(spec))

    def test_noop_when_already_running(self):
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("backends.common.server_running", return_value=True), \
                patch("subprocess.Popen") as mk:
            self.assertTrue(servers.start(self.spec))
        mk.assert_not_called()

    def test_refuses_to_double_start_while_previous_boot_is_alive(self):
        """A live pid file blocks a second spawn of the same server.

        A previous ``start`` whose server is still booting must not be
        orphaned by a duplicate process on the same port.
        """
        pid_file = self.dir / "test-server.pid"
        pid_file.write_text("4242", encoding="utf-8")
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True), \
                patch("subprocess.Popen") as mk, \
                patch("backends.common.server_running", return_value=False):
            ok = servers.start(self.spec)
        self.assertFalse(ok)
        mk.assert_not_called()
        self.assertTrue(pid_file.exists())

    def test_happy_path_spawns_and_polls_until_ready(self):
        proc = MagicMock()
        proc.pid = 4242
        proc.poll.return_value = None  # process still running
        # server_running: False on the pre-check, True once inside the loop.
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc) as mk, \
                patch("backends.common.server_running",
                      side_effect=[False, True]), \
                patch("time.sleep"):
            ok = servers.start(self.spec)
        self.assertTrue(ok)
        mk.assert_called_once()
        # Pid file written (first field is the pid; the optional second
        # field is the start-time ownership token, absent on this platform).
        self.assertEqual(
            (self.dir / "test-server.pid").read_text(encoding="utf-8")
            .split()[0], "4242")

    def test_pid_file_records_a_start_time_token_where_available(self):
        proc = MagicMock()
        proc.pid = 5150
        proc.poll.return_value = None
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch.object(servers, "_process_start_token",
                             return_value="12345"), \
                patch("backends.common.server_running",
                      side_effect=[False, True]), \
                patch("time.sleep"):
            self.assertTrue(servers.start(self.spec))
        fields = (self.dir / "test-server.pid") \
            .read_text(encoding="utf-8").split()
        self.assertEqual(fields, ["5150", "12345"])

    def test_recycled_pid_with_mismatched_token_is_not_ours(self):
        # The pid is alive but its start time differs from the recorded
        # token: an unrelated process now owns this pid, so manages/alive
        # must report not-ours (and never kill it).
        pid_file = self.dir / "test-server.pid"
        pid_file.write_text("4242 111\n", encoding="utf-8")
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True), \
                patch.object(servers, "_process_start_token",
                             return_value="999"):
            self.assertFalse(servers.alive("test"))
            self.assertFalse(servers.manages([self.spec]))
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True), \
                patch.object(servers, "_process_start_token",
                             return_value="111"):
            self.assertTrue(servers.alive("test"))

    def test_returns_false_when_process_exits_early(self):
        proc = MagicMock()
        proc.pid = 99
        proc.poll.return_value = 1  # exited with code 1
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=False), \
                patch("time.sleep"):
            ok = servers.start(self.spec)
        self.assertFalse(ok)
        # Pid file cleaned up after early exit.
        self.assertFalse((self.dir / "test-server.pid").exists())

    def test_exited_event_carries_a_known_crash_hint(self):
        """The exited event's log tail is scanned for known signatures."""
        (self.dir / "test-server.log").write_text(
            "triton.compiler.errors.CompilationError:\n"
            'ValueError("type fp8e4nv not supported in this architecture. '
            'The supported fp8 dtypes are")\n', encoding="utf-8")
        proc = MagicMock()
        proc.pid = 99
        proc.poll.return_value = 1
        events = []
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=False), \
                patch("time.sleep"):
            ok = servers.start(self.spec, progress=events.append)
        self.assertFalse(ok)
        exited = next(e for e in events if e.get("kind") == "exited")
        self.assertIn("FP8", exited["hint"])
        self.assertIn("8.9", exited["hint"])

    def test_exited_event_has_no_hint_for_unknown_crashes(self):
        proc = MagicMock()
        proc.pid = 99
        proc.poll.return_value = 1
        events = []
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=False), \
                patch("time.sleep"):
            servers.start(self.spec, progress=events.append)
        exited = next(e for e in events if e.get("kind") == "exited")
        self.assertIsNone(exited["hint"])

    def test_boot_hint_reads_the_log_tail(self):
        self.assertIsNone(servers._boot_hint(["everything fine"]))
        self.assertIn("FP8", servers._boot_hint(
            ["x", 'ValueError("type fp8e4nv not supported in this '
            'architecture")', "y"]))
        self.assertIsNone(servers._boot_hint([]))

    def test_boot_hint_names_missing_companion_packages(self):
        # A model whose companion pip packages are absent from the backend
        # venv (e.g. sglang-omni's Qwen3-TTS models need qwen_tts) dies on
        # the import; the hint says what to do about it.
        self.assertIn("companion", servers._boot_hint(
            ['File "...", in resolve_checkpoint',
             "ModuleNotFoundError: No module named 'qwen_tts'"]))

    def test_console_progress_prints_the_hint(self):
        out = io.StringIO()
        with redirect_stdout(out):
            servers._console_progress({
                "kind": "exited", "name": "test", "returncode": 1,
                "log_tail": ["boom"], "hint": "FP8 needs compute "
                                              "capability 8.9+"})
        self.assertIn("hint: FP8 needs compute capability 8.9+",
                      out.getvalue())

    def test_returns_false_on_timeout(self):
        proc = MagicMock()
        proc.pid = 7
        proc.poll.return_value = None
        # time.time: first call < deadline loop entry, then past deadline.
        times = iter([0.0, float(servers.SERVER_START_TIMEOUT + 1)])
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=False), \
                patch("time.sleep"), \
                patch("time.time", side_effect=lambda: next(times)):
            ok = servers.start(self.spec)
        self.assertFalse(ok)

    def _boot_proc(self):
        proc = MagicMock()
        proc.pid = 4242
        proc.poll.return_value = None
        return proc

    def test_cwd_passed_to_popen(self):
        """A spec with a cwd spawns the server in that working directory.

        audio.cpp resolves model_specs/<family>.json relative to its
        process working directory, so the hub must start it from the
        checkout.
        """
        spec = ServerSpec("test", "http://127.0.0.1:9999",
                          [str(self.exe)], cwd=Path("/opt/audio.cpp"))
        proc = self._boot_proc()
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc) as mk, \
                patch("backends.common.server_running",
                      side_effect=[False, True]), \
                patch("time.sleep"):
            ok = servers.start(spec)
        self.assertTrue(ok)
        kwargs = mk.call_args.kwargs
        self.assertEqual(kwargs.get("cwd"), "/opt/audio.cpp")

    def test_identity_spec_waits_for_http_identity(self):
        """Readiness needs the server to answer HTTP as its identity.

        A TCP-accepting but still-booting server (lazy model load, slow
        listen-before-serve) must not count as ready.
        """
        spec = ServerSpec("test", "http://127.0.0.1:9999",
                          [str(self.exe)], identity="audiocpp")
        proc = self._boot_proc()
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=True), \
                patch.object(servers.probe, "identify_server",
                             side_effect=[None, None, "audiocpp"]), \
                patch("time.sleep"):
            ok = servers.start(spec)
        self.assertTrue(ok)

    def test_faster_identity_requires_model_loaded(self):
        """The faster identity additionally waits for /health model_loaded."""
        spec = ServerSpec("test", "http://127.0.0.1:9999",
                          [str(self.exe)], identity="faster")
        proc = self._boot_proc()
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=True), \
                patch.object(servers.probe, "identify_server",
                             return_value="faster"), \
                patch.object(servers.probe, "faster_model_loaded",
                             side_effect=[False, True]), \
                patch("time.sleep"):
            ok = servers.start(spec)
        self.assertTrue(ok)

    def test_progress_receives_boot_events(self):
        events = []
        proc = self._boot_proc()
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running",
                      side_effect=[False, True]), \
                patch("time.sleep"):
            ok = servers.start(self.spec, progress=events.append)
        self.assertTrue(ok)
        kinds = [event["kind"] for event in events]
        self.assertEqual(kinds, ["starting", "ready"])
        self.assertEqual(events[0]["pid"], 4242)
        self.assertIn("--port", events[0]["argv"])

    def test_cancel_aborts_boot_kills_process_and_reports(self):
        import threading
        cancel = threading.Event()
        cancel.set()
        proc = self._boot_proc()
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch("subprocess.Popen", return_value=proc), \
                patch("backends.common.server_running", return_value=False), \
                patch.object(servers, "_kill_pid", return_value=True) as mk, \
                patch("time.sleep"):
            ok = servers.start(self.spec, cancel=cancel)
        self.assertFalse(ok)
        mk.assert_called_once_with(4242)
        self.assertFalse((self.dir / "test-server.pid").exists())

    def test_console_progress_prints_events(self):
        import io
        from contextlib import redirect_stdout
        buf = io.StringIO()
        with redirect_stdout(buf):
            servers._console_progress({"kind": "running", "name": "test",
                                       "url": "http://127.0.0.1:9999"})
            servers._console_progress({"kind": "error", "message": "boom"})
        out = buf.getvalue()
        self.assertIn("already running", out)
        self.assertIn("boom", out)


class StopTests(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self._tmp.name)

    def tearDown(self):
        self._tmp.cleanup()

    def _write_pid(self, name, pid):
        (self.dir / f"{name}-server.pid").write_text(str(pid),
                                                     encoding="utf-8")

    def test_returns_false_when_no_pid_file(self):
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertFalse(servers.stop("test"))

    def test_stops_alive_process_and_removes_pid_file(self):
        self._write_pid("test", 1234)
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True), \
                patch.object(servers, "_kill_pid", return_value=True) as mk:
            ok = servers.stop("test")
        self.assertTrue(ok)
        mk.assert_called_once_with(1234)
        self.assertFalse((self.dir / "test-server.pid").exists())

    def test_already_dead_returns_true_and_cleans_pid_file(self):
        self._write_pid("test", 1234)
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=False), \
                patch.object(servers, "_kill_pid") as mk:
            ok = servers.stop("test")
        self.assertTrue(ok)
        mk.assert_not_called()
        self.assertFalse((self.dir / "test-server.pid").exists())

    def test_corrupt_pid_file_returns_false_and_cleans(self):
        (self.dir / "test-server.pid").write_text("not-a-number",
                                                   encoding="utf-8")
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertFalse(servers.stop("test"))
        self.assertFalse((self.dir / "test-server.pid").exists())


class ReapTests(unittest.TestCase):
    """_reap_exited: an exited child must stop counting as alive."""

    def test_true_when_child_exited(self):
        with patch("os.waitpid", return_value=(4242, 0)) as mk:
            self.assertTrue(servers._reap_exited(4242))
        mk.assert_called_once_with(4242, os.WNOHANG)

    def test_false_while_still_running(self):
        # (0, 0) is WNOHANG's "still running" answer.
        with patch("os.waitpid", return_value=(0, 0)):
            self.assertFalse(servers._reap_exited(4242))

    def test_false_when_not_our_child(self):
        with patch("os.waitpid", side_effect=ChildProcessError):
            self.assertFalse(servers._reap_exited(4242))


class KillPidTests(unittest.TestCase):
    """_kill_pid: the reap check ends the grace wait before SIGKILL."""

    def test_reaped_child_ends_wait_without_sigkill(self):
        with patch("os.getpgid", return_value=4242), \
                patch("os.killpg") as mk_killpg, \
                patch("os.waitpid", return_value=(4242, 0)) as mk_waitpid, \
                patch("time.sleep") as mk_sleep:
            ok = servers._kill_pid(4242)
        self.assertTrue(ok)
        mk_killpg.assert_called_once_with(4242, signal.SIGTERM)
        mk_waitpid.assert_called_once_with(4242, os.WNOHANG)
        mk_sleep.assert_not_called()

    def test_escalates_to_sigkill_when_child_stays_alive(self):
        with patch("os.getpgid", return_value=4242), \
                patch("os.killpg") as mk_killpg, \
                patch("os.waitpid", return_value=(0, 0)), \
                patch("time.sleep"):
            ok = servers._kill_pid(4242)
        self.assertTrue(ok)
        calls = mk_killpg.call_args_list
        self.assertEqual(calls[0].args, (4242, signal.SIGTERM))
        self.assertEqual(calls[-1].args, (4242, signal.SIGKILL))

    def test_foreign_child_falls_back_to_group_probe(self):
        # ChildProcessError from waitpid (not our child / already reaped):
        # the killpg(0) probe decides; a vanished group ends the wait.
        with patch("os.getpgid", return_value=4242), \
                patch("os.killpg",
                      side_effect=[None, ProcessLookupError]) as mk_killpg, \
                patch("os.waitpid", side_effect=ChildProcessError), \
                patch("time.sleep"):
            ok = servers._kill_pid(4242)
        self.assertTrue(ok)
        calls = mk_killpg.call_args_list
        self.assertEqual(calls[0].args, (4242, signal.SIGTERM))
        self.assertEqual(calls[-1].args, (4242, 0))


class ManagesTests(unittest.TestCase):
    """manages(): a live recorded pid marks a server as ours."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self._tmp.name)
        self.specs = [ServerSpec("test", "http://127.0.0.1:9999", [])]

    def tearDown(self):
        self._tmp.cleanup()

    def _write_pid(self, name, pid):
        (self.dir / f"{name}-server.pid").write_text(str(pid),
                                                     encoding="utf-8")

    def test_false_without_pid_file(self):
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertFalse(servers.manages(self.specs))

    def test_true_with_live_recorded_pid(self):
        self._write_pid("test", 4242)
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True):
            self.assertTrue(servers.manages(self.specs))

    def test_false_with_dead_recorded_pid(self):
        self._write_pid("test", 4242)
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=False):
            self.assertFalse(servers.manages(self.specs))

    def test_false_with_corrupt_pid_file(self):
        (self.dir / "test-server.pid").write_text("junk",
                                                  encoding="utf-8")
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive",
                             return_value=True) as mk_alive:
            self.assertFalse(servers.manages(self.specs))
        mk_alive.assert_not_called()

    def test_true_when_any_spec_is_ours(self):
        other = ServerSpec("other", "http://127.0.0.1:9998", [])
        with patch.object(servers, "LOG_DIR", self.dir), \
                patch.object(servers, "_pid_alive", return_value=True):
            self._write_pid("test", 4242)
            self.assertTrue(servers.manages([other] + self.specs))
        # The live pid belongs to 'test'; 'other' alone stays unmanaged.
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertFalse(servers.manages([other]))


class PidForTests(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.dir = Path(self._tmp.name)

    def tearDown(self):
        self._tmp.cleanup()

    def test_none_when_no_pid_file(self):
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertIsNone(servers.pid_for("test"))

    def test_returns_pid_from_file(self):
        (self.dir / "test-server.pid").write_text("555\n",
                                                   encoding="utf-8")
        with patch.object(servers, "LOG_DIR", self.dir):
            self.assertEqual(servers.pid_for("test"), 555)


if __name__ == "__main__":
    unittest.main()