from __future__ import annotations

import os
import tempfile
import unittest
from unittest import mock

from rlm.bash import bash
from rlm import sleep_guard


class SleepGuardTests(unittest.TestCase):
    def setUp(self) -> None:
        sleep_guard.install_sleep_guard()
        os.environ.pop("RLM_CELL_SLEEP_LIMIT", None)

    def tearDown(self) -> None:
        sleep_guard.install_sleep_guard()
        os.environ.pop("RLM_CELL_SLEEP_LIMIT", None)

    def test_long_literal_sleep_refused_with_options(self) -> None:
        code = compile("import time; time.sleep(15)", "<cell-test-long>", "exec")
        with self.assertRaises(ValueError) as ctx:
            eval(code, {"__name__": "__main__"})
        msg = str(ctx.exception)
        self.assertIn("r = await h", msg)
        self.assertIn("sentinel `bash()`", msg)

    def test_short_sleep_allowed(self) -> None:
        code = compile("import time; time.sleep(0.001)", "<cell-test-short>", "exec")
        eval(code, {"__name__": "__main__"})

    def test_library_frame_sleep_allowed(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "helper_lib.py")
            with open(mod_file, "w") as f:
                f.write("import time\ndef run_sleep():\n    time.sleep(15)\n")
            import sys
            sys.path.insert(0, tmpdir)
            try:
                import helper_lib
                with mock.patch.object(sleep_guard, "_ORIG_TIME_SLEEP") as mock_orig:
                    code = compile("helper_lib.run_sleep()", "<cell-test-lib>", "exec")
                    eval(code, {"helper_lib": helper_lib, "__name__": "__main__"})
                    mock_orig.assert_called_once_with(15)
            finally:
                sys.path.remove(tmpdir)
                sys.modules.pop("helper_lib", None)

    def test_escape_hatch_allows_sleep(self) -> None:
        os.environ["RLM_CELL_SLEEP_LIMIT"] = "0"
        with mock.patch.object(sleep_guard, "_ORIG_TIME_SLEEP") as mock_orig:
            code = compile("import time; time.sleep(15)", "<cell-test-escape>", "exec")
            eval(code, {"__name__": "__main__"})
            mock_orig.assert_called_once_with(15)

    def test_cell_awaiting_handle_does_not_fire_guard(self) -> None:
        import asyncio
        async def run_cmd():
            h = bash("echo test")
            return await h
        res = asyncio.run(run_cmd())
        self.assertEqual(res.exit_code, 0)


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