from __future__ import annotations

import os
import subprocess
import unittest

from rlm.subprocess_guard import install_subprocess_guards


class SubprocessGuardTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        install_subprocess_guards()

    def test_cell_caller_rejects_search_commands(self):
        # Simulate call from user cell by executing in a frame named __main__
        code = compile("import subprocess; subprocess.run(['rg', 'foo'])", "<cell-test-1>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("use search('pattern', path='...')", str(ctx.exception))

    def test_cell_caller_rejects_vcs_commands(self):
        code = compile("import subprocess; subprocess.run(['git', 'status'])", "<cell-test-2>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("git.status(...)", str(ctx.exception))

        code_grep = compile("import subprocess; subprocess.run(['git', 'grep', 'foo'])", "<cell-test-2b>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code_grep, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("git.grep(...)", str(ctx.exception))

    def test_cell_caller_rejects_bash_lc_with_rg(self):
        code = compile("import subprocess; subprocess.run(['bash', '-lc', 'rg -l line-height'])", "<cell-test-3>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("bash() rejects `rg`", str(ctx.exception))

    def test_cell_caller_rejects_find_commands(self):
        code = compile("import subprocess; subprocess.run(['find', '.'])", "<cell-test-find>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("search.files", str(ctx.exception))

    def test_cell_caller_rejects_ls_commands(self):
        code = compile("import subprocess; subprocess.run(['ls', '-la'])", "<cell-test-ls>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("search.ls", str(ctx.exception))

    def test_cell_caller_rejects_general_shell_commands(self):
        code = compile("import subprocess; subprocess.run(['whoami'])", "<cell-test-4>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("use await bash(...) instead", str(ctx.exception))

    def test_popen_remains_subscriptable_class(self):
        # mcp.os.win32.utilities evaluates `subprocess.Popen[bytes]` in a
        # signature at import time; the guard must keep Popen class-like.
        self.assertTrue(issubclass(subprocess.Popen, subprocess.Popen))
        subprocess.Popen[bytes]
        import mcp.os.win32.utilities  # noqa: F401

    def test_cell_caller_rejects_nix_commands(self):
        code_shell = compile("import subprocess; subprocess.run(['nix-shell', '-p', 'cargo'])", "<cell-test-nix-1>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code_shell, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("use `await nix.develop(...)`", str(ctx.exception))

        code_dev = compile("import subprocess; subprocess.run(['nix', 'develop'])", "<cell-test-nix-2>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code_dev, {"subprocess": subprocess, "__name__": "__main__"})
        self.assertIn("nix.develop", str(ctx.exception))

    def test_cell_caller_rejects_private_helper_run(self):
        # Even if a helper in an exempt module calls subprocess, if it is a private helper
        # called directly from a cell, it must not bypass guards
        def _fake_skill_run(cmd):
            return subprocess.run(cmd)

        _fake_skill_run.__module__ = "jj"
        code = compile("_fake_skill_run(['git', 'status'])", "<cell-test-priv>", "exec")
        with self.assertRaises(RuntimeError) as ctx:
            eval(code, {"_fake_skill_run": _fake_skill_run, "__name__": "__main__"})
        self.assertIn("git.status(...)", str(ctx.exception))

    def test_exempt_caller_allows_subprocess(self):
        # Calls from unittest (an exempt module) must run normally
        res = subprocess.run(["echo", "hello"], capture_output=True, text=True)
        self.assertEqual(res.stdout.strip(), "hello")


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