from unittest.mock import MagicMock, AsyncMock, patch

import pytest

from costrouter import BaseRouter, CostRouter, CostTier, ModelConfig, RoutedResponse
from costrouter.exceptions import AllModelsFailedError, ScreenerError


def _make_llm_response(content: str = "Hello!") -> MagicMock:
    msg = MagicMock()
    msg.content = content
    choice = MagicMock()
    choice.message = msg
    resp = MagicMock()
    resp.choices = [choice]
    return resp


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


class TestCostRouterInit:
    def test_dict_config(self):
        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
            screener="gpt-4o-mini",
            strategy="heuristic",
        )
        assert router._models[0].name == "gpt-3.5-turbo"
        assert router._models[0].tier == CostTier.SIMPLE

    def test_list_config(self):
        router = CostRouter(
            models=[
                ModelConfig(name="gpt-3.5-turbo", tier=CostTier.SIMPLE),
                ModelConfig(name="gpt-4o", tier=CostTier.COMPLEX, api_key="sk-test"),
            ],
            strategy="heuristic",
        )
        assert len(router._models) == 2

    def test_ordered_list_2_models(self):
        router = CostRouter(
            models=["gpt-4o-mini", "gpt-4o"],
            strategy="heuristic",
        )
        assert router._models[0].name == "gpt-4o-mini"
        assert router._models[0].tier == CostTier.SIMPLE
        assert router._models[1].name == "gpt-4o"
        assert router._models[1].tier == CostTier.COMPLEX

    def test_ordered_list_3_models(self):
        router = CostRouter(
            models=["gpt-4o-mini", "gpt-4o", "claude-opus-4-6"],
            strategy="heuristic",
        )
        assert router._models[0].tier == CostTier.SIMPLE
        assert router._models[1].tier == CostTier.MODERATE
        assert router._models[2].tier == CostTier.COMPLEX

    def test_ordered_list_4_models(self):
        router = CostRouter(
            models=["a", "b", "c", "d"],
            strategy="heuristic",
        )
        assert router._models[0].tier == CostTier.SIMPLE
        assert router._models[1].tier == CostTier.MODERATE
        assert router._models[2].tier == CostTier.COMPLEX
        assert router._models[3].tier == CostTier.EXPERT

    def test_ordered_list_1_model(self):
        router = CostRouter(
            models=["gpt-4o"],
            strategy="heuristic",
        )
        assert router._models[0].tier == CostTier.MODERATE

    def test_screener_defaults_to_cheapest(self):
        router = CostRouter(
            models=["gpt-4o-mini", "gpt-4o"],
            strategy="heuristic",
        )
        # No screener specified — should default to first (cheapest) model
        assert router._models[0].name == "gpt-4o-mini"

    def test_empty_models_raises(self):
        with pytest.raises(ValueError, match="At least one model"):
            CostRouter(models={}, strategy="heuristic")

    def test_strategies(self):
        for strategy in ("heuristic", "llm", "cascade"):
            router = CostRouter(
                models={"simple": "gpt-3.5-turbo"},
                strategy=strategy,
            )
            assert router._strategy == strategy


class TestCostRouterCompletion:
    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_heuristic_simple_greeting(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response("Hi there!")

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
            strategy="heuristic",
        )
        result = router.completion(
            messages=[{"role": "user", "content": "Hello!"}]
        )

        assert isinstance(result, RoutedResponse)
        assert result.routing.tier == CostTier.SIMPLE
        assert result.routing.model == "gpt-3.5-turbo"
        assert result.routing.router_type == "heuristic"
        assert result.cost == 0.001
        assert result.routing.screener_cost == 0.0

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.01)
    @patch("costrouter.client.litellm")
    def test_heuristic_code_routes_to_complex(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response("def fib()...")

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
            strategy="heuristic",
        )
        result = router.completion(
            messages=[{"role": "user", "content": "Write a function to sort a list"}]
        )

        assert result.routing.tier == CostTier.COMPLEX
        assert result.routing.model == "gpt-4o"

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_tier_falls_back_to_nearest_cheaper(self, mock_litellm, mock_calc, mock_est):
        """If no model for MODERATE, should prefer cheaper (SIMPLE) over expensive (EXPERT)."""
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "expert": "gpt-4o"},
            strategy="heuristic",
        )
        # Default heuristic for ambiguous text -> MODERATE
        result = router.completion(
            messages=[{"role": "user", "content": "Tell me something interesting about dolphins and their social behaviors"}]
        )
        # MODERATE has no model, should fall back to SIMPLE (cheaper) first
        assert result.routing.model == "gpt-3.5-turbo"

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_fallback_on_model_failure(self, mock_litellm, mock_calc, mock_est):
        """If primary model fails, should try next tier up."""
        mock_litellm.completion.side_effect = [
            Exception("Rate limited"),
            _make_llm_response("Fallback response"),
        ]

        router = CostRouter(
            models={
                "simple": "gpt-3.5-turbo",
                "complex": "gpt-4o",
            },
            strategy="heuristic",
            fallback=True,
        )
        result = router.completion(
            messages=[{"role": "user", "content": "Hello!"}]
        )
        # Should have fallen back from simple (gpt-3.5-turbo) to complex (gpt-4o)
        assert result.routing.model == "gpt-4o"

    @patch("costrouter.client.litellm")
    def test_all_models_fail_raises(self, mock_litellm):
        mock_litellm.completion.side_effect = Exception("All down")

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo"},
            strategy="heuristic",
            fallback=True,
        )
        with pytest.raises(AllModelsFailedError):
            router.completion(
                messages=[{"role": "user", "content": "Hello!"}]
            )

    @patch("costrouter.client.litellm")
    def test_no_fallback_raises_immediately(self, mock_litellm):
        mock_litellm.completion.side_effect = Exception("Failed")

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
            strategy="heuristic",
            fallback=False,
        )
        with pytest.raises(AllModelsFailedError):
            router.completion(
                messages=[{"role": "user", "content": "Hello!"}]
            )
        # Should only have tried once (no fallback)
        assert mock_litellm.completion.call_count == 1

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_analytics_recorded(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo"},
            strategy="heuristic",
        )
        router.completion(messages=[{"role": "user", "content": "Hi"}])
        router.completion(messages=[{"role": "user", "content": "Hey"}])

        assert len(router.analytics.records) == 2

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_extra_kwargs_passed_through(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo"},
            strategy="heuristic",
        )
        router.completion(
            messages=[{"role": "user", "content": "Hi"}],
            temperature=0.5,
            max_tokens=100,
        )

        call_kwargs = mock_litellm.completion.call_args[1]
        assert call_kwargs["temperature"] == 0.5
        assert call_kwargs["max_tokens"] == 100

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_model_api_key_passed(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models=[
                ModelConfig(name="gpt-3.5-turbo", tier=CostTier.SIMPLE, api_key="sk-model-key"),
            ],
            strategy="heuristic",
        )
        router.completion(messages=[{"role": "user", "content": "Hi"}])

        call_kwargs = mock_litellm.completion.call_args[1]
        assert call_kwargs["api_key"] == "sk-model-key"


    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_screener_failure_uses_default_tier(self, mock_litellm, mock_calc, mock_est):
        """If screener fails, should fall back to default_tier instead of crashing."""
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
            strategy="llm",
            default_tier=CostTier.COMPLEX,
        )
        # Make the screener LLM fail
        with patch.object(router._router, "classify", side_effect=ScreenerError("down")):
            result = router.completion(
                messages=[{"role": "user", "content": "Hello"}]
            )
        assert result.routing.tier == CostTier.COMPLEX
        assert result.routing.confidence == 0.0

    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_extra_params_cannot_overwrite_model(self, mock_litellm, mock_calc, mock_est):
        """extra_params should not be able to overwrite the routed model name."""
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models=[
                ModelConfig(
                    name="gpt-3.5-turbo",
                    tier=CostTier.SIMPLE,
                    extra_params={"model": "should-be-ignored"},
                ),
            ],
            strategy="heuristic",
        )
        router.completion(messages=[{"role": "user", "content": "Hi"}])

        call_kwargs = mock_litellm.completion.call_args[1]
        assert call_kwargs["model"] == "gpt-3.5-turbo"


class TestCostRouterAsync:
    @pytest.mark.asyncio
    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    async def test_acompletion_heuristic(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.acompletion = AsyncMock(return_value=_make_llm_response())

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo"},
            strategy="heuristic",
        )
        result = await router.acompletion(
            messages=[{"role": "user", "content": "Hi"}]
        )

        assert isinstance(result, RoutedResponse)
        assert result.routing.tier == CostTier.SIMPLE

    @pytest.mark.asyncio
    @patch("costrouter.client.litellm")
    async def test_acompletion_fallback(self, mock_litellm):
        mock_litellm.acompletion = AsyncMock(
            side_effect=[Exception("Fail"), _make_llm_response()]
        )

        with patch("costrouter.client.calculate_cost", return_value=0.001), \
             patch("costrouter.client.estimate_max_cost", return_value=0.05):
            router = CostRouter(
                models={"simple": "gpt-3.5-turbo", "complex": "gpt-4o"},
                strategy="heuristic",
                fallback=True,
            )
            result = await router.acompletion(
                messages=[{"role": "user", "content": "Hi"}]
            )
            assert result.routing.model == "gpt-4o"


class TestCustomRouter:
    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_custom_router(self, mock_litellm, mock_calc, mock_est):
        """Users can pass their own BaseRouter subclass."""

        class AlwaysExpert(BaseRouter):
            def classify(self, messages):
                return CostTier.EXPERT, 1.0, 0.0

            async def aclassify(self, messages):
                return self.classify(messages)

        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "expert": "gpt-4o"},
            strategy=AlwaysExpert(),
        )
        result = router.completion(
            messages=[{"role": "user", "content": "Hi"}]
        )
        assert result.routing.tier == CostTier.EXPERT
        assert result.routing.model == "gpt-4o"
        assert result.routing.router_type == "custom"


class TestRoutedResponseShortcuts:
    @patch("costrouter.client.estimate_max_cost", return_value=0.05)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_content_shortcut(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response("Hello there!")

        router = CostRouter(models=["gpt-4o-mini", "gpt-4o"], strategy="heuristic")
        result = router.completion(messages=[{"role": "user", "content": "Hi"}])

        assert result.content == "Hello there!"
        assert result.model == "gpt-4o-mini"
        assert result.tier == CostTier.SIMPLE

    def test_content_shortcut_none_on_bad_response(self):
        from costrouter.models import RoutedResponse, RoutingResult

        result = RoutedResponse(
            response=None,
            routing=RoutingResult(
                tier=CostTier.SIMPLE, model="x", confidence=1.0,
                router_type="test", screener_cost=0.0,
            ),
            cost=0.0, total_cost=0.0, saved=0.0,
        )
        assert result.content is None


class TestCostRouterSavings:
    @patch("costrouter.client.estimate_max_cost", return_value=0.10)
    @patch("costrouter.client.calculate_cost", return_value=0.001)
    @patch("costrouter.client.litellm")
    def test_saved_calculation(self, mock_litellm, mock_calc, mock_est):
        mock_litellm.completion.return_value = _make_llm_response()

        router = CostRouter(
            models={"simple": "gpt-3.5-turbo", "expert": "gpt-4o"},
            strategy="heuristic",
        )
        result = router.completion(
            messages=[{"role": "user", "content": "Hi"}]
        )
        # saved = max_cost(0.10) - total_cost(0.001) = 0.099
        assert result.saved == pytest.approx(0.099, abs=0.001)
        assert result.total_cost == pytest.approx(0.001)
