import unittest
from unittest.mock import Mock

from rcl_rlm.execution_context import create_execution_context
from rcl_rlm.exceptions import BatchLLMError, SandboxError
from rcl_rlm.llm_client import LLMClient
from rcl_rlm.llm_orchestrator import LlmOrchestrator
from rcl_rlm.modelrelay import BatchError, BatchResponse, BatchResult, BatchUsage


class TestBatchLLMQuery(unittest.TestCase):
    def _make_orchestrator(self, max_calls: int = 10, max_depth: int = 1):
        context = create_execution_context(
            max_calls=max_calls,
            max_depth=max_depth,
            on_progress=lambda _status: None,
        )
        client = Mock(spec=LLMClient)
        orchestrator = LlmOrchestrator(context, client=client, subcall_model="test-model")
        return context, orchestrator, client

    def test_batch_llm_query_updates_calls_and_tokens(self):
        context, llm, client = self._make_orchestrator(max_calls=5)
        client.batch_responses_typed.return_value = BatchResponse(
            id="batch_1",
            results=[
                BatchResult(
                    id="subcall_0",
                    status="success",
                    response={
                        "output": [
                            {
                                "type": "message",
                                "role": "assistant",
                                "content": [{"type": "text", "text": "ok-1"}],
                            }
                        ]
                    },
                ),
                BatchResult(
                    id="subcall_1",
                    status="success",
                    response={
                        "output": [
                            {
                                "type": "message",
                                "role": "assistant",
                                "content": [{"type": "text", "text": "ok-2"}],
                            }
                        ]
                    },
                ),
            ],
            usage=BatchUsage(
                total_input_tokens=3,
                total_output_tokens=2,
                total_requests=2,
                successful_requests=2,
                failed_requests=0,
            ),
        )
        results = llm.batch_llm_query(["p1", "p2"])
        self.assertEqual(results, ["ok-1", "ok-2"])
        self.assertEqual(context.stats.calls_made, 2)
        self.assertEqual(context.stats.total_tokens, 5)
        self.assertEqual(context.stats.depth, 0)
        client.batch_responses_typed.assert_called_once()

    def test_batch_llm_query_respects_max_calls(self):
        context, llm, client = self._make_orchestrator(max_calls=1)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1", "p2"])
        self.assertEqual(context.stats.calls_made, 0)
        client.batch_responses_typed.assert_not_called()

    def test_batch_llm_query_aggregates_errors(self):
        _, llm, client = self._make_orchestrator(max_calls=5)
        client.batch_responses_typed.return_value = BatchResponse(
            id="batch_1",
            results=[
                BatchResult(
                    id="subcall_0",
                    status="success",
                    response={
                        "output": [
                            {
                                "type": "message",
                                "role": "assistant",
                                "content": [{"type": "text", "text": "ok-1"}],
                            }
                        ]
                    },
                ),
                BatchResult(
                    id="subcall_1",
                    status="error",
                    error=BatchError(status=400, message="bad request", code="INVALID"),
                ),
            ],
            usage=BatchUsage(
                total_input_tokens=1,
                total_output_tokens=1,
                total_requests=2,
                successful_requests=1,
                failed_requests=1,
            ),
        )
        with self.assertRaises(BatchLLMError) as ctx:
            llm.batch_llm_query(["p1", "p2"])
        self.assertIn("subcall_1", str(ctx.exception))
        self.assertEqual(len(ctx.exception.errors), 1)

    def test_batch_llm_query_rejects_invalid_max_concurrent(self):
        _, llm, client = self._make_orchestrator(max_calls=5)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], max_concurrent=0)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], max_concurrent=6)
        client.batch_responses_typed.assert_not_called()

    def test_batch_llm_query_rejects_invalid_fail_fast(self):
        _, llm, client = self._make_orchestrator(max_calls=5)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], fail_fast="yes")  # type: ignore[arg-type]
        client.batch_responses_typed.assert_not_called()

    def test_batch_llm_query_rejects_invalid_timeout(self):
        _, llm, client = self._make_orchestrator(max_calls=5)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], timeout_ms=0)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], timeout_ms=-5)
        with self.assertRaises(SandboxError):
            llm.batch_llm_query(["p1"], timeout_ms="10")  # type: ignore[arg-type]
        client.batch_responses_typed.assert_not_called()


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