from unittest.mock import MagicMock, patch

import pytest

from costrouter.pricing import calculate_cost, estimate_max_cost, get_model_input_cost_per_token


class TestCalculateCost:
    @patch("costrouter.pricing.litellm")
    def test_returns_cost(self, mock_litellm):
        mock_litellm.completion_cost.return_value = 0.0042
        response = MagicMock()
        cost = calculate_cost(response)
        assert cost == 0.0042
        mock_litellm.completion_cost.assert_called_once_with(completion_response=response)

    @patch("costrouter.pricing.litellm")
    def test_returns_zero_on_error(self, mock_litellm):
        mock_litellm.completion_cost.side_effect = Exception("No pricing data")
        cost = calculate_cost(MagicMock())
        assert cost == 0.0


class TestEstimateMaxCost:
    @patch("costrouter.pricing._get_model_output_cost_per_token", return_value=0.00006)
    @patch("costrouter.pricing.get_model_input_cost_per_token", return_value=0.00003)
    def test_estimates_cost_from_messages(self, mock_input, mock_output):
        messages = [{"role": "user", "content": "Hello world"}]
        cost = estimate_max_cost("gpt-4o", messages=messages)
        assert cost > 0

    @patch("costrouter.pricing._get_model_output_cost_per_token", return_value=0.00006)
    @patch("costrouter.pricing.get_model_input_cost_per_token", return_value=0.00003)
    def test_estimates_cost_from_response(self, mock_input, mock_output):
        response = MagicMock()
        response.usage.prompt_tokens = 100
        response.usage.completion_tokens = 50
        cost = estimate_max_cost("gpt-4o", response=response)
        # 100 * 0.00003 + 50 * 0.00006 = 0.003 + 0.003 = 0.006
        assert cost == pytest.approx(0.006)

    @patch("costrouter.pricing._get_model_output_cost_per_token", return_value=0.00006)
    @patch("costrouter.pricing.get_model_input_cost_per_token", return_value=0.00003)
    def test_response_preferred_over_messages(self, mock_input, mock_output):
        """When both response and messages are given, response tokens should be used."""
        response = MagicMock()
        response.usage.prompt_tokens = 10
        response.usage.completion_tokens = 5
        messages = [{"role": "user", "content": "x" * 10000}]
        cost = estimate_max_cost("gpt-4o", response=response, messages=messages)
        # Should use the small response token counts, not the huge message
        assert cost == pytest.approx(10 * 0.00003 + 5 * 0.00006)

    @patch("costrouter.pricing.get_model_input_cost_per_token", return_value=None)
    def test_returns_zero_when_no_pricing(self, mock_input):
        messages = [{"role": "user", "content": "Hello"}]
        cost = estimate_max_cost("unknown-model", messages=messages)
        assert cost == 0.0

    def test_returns_zero_when_no_args(self):
        with patch("costrouter.pricing.get_model_input_cost_per_token", return_value=0.00003):
            cost = estimate_max_cost("gpt-4o")
            assert cost == 0.0


class TestGetModelInputCostPerToken:
    @patch("costrouter.pricing.litellm")
    def test_returns_cost(self, mock_litellm):
        mock_litellm.get_model_info.return_value = {"input_cost_per_token": 0.00003}
        cost = get_model_input_cost_per_token("gpt-4o")
        assert cost == 0.00003

    @patch("costrouter.pricing.litellm")
    def test_returns_none_on_error(self, mock_litellm):
        mock_litellm.get_model_info.side_effect = Exception("Unknown model")
        cost = get_model_input_cost_per_token("nonexistent")
        assert cost is None

    @patch("costrouter.pricing.litellm")
    def test_returns_none_when_no_key(self, mock_litellm):
        mock_litellm.get_model_info.return_value = {}
        cost = get_model_input_cost_per_token("gpt-4o")
        assert cost is None
