from unittest.mock import MagicMock, AsyncMock, patch

import pytest

from costrouter.exceptions import ScreenerError
from costrouter.models import CostTier
from costrouter.router.llm import LLMRouter


def _make_response(content: str) -> MagicMock:
    msg = MagicMock()
    msg.content = content
    choice = MagicMock()
    choice.message = msg
    resp = MagicMock()
    resp.choices = [choice]
    return resp


def _msgs(text: str) -> list[dict]:
    return [{"role": "user", "content": text}]


class TestLLMRouter:
    def setup_method(self):
        self.router = LLMRouter(screener_model="gpt-4o-mini")

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_classify_simple(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("SIMPLE")
        tier, conf, cost = self.router.classify(_msgs("Hi"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.85
        assert cost == 0.0001

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_classify_complex(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("COMPLEX")
        tier, conf, _ = self.router.classify(_msgs("Write a function"))
        assert tier == CostTier.COMPLEX

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_classify_expert(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("EXPERT")
        tier, conf, _ = self.router.classify(_msgs("Prove the Riemann hypothesis"))
        assert tier == CostTier.EXPERT

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_classify_moderate(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("MODERATE")
        tier, conf, _ = self.router.classify(_msgs("Summarize this article"))
        assert tier == CostTier.MODERATE

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_unparseable_defaults_to_moderate(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("I think this is complex because...")
        tier, conf, _ = self.router.classify(_msgs("something"))
        assert tier == CostTier.MODERATE
        assert conf == 0.5

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_whitespace_stripped(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("  SIMPLE  \n")
        tier, conf, _ = self.router.classify(_msgs("Hi"))
        assert tier == CostTier.SIMPLE

    @patch("costrouter.router.llm.litellm")
    def test_api_error_raises_screener_error(self, mock_litellm):
        mock_litellm.completion.side_effect = Exception("API down")
        with pytest.raises(ScreenerError, match="LLM screening failed"):
            self.router.classify(_msgs("test"))

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_sends_full_conversation(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("MODERATE")
        messages = [
            {"role": "user", "content": "First message"},
            {"role": "assistant", "content": "Response"},
            {"role": "user", "content": "Follow up"},
        ]
        self.router.classify(messages)

        call_args = mock_litellm.completion.call_args
        sent_messages = call_args[1]["messages"]
        # System prompt + 3 conversation messages
        assert len(sent_messages) == 4
        assert sent_messages[0]["role"] == "system"

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_uses_api_key(self, mock_litellm, mock_cost):
        router = LLMRouter(screener_model="gpt-4o-mini", api_key="sk-test")
        mock_litellm.completion.return_value = _make_response("SIMPLE")
        router.classify(_msgs("Hi"))

        call_args = mock_litellm.completion.call_args
        assert call_args[1]["api_key"] == "sk-test"

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    def test_max_tokens_and_temperature(self, mock_litellm, mock_cost):
        mock_litellm.completion.return_value = _make_response("SIMPLE")
        self.router.classify(_msgs("Hi"))

        call_args = mock_litellm.completion.call_args
        assert call_args[1]["max_tokens"] == 10
        assert call_args[1]["temperature"] == 0


class TestLLMRouterAsync:
    def setup_method(self):
        self.router = LLMRouter(screener_model="gpt-4o-mini")

    @pytest.mark.asyncio
    @patch("costrouter.router.llm.calculate_cost", return_value=0.0001)
    @patch("costrouter.router.llm.litellm")
    async def test_aclassify_simple(self, mock_litellm, mock_cost):
        mock_litellm.acompletion = AsyncMock(return_value=_make_response("SIMPLE"))
        tier, conf, cost = await self.router.aclassify(_msgs("Hi"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.85

    @pytest.mark.asyncio
    @patch("costrouter.router.llm.litellm")
    async def test_aclassify_error(self, mock_litellm):
        mock_litellm.acompletion = AsyncMock(side_effect=Exception("API down"))
        with pytest.raises(ScreenerError):
            await self.router.aclassify(_msgs("test"))
