import importlib.util
import io
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

spec = importlib.util.spec_from_file_location("watchdog", Path(__file__).resolve().parents[1] / "scripts/llama_watchdog.py")
watchdog = importlib.util.module_from_spec(spec)
spec.loader.exec_module(watchdog)


def busy(task=1, decoded=20):
    return [{"id": 0, "id_task": task, "is_processing": True,
             "next_token": [{"n_decoded": decoded}]}]


class WatchdogTests(unittest.TestCase):
    def test_busy_request_is_allowed_then_expires_despite_progress(self):
        state, reasons = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
        self.assertFalse(reasons)
        state, reasons = watchdog.observe(state, busy(decoded=1000), "server1", 699, 600, 2048)
        self.assertFalse(reasons)
        _, reasons = watchdog.observe(state, busy(decoded=1100), "server1", 700, 600, 2048)
        self.assertIn("600s", reasons[0])

    def test_new_request_and_restarted_server_do_not_inherit_deadline(self):
        state, _ = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
        for slots, invocation in [(busy(2), "server1"), (busy(), "server2")]:
            _, reasons = watchdog.observe(state, slots, invocation, 900, 600, 2048)
            self.assertFalse(reasons)

    def test_idle_slot_drops_old_deadline(self):
        state, _ = watchdog.observe({}, busy(), "server1", 100, 600, 2048)
        state, reasons = watchdog.observe(state, [{"id": 0, "is_processing": False}], "server1", 900, 600, 2048)
        self.assertEqual(state["active"], {})
        self.assertFalse(reasons)

    def test_excess_output_is_stopped_even_when_client_overrides_default(self):
        _, reasons = watchdog.observe({}, busy(decoded=2049), "server1", 100, 600, 2048)
        self.assertIn("2048 tokens", reasons[0])

    def test_prompt_and_generated_text_are_not_persisted(self):
        slots = busy()
        slots[0].update(prompt="private prompt", generated="private output")
        state, _ = watchdog.observe({}, slots, "server1", 100, 600, 2048)
        self.assertNotIn("private", str(state))


class WatchdogIntegrationTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.environment = patch.dict(watchdog.os.environ, {"XDG_STATE_HOME": self.temp.name,
            "LLAMA_MAX_REQUEST_SECONDS": "600", "LLAMA_MAX_OUTPUT_TOKENS": "2048"})
        self.environment.start()
        self.addCleanup(self.environment.stop)
        self.path = Path(self.temp.name) / "hermes-stack-maintenance/llama-watchdog.json"
        self.boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip()
        self.calls = []

    def control(self, *args):
        self.calls.append(args)
        return "ActiveState=active\nInvocationID=test-server" if args[0] == "show" else ""

    def tick(self, now, slots=None, error=False):
        response = io.BytesIO(json.dumps(slots or []).encode())
        with patch.object(watchdog, "systemctl", self.control), \
             patch.object(watchdog.time, "monotonic", return_value=now), \
             patch.object(watchdog, "urlopen", side_effect=OSError() if error else None, return_value=response):
            watchdog.main()

    def restarts(self):
        return [call for call in self.calls if call[0] == "restart"]

    def test_deadline_restarts_actual_configured_service_once(self):
        self.tick(100, busy())
        self.tick(699, busy())
        self.assertFalse(self.restarts())
        self.tick(700, busy())
        self.assertEqual(self.restarts(), [("restart", "--no-block", "llama-qwen.service")])
        self.tick(701, busy(decoded=3000))
        self.assertEqual(len(self.restarts()), 1)
        self.assertEqual(self.path.stat().st_mode & 0o777, 0o600)

    def test_only_persistent_probe_failure_triggers_recovery(self):
        self.tick(100, error=True)
        self.tick(699, error=True)
        self.assertFalse(self.restarts())
        self.tick(700, error=True)
        self.assertEqual(len(self.restarts()), 1)

    def test_successful_probe_resets_failure_deadline(self):
        self.tick(100, error=True)
        self.tick(500)
        self.tick(700, error=True)
        self.assertFalse(self.restarts())

    def test_inactive_service_never_started_or_probed(self):
        with patch.object(watchdog, "systemctl", return_value="ActiveState=inactive\nInvocationID=") as control, \
             patch.object(watchdog, "urlopen") as probe:
            watchdog.main()
        probe.assert_not_called()
        self.assertEqual(control.call_count, 1)

    def test_reboot_clears_old_monotonic_deadlines_and_cooldown(self):
        self.path.parent.mkdir()
        self.path.write_text(json.dumps({"boot": "previous-boot", "invocation": "test-server",
            "active": {"0:1": 100}, "last_restart": 10000, "unreachable_since": 50}))
        self.tick(100, busy())
        self.tick(700, busy())
        self.assertEqual(len(self.restarts()), 1)
        self.assertEqual(json.loads(self.path.read_text())["boot"], self.boot)


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