from unittest.mock import MagicMock, AsyncMock, patch

import pytest

from costrouter.models import CostTier
from costrouter.router.cascade import CascadeRouter


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 TestCascadeRouter:
    def setup_method(self):
        self.router = CascadeRouter(
            screener_model="gpt-4o-mini",
            confidence_threshold=0.7,
        )

    def test_high_confidence_heuristic_skips_llm(self):
        """Greeting has 0.9 confidence -> should skip LLM."""
        tier, conf, cost = self.router.classify(_msgs("Hello, how are you today?"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.9
        assert cost == 0.0  # No LLM call

    def test_short_prompt_skips_llm(self):
        """Short prompt has 0.7 confidence -> should skip LLM."""
        tier, conf, cost = self.router.classify(_msgs("Hi"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.7
        assert cost == 0.0

    def test_code_keyword_skips_llm(self):
        """Code keyword has 0.7 confidence -> should skip LLM."""
        tier, conf, cost = self.router.classify(
            _msgs("Write a function to calculate fibonacci numbers")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.7
        assert cost == 0.0

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0002)
    @patch("costrouter.router.llm.litellm")
    def test_low_confidence_calls_llm(self, mock_litellm, mock_cost):
        """Default heuristic has 0.3 confidence -> should call LLM."""
        mock_litellm.completion.return_value = _make_response("COMPLEX")

        tier, conf, cost = self.router.classify(
            _msgs("Tell me something interesting about octopuses and their behavior")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.85
        assert cost == 0.0002
        mock_litellm.completion.assert_called_once()

    @patch("costrouter.router.llm.calculate_cost", return_value=0.0002)
    @patch("costrouter.router.llm.litellm")
    def test_analysis_below_threshold_calls_llm(self, mock_litellm, mock_cost):
        """Analysis has 0.5 confidence -> below 0.7 threshold -> calls LLM."""
        mock_litellm.completion.return_value = _make_response("MODERATE")

        tier, conf, cost = self.router.classify(
            _msgs("Compare the economic policies of these two candidates")
        )
        assert tier == CostTier.MODERATE
        assert conf == 0.85
        mock_litellm.completion.assert_called_once()

    def test_custom_threshold(self):
        """With threshold 0.95, even greetings should trigger LLM."""
        router = CascadeRouter(
            screener_model="gpt-4o-mini",
            confidence_threshold=0.95,
        )
        with patch("costrouter.router.llm.litellm") as mock_litellm, \
             patch("costrouter.router.llm.calculate_cost", return_value=0.0):
            mock_litellm.completion.return_value = _make_response("SIMPLE")
            tier, conf, _ = router.classify(_msgs("Hello there!"))
            mock_litellm.completion.assert_called_once()

    @pytest.mark.asyncio
    async def test_async_high_confidence_skips_llm(self):
        tier, conf, cost = await self.router.aclassify(_msgs("Hello there!"))
        assert tier == CostTier.SIMPLE
        assert cost == 0.0

    @pytest.mark.asyncio
    @patch("costrouter.router.llm.calculate_cost", return_value=0.0002)
    @patch("costrouter.router.llm.litellm")
    async def test_async_low_confidence_calls_llm(self, mock_litellm, mock_cost):
        mock_litellm.acompletion = AsyncMock(return_value=_make_response("EXPERT"))

        tier, conf, cost = await self.router.aclassify(
            _msgs("Tell me something interesting about octopuses and their behavior")
        )
        assert tier == CostTier.EXPERT
        assert conf == 0.85
