import os
import tempfile
import unittest

from rcl_rlm.exceptions import SandboxError
from rcl_rlm.message_database import MessageDatabase
from rcl_rlm.restricted_executor import RestrictedExecutor


class TestSandboxOutputBudget(unittest.TestCase):
    def _make_executor(self, max_output_chars: int) -> RestrictedExecutor:
        temp = tempfile.NamedTemporaryFile(delete=False)
        temp.close()
        self.addCleanup(lambda: os.unlink(temp.name))
        db = MessageDatabase(temp.name, quiet=True)
        return RestrictedExecutor(db, max_output_chars=max_output_chars)

    def test_output_under_limit(self) -> None:
        executor = self._make_executor(5)
        output = executor.execute('print("1234", end="")')
        self.assertEqual(output, "1234")

    def test_output_at_limit(self) -> None:
        executor = self._make_executor(4)
        output = executor.execute('print("1234", end="")')
        self.assertEqual(output, "1234")

    def test_output_over_limit(self) -> None:
        executor = self._make_executor(4)
        with self.assertRaises(SandboxError) as ctx:
            executor.execute('print("12345", end="")')
        self.assertIn("Summarize", str(ctx.exception))
        self.assertIn("attempted", str(ctx.exception))


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