import os
import tempfile
import unittest
from typing import Any

from rcl_rlm.modelrelay import TokenUsage, BatchResponse
from rcl_rlm.rlm import run_rlm


class MockLLMClient:
    """Mock LLM client for testing that returns predetermined responses."""

    def __init__(self, responses: list[tuple[str, TokenUsage]]) -> None:
        self._responses = responses
        self._call_count = 0

    def responses_create(
        self,
        messages: list[dict[str, Any]],
        model: str,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        return_usage: bool = False,
    ) -> str | tuple[str, TokenUsage]:
        if self._call_count >= len(self._responses):
            raise RuntimeError(f"MockLLMClient: no more responses (called {self._call_count + 1} times)")
        response, usage = self._responses[self._call_count]
        self._call_count += 1
        if return_usage:
            return response, usage
        return response

    def batch_responses_typed(
        self,
        requests: list[dict[str, Any]],
        options: dict[str, Any] | None = None,
    ) -> BatchResponse:
        raise NotImplementedError("batch_responses_typed not needed for this test")

    def get_model_context_window(self, model_id: str) -> int | None:
        """Return default context window for tests."""
        return 32000

    def get_model_max_output_tokens(self, model_id: str) -> int | None:
        """Return default max output tokens for tests."""
        return 4096


class TestRLMOutputBudget(unittest.TestCase):
    def test_run_rlm_propagates_output_budget_error(self) -> None:
        temp = tempfile.NamedTemporaryFile(delete=False)
        temp.close()
        self.addCleanup(lambda: os.unlink(temp.name))

        response = "```python\nprint('x' * 10, end='')\n```"
        usage = TokenUsage(1, 1)

        mock_client = MockLLMClient(responses=[(response, usage), (response, usage)])

        result = run_rlm(
            question="test",
            db_path=temp.name,
            max_output_chars=5,
            max_iterations=1,
            llm_client=mock_client,
        )

        self.assertIsNotNone(result.error)
        self.assertEqual(result.error.get("type"), "SandboxError")
        self.assertIn("Summarize", result.error.get("message", ""))
        self.assertIn("attempted", result.error.get("message", ""))


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