from __future__ import annotations

import asyncio
import os
import signal
import subprocess
import time
import unittest
from unittest import mock

import rlm
from rlm.bash import (
    BashHandle,
    HostProfile,
    _check_command_guards,
    adopt_process,
    bash,
    get_host_profile,
    is_pid_owned,
    kill_process,
    list_owned_pids,
    list_owned_processes,
    preflight_vram,
    reset_host_profile,
    set_host_profile,
    signal_process,
)
from rlm.host_safety import (
    ProcessIdentity,
    _ANCESTRY_CACHE,
    _OWNED_LEASES,
    clear_owned_processes,
    register_owned_process,
)
from rlm.subprocess_guard import install_subprocess_guards


class HostSafetyTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        clear_owned_processes()
        reset_host_profile()
        install_subprocess_guards()

    def tearDown(self):
        clear_owned_processes()
        reset_host_profile()

    # -----------------------------------------------------------------------
    # 1. Scoped Process Control
    # -----------------------------------------------------------------------

    async def test_owned_process_can_be_signalled_through_scoped_path(self):
        """An owned process started via bash() can be signalled through kill_process()."""
        # Start a long-running process with bash()
        handle = bash("sleep 30")
        pid = handle.pid
        self.assertTrue(is_pid_owned(pid))
        self.assertIn(pid, list_owned_pids())

        # Signal via scoped process control function
        signalled = kill_process(pid, sig=signal.SIGTERM)
        self.assertTrue(signalled)

        # Await completion to verify process was terminated
        result = await handle
        self.assertNotEqual(result.exit_code, 0)
        self.assertTrue(result.timed_out or result.exit_code < 0 or result.exit_code == 143)

    async def test_owned_process_can_be_signalled_by_handle(self):
        """kill_process() also accepts the BashHandle directly."""
        handle = bash("sleep 30")
        signalled = kill_process(handle, sig=signal.SIGTERM)
        self.assertTrue(signalled)
        result = await handle
        self.assertNotEqual(result.exit_code, 0)

    def test_foreign_pid_is_refused_with_actionable_message(self):
        """Signalling a foreign PID is refused naming ownership and action."""
        foreign_pid = 999999
        self.assertFalse(is_pid_owned(foreign_pid))

        with self.assertRaises(ValueError) as ctx:
            kill_process(foreign_pid)

        err = str(ctx.exception)
        self.assertIn("kill_process rejects pid 999999", err)
        self.assertIn("not owned by this session", err)
        self.assertIn("rlm.adopt_process", err)
        self.assertIn("h.kill()", err)
        self.assertIn("Active owned PIDs", err)

    def test_falsifier_cannot_signal_unowned_process(self):
        """Falsifier test: verify scoped primitive strictly refuses foreign PIDs even with mock/real."""
        fake_pid = 12345
        with self.assertRaises(ValueError) as ctx:
            signal_process(fake_pid, sig=signal.SIGKILL)
        self.assertIn("not owned by this session", str(ctx.exception))

    def test_user_cell_os_kill_on_foreign_pid_is_refused(self):
        """Direct os.kill from a simulated user cell on a foreign PID is blocked."""
        code = compile("import os; os.kill(999999, 15)", "<cell-foreign-kill>", "exec")
        with self.assertRaises(PermissionError) as ctx:
            eval(code, {"os": os, "__name__": "__main__"})
        self.assertIn("os.kill rejects pid 999999", str(ctx.exception))
        self.assertIn("not owned by this session", str(ctx.exception))
        self.assertIn("rlm.adopt_process", str(ctx.exception))
        self.assertIn("kill_process(pid)", str(ctx.exception))

    def test_user_cell_os_kill_on_owned_pid_is_permitted(self):
        """Direct os.kill from a user cell on an owned PID is permitted."""
        dummy_proc = subprocess.Popen(["sleep", "30"])
        try:
            register_owned_process(dummy_proc.pid, "sleep 30")
            code = compile(f"import os; os.kill({dummy_proc.pid}, 15)", "<cell-owned-kill>", "exec")
            # Should not raise PermissionError
            eval(code, {"os": os, "__name__": "__main__"})
            dummy_proc.wait(timeout=5)
        finally:
            if dummy_proc.poll() is None:
                dummy_proc.kill()
                dummy_proc.wait(timeout=5)

    def test_bash_refusal_names_scoped_process_control(self):
        """Broad pkill and kill commands in bash() report kill_process in actionable error."""
        for cmd in ["pkill -f cargo", "killall node", "kill 12345"]:
            with self.subTest(cmd=cmd):
                with self.assertRaises(ValueError) as ctx:
                    _check_command_guards(cmd)
                msg = str(ctx.exception)
                self.assertIn("kill_process(pid)", msg)
                self.assertIn("rlm.adopt_process", msg)
                self.assertIn("h.kill()", msg)

    def test_descendant_process_is_owned_via_ancestry_walk(self):
        """A grandchild process spawned inside an owned process is recognized as owned."""
        proc = subprocess.Popen(
            ["sh", "-c", "sleep 60 & echo $!; wait"],
            stdout=subprocess.PIPE,
            text=True,
        )
        try:
            register_owned_process(proc.pid, "sh -c 'sleep 60 & echo $!; wait'")
            assert proc.stdout is not None
            first_line = proc.stdout.readline().strip()
            self.assertTrue(first_line, "No output received from grandchild spawn")
            grandchild_pid = int(first_line)

            self.assertTrue(is_pid_owned(grandchild_pid))
            signalled = kill_process(grandchild_pid, sig=signal.SIGTERM)
            self.assertTrue(signalled)
        finally:
            proc.kill()
            proc.wait(timeout=5)

    def test_pid_1_and_unrelated_process_refused(self):
        """PID 1 and unrelated system processes are strictly refused."""
        self.assertFalse(is_pid_owned(1))
        with self.assertRaises(ValueError) as ctx:
            kill_process(1)
        self.assertIn("not owned by this session", str(ctx.exception))

        fake_pid = 888888
        self.assertFalse(is_pid_owned(fake_pid))
        with self.assertRaises(ValueError) as ctx:
            kill_process(fake_pid)
        self.assertIn("not owned by this session", str(ctx.exception))

    def test_adopt_process_records_and_enables_signaling(self):
        """rlm.adopt_process(pid, *, reason=...) enables signaling for external test subject."""
        dummy_proc = subprocess.Popen(["sleep", "60"])
        try:
            pid = dummy_proc.pid
            self.assertFalse(is_pid_owned(pid))

            lease = rlm.adopt_process(pid, reason="private daemon test subject in tmux")
            self.assertEqual(lease.pid, pid)
            self.assertEqual(lease.reason, "private daemon test subject in tmux")
            self.assertTrue(is_pid_owned(pid))

            records = list_owned_processes()
            adopted_record = next((r for r in records if r["pid"] == pid), None)
            self.assertIsNotNone(adopted_record)
            self.assertEqual(adopted_record["reason"], "private daemon test subject in tmux")

            signalled = kill_process(pid, sig=signal.SIGTERM)
            self.assertTrue(signalled)
            dummy_proc.wait(timeout=5)
            self.assertIsNotNone(dummy_proc.poll())
        finally:
            if dummy_proc.poll() is None:
                dummy_proc.kill()
                dummy_proc.wait(timeout=5)

    def test_adopt_process_without_reason_is_refused(self):
        """Adoption without a non-empty reason is refused."""
        dummy_proc = subprocess.Popen(["sleep", "10"])
        try:
            with self.assertRaises(ValueError) as ctx:
                rlm.adopt_process(dummy_proc.pid, reason="")
            self.assertIn("non-empty reason", str(ctx.exception))

            with self.assertRaises(ValueError) as ctx:
                rlm.adopt_process(dummy_proc.pid, reason="   ")
            self.assertIn("non-empty reason", str(ctx.exception))
        finally:
            dummy_proc.kill()
            dummy_proc.wait(timeout=5)

    def test_adopt_process_invalid_pid_is_refused(self):
        """Adoption of PID 1 or invalid PID is refused."""
        with self.assertRaises(ValueError) as ctx:
            rlm.adopt_process(1, reason="cannot adopt init")
        self.assertIn("valid non-system PID", str(ctx.exception))

        with self.assertRaises(ValueError) as ctx:
            rlm.adopt_process(-5, reason="invalid negative pid")
        self.assertIn("valid non-system PID", str(ctx.exception))

    def test_adoption_is_journaled(self):
        """Adopting a process appends a record to the session orphan journal."""
        dummy_proc = subprocess.Popen(["sleep", "10"])
        import tempfile, json
        with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tf:
            journal_path = tf.name
        try:
            with mock.patch.dict(os.environ, {"PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL": journal_path}):
                rlm.adopt_process(dummy_proc.pid, reason="auditable daemon adoption")

            with open(journal_path) as f:
                lines = f.readlines()
            records = [json.loads(line) for line in lines if line.strip()]
            self.assertTrue(records)
            rec = records[-1]
            self.assertEqual(rec["pid"], dummy_proc.pid)
            self.assertEqual(rec["action"], "adopted")
            self.assertEqual(rec["reason"], "auditable daemon adoption")
        finally:
            dummy_proc.kill()
            dummy_proc.wait(timeout=5)
            if os.path.exists(journal_path):
                os.remove(journal_path)

    def test_refusal_messages_name_adopt_process(self):
        """Refusal messages across kill_process, os.kill, and bash explicitly name adopt_process."""
        with self.assertRaises(ValueError) as ctx:
            kill_process(999999)
        self.assertIn("rlm.adopt_process", str(ctx.exception))

        code = compile("import os; os.kill(999999, 15)", "<cell-foreign-kill-adopt>", "exec")
        with self.assertRaises(PermissionError) as ctx:
            eval(code, {"os": os, "__name__": "__main__"})
        self.assertIn("rlm.adopt_process", str(ctx.exception))

        with self.assertRaises(ValueError) as ctx:
            _check_command_guards("kill 12345")
        self.assertIn("rlm.adopt_process", str(ctx.exception))

    # -----------------------------------------------------------------------
    # 2. Power-Control Command Guards
    # -----------------------------------------------------------------------

    def test_power_control_commands_refused_without_explicit_intent(self):
        """Standalone power commands and systemctl power actions are refused."""
        refused_cmds = [
            "reboot",
            "poweroff",
            "shutdown -h now",
            "halt",
            "init 0",
            "init 6",
            "systemctl reboot",
            "systemctl poweroff",
            "systemctl halt",
            "systemctl suspend",
            "systemctl hibernate",
            "sudo reboot",
            "sudo systemctl poweroff",
        ]
        for cmd in refused_cmds:
            with self.subTest(cmd=cmd):
                with self.assertRaises(ValueError) as ctx:
                    _check_command_guards(cmd)
                msg = str(ctx.exception)
                self.assertIn("power-control commands", msg)
                self.assertIn("require explicit user intent", msg)
                self.assertIn("RLM_ALLOW_POWER_CONTROL=1", msg)

    def test_power_control_permitted_with_explicit_intent(self):
        """Power commands succeed validation when explicit intent is declared."""
        for cmd in [
            "RLM_ALLOW_POWER_CONTROL=1 reboot",
            "RLM_ALLOW_POWER_CONTROL=1 systemctl poweroff",
            "PRIME_ALLOW_POWER_CONTROL=1 shutdown -h now",
        ]:
            with self.subTest(cmd=cmd):
                self.assertIsNone(_check_command_guards(cmd))

    def test_protected_services_cannot_be_disabled_without_explicit_intent(self):
        """systemctl disable/stop/mask on protected desktop services is refused."""
        set_host_profile(
            name="desktop",
            services_to_survive=("display-manager", "sway", "hyprland", "pipewire"),
        )
        for cmd in [
            "systemctl disable display-manager",
            "systemctl stop sway",
            "systemctl mask pipewire.service",
        ]:
            with self.subTest(cmd=cmd):
                with self.assertRaises(ValueError) as ctx:
                    _check_command_guards(cmd)
                msg = str(ctx.exception)
                self.assertIn("protected desktop service", msg)
                self.assertIn("RLM_ALLOW_POWER_CONTROL=1", msg)

    # -----------------------------------------------------------------------
    # 3. Desktop VRAM Safety & Preflight
    # -----------------------------------------------------------------------

    def test_vram_preflight_refuses_overbudget_naming_budget(self):
        """preflight_vram() rejects requested VRAM exceeding budget naming budget and headroom."""
        set_host_profile(
            name="test-gpu",
            total_vram_gb=24.0,
            desktop_headroom_gb=4.0,  # Available budget = 20.0 GB
        )
        with self.assertRaises(ValueError) as ctx:
            preflight_vram(24.0)

        err = str(ctx.exception)
        self.assertIn("VRAM preflight rejected", err)
        self.assertIn("requested 24.0 GB exceeds host VRAM budget", err)
        self.assertIn("20.0 GB available", err)
        self.assertIn("4.0 GB desktop headroom reserved", err)
        self.assertIn("24.0 GB total", err)
        self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", err)

    def test_vram_preflight_accepts_within_budget(self):
        """preflight_vram() succeeds when within budget."""
        set_host_profile(
            name="test-gpu",
            total_vram_gb=24.0,
            desktop_headroom_gb=4.0,  # Available budget = 20.0 GB
        )
        res = preflight_vram(16.0)
        self.assertEqual(res["status"], "accepted")
        self.assertEqual(res["available_budget_gb"], 20.0)

    def test_bash_guard_refuses_overbudget_gpu_command_naming_budget(self):
        """bash() rejects over-budget GPU commands naming the budget."""
        set_host_profile(
            name="test-gpu",
            total_vram_gb=24.0,
            desktop_headroom_gb=4.0,  # Budget = 20.0 GB
        )
        # 1. 70B model requires estimated 70.0 GB > 20.0 GB budget
        with self.assertRaises(ValueError) as ctx:
            _check_command_guards("vllm serve meta-llama/Llama-3-70B")
        err = str(ctx.exception)
        self.assertIn("bash() rejects `vllm`", err)
        self.assertIn("estimated VRAM requirement (70.0 GB) exceeds host VRAM budget (20.0 GB", err)
        self.assertIn("4.0 GB desktop headroom reserved", err)
        self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", err)

        # 2. --gpu-memory-utilization 0.95 on 24 GB = 22.8 GB > 20.0 GB budget
        with self.assertRaises(ValueError) as ctx:
            _check_command_guards("python -m vllm --model meta-llama/Llama-3-8B --gpu-memory-utilization 0.95")
        err = str(ctx.exception)
        self.assertIn("exceeds host VRAM budget (20.0 GB", err)

    def test_bash_guard_allows_overbudget_gpu_with_explicit_opt_in(self):
        """Over-budget GPU command is allowed when explicit opt-in is provided."""
        set_host_profile(
            name="test-gpu",
            total_vram_gb=24.0,
            desktop_headroom_gb=4.0,
        )
        cmd = "RLM_ALLOW_OVERBUDGET_VRAM=1 vllm serve meta-llama/Llama-3-70B"
        self.assertIsNone(_check_command_guards(cmd))

    def test_bash_guard_allows_within_budget_gpu_command(self):
        """GPU command within budget is permitted."""
        set_host_profile(
            name="test-gpu",
            total_vram_gb=24.0,
            desktop_headroom_gb=4.0,
        )
        # 8B model estimated 6.4 GB <= 20.0 GB
        cmd = "vllm serve meta-llama/Llama-3-8B --gpu-memory-utilization 0.7"
        self.assertIsNone(_check_command_guards(cmd))

    # -----------------------------------------------------------------------
    # 4. Hardened Process Ownership & Identity (Defects 1, 2, 3)
    # -----------------------------------------------------------------------

    def test_dead_process_lease_expires_and_refuses_signal(self):
        """A lease for a reaped/dead process expires and cannot authorize signals."""
        proc = subprocess.Popen(["sleep", "0.05"])
        register_owned_process(proc.pid, "sleep 0.05")
        proc.wait(timeout=5)
        self.assertFalse(is_pid_owned(proc.pid))
        with self.assertRaises(ValueError) as ctx:
            kill_process(proc.pid)
        self.assertIn("not owned by this session", str(ctx.exception))

    def test_pid_reuse_identity_mismatch_refuses_signal(self):
        """PID reuse with mismatched process identity expires lease and refuses signal."""
        proc = subprocess.Popen(["sleep", "30"])
        try:
            lease = register_owned_process(proc.pid, "sleep 30")
            lease.identity = ProcessIdentity(pid=proc.pid, start_time="mismatched_start_123", comm="stranger")
            self.assertFalse(is_pid_owned(proc.pid))
            with self.assertRaises(ValueError) as ctx:
                kill_process(proc.pid)
            self.assertIn("not owned by this session", str(ctx.exception))
        finally:
            proc.kill()
            proc.wait(timeout=5)

    def test_ancestry_walk_does_not_cache_visited_strangers(self):
        """Ancestry walk does not permanently enroll visited intermediate PIDs into _OWNED_LEASES."""
        proc = subprocess.Popen(["sh", "-c", "sleep 60 & echo $!; wait"], stdout=subprocess.PIPE, text=True)
        try:
            register_owned_process(proc.pid, "parent")
            assert proc.stdout is not None
            grandchild_pid = int(proc.stdout.readline().strip())
            self.assertTrue(is_pid_owned(grandchild_pid))
            self.assertNotIn(grandchild_pid, _OWNED_LEASES)
        finally:
            proc.kill()
            proc.wait(timeout=5)

    def test_adopted_process_signals_pid_alone_not_entire_group(self):
        """An adopted process signals the PID alone without os.killpg when group is not fully owned."""
        proc = subprocess.Popen(["sleep", "30"])
        try:
            adopt_process(proc.pid, reason="test adoption")
            with mock.patch("os.kill") as mock_kill, mock.patch("os.killpg") as mock_killpg:
                kill_process(proc.pid, sig=signal.SIGTERM)
                mock_kill.assert_called_once_with(proc.pid, signal.SIGTERM)
                mock_killpg.assert_not_called()
        finally:
            proc.kill()
            proc.wait(timeout=5)

    def test_adopt_lifeline_refused_without_opt_in(self):
        """Adopting kernel self, ppid, or owner PID is refused without self_lifeline_ok=True."""
        with self.assertRaises(ValueError) as ctx:
            adopt_process(os.getpid(), reason="attempt self adopt")
        self.assertIn("critical session lifeline", str(ctx.exception))
        self.assertIn("self_lifeline_ok=True", str(ctx.exception))

        with self.assertRaises(ValueError) as ctx:
            adopt_process(os.getppid(), reason="attempt ppid adopt")
        self.assertIn("critical session lifeline", str(ctx.exception))

    def test_adopt_lifeline_permitted_with_opt_in_and_audited(self):
        """Adopting lifeline with self_lifeline_ok=True succeeds and emits audit entry."""
        import tempfile, json
        with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tf:
            journal_path = tf.name
        try:
            with mock.patch.dict(os.environ, {"PRIME_AGENT_INTERNAL_ORPHAN_PROCESS_JOURNAL": journal_path}):
                lease = adopt_process(os.getppid(), reason="explicit supervisor restart", self_lifeline_ok=True)
                self.assertEqual(lease.pid, os.getppid())
            with open(journal_path) as f:
                lines = [json.loads(line) for line in f if line.strip()]
            self.assertTrue(any(r.get("selfLifelineOptIn") is True for r in lines))
            self.assertTrue(any(r.get("action") == "self_lifeline_override" for r in lines))
        finally:
            if os.path.exists(journal_path):
                os.remove(journal_path)

    # -----------------------------------------------------------------------
    # 5. systemctl Restart Verbs Guard (Defect 4)
    # -----------------------------------------------------------------------

    def test_systemctl_restart_verbs_refused_on_desktop_services(self):
        """systemctl restart, try-restart, reload-or-restart, isolate are guarded for protected desktop services."""
        set_host_profile(services_to_survive=("display-manager", "gdm", "logind"))
        for verb in ["restart", "try-restart", "reload-or-restart", "isolate"]:
            for target in ["display-manager", "gdm.service", "logind"]:
                with self.subTest(verb=verb, target=target):
                    with self.assertRaises(ValueError) as ctx:
                        _check_command_guards(f"systemctl {verb} {target}")
                    msg = str(ctx.exception)
                    self.assertIn("protected desktop service", msg)
                    self.assertIn("RLM_ALLOW_POWER_CONTROL=1", msg)

    def test_systemctl_restart_permitted_for_unrelated_services(self):
        """systemctl restart on normal/unrelated services is permitted."""
        set_host_profile(services_to_survive=("display-manager", "gdm"))
        for cmd in ["systemctl restart nginx", "systemctl try-restart my-worker.service"]:
            with self.subTest(cmd=cmd):
                self.assertIsNone(_check_command_guards(cmd))

    # -----------------------------------------------------------------------
    # 6. Unconfigured VRAM Preflight Fail-Closed (Defect 5)
    # -----------------------------------------------------------------------

    def test_unconfigured_vram_fails_closed_for_large_model(self):
        """Unconfigured host profile fails closed for large model requests exceeding default headroom/budget."""
        reset_host_profile()
        with self.assertRaises(ValueError) as ctx:
            _check_command_guards("vllm serve meta-llama/Llama-3-70B")
        msg = str(ctx.exception)
        self.assertIn("exceeds host VRAM budget", msg)
        self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", msg)
        self.assertIn("set_host_profile", msg)

    def test_unconfigured_vram_permits_small_model(self):
        """Unconfigured host profile permits small models within conservative headroom."""
        reset_host_profile()
        self.assertIsNone(_check_command_guards("vllm serve meta-llama/Llama-3-8B --gpu-memory-utilization 0.5"))

    def test_unconfigured_vram_permits_large_model_with_opt_in(self):
        """Unconfigured host profile permits large models with explicit opt-in."""
        reset_host_profile()
        self.assertIsNone(_check_command_guards("RLM_ALLOW_OVERBUDGET_VRAM=1 vllm serve meta-llama/Llama-3-70B"))

    def test_unconfigured_vram_headless_allows_normal_work(self):
        """On a genuinely headless host, no desktop headroom is reserved and normal server work is permitted."""
        reset_host_profile()
        with mock.patch("rlm.host_safety.is_headless", return_value=True):
            # Headless unconfigured budget is 80.0 GB, so 70B model passes
            self.assertIsNone(_check_command_guards("vllm serve meta-llama/Llama-3-70B"))
            # Runaway request (> 80 GB) fails closed even on headless
            with self.assertRaises(ValueError) as ctx:
                _check_command_guards("vllm serve meta-llama/Llama-3-405B")
            self.assertIn("RLM_ALLOW_OVERBUDGET_VRAM=1", str(ctx.exception))


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