"""
Integration tests for Perplexity cost calculation and transformation.

Tests the end-to-end functionality of Perplexity cost calculation 
including integration with the main LiteLLM cost calculator.
"""

import json
import math
import os

import pytest

# Add the project root to Python path

import litellm
from litellm import ModelResponse
from litellm.cost_calculator import completion_cost, cost_per_token
from litellm.llms.perplexity.chat.transformation import PerplexityChatConfig
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
from litellm.utils import get_model_info


class TestPerplexityIntegration:
    """Integration test suite for Perplexity functionality."""

    @pytest.fixture(autouse=True)
    def setup_model_cost_map(self):
        """Set up the model cost map for testing."""
        # Ensure we use local model cost map for consistent testing
        os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"

        # Load the model cost map
        try:
            with open("model_prices_and_context_window.json", "r") as f:
                model_cost_map = json.load(f)
            litellm.model_cost = model_cost_map
        except FileNotFoundError:
            # Fallback to ensure we have the Perplexity model configuration
            litellm.model_cost = {
                "perplexity/sonar-deep-research": {
                    "max_tokens": 128000,
                    "max_input_tokens": 128000,
                    "input_cost_per_token": 2e-06,
                    "output_cost_per_token": 8e-06,
                    "output_cost_per_reasoning_token": 3e-06,
                    "citation_cost_per_token": 2e-06,
                    "search_context_cost_per_query": {
                        "search_context_size_low": 0.005,
                        "search_context_size_medium": 0.005,
                        "search_context_size_high": 0.005,
                    },
                    "litellm_provider": "perplexity",
                    "mode": "chat",
                    "supports_reasoning": True,
                    "supports_web_search": True,
                }
            }

    def test_end_to_end_cost_calculation_with_transformation(self):
        """Test end-to-end cost calculation with response transformation."""
        # Create a Perplexity API response that includes citations and search queries
        config = PerplexityChatConfig()

        # Create a ModelResponse with basic usage (before transformation)
        model_response = ModelResponse()
        model_response.model = "sonar-deep-research"
        model_response.usage = Usage(
            prompt_tokens=100,
            completion_tokens=50,
            total_tokens=150,
            reasoning_tokens=10,
        )

        # Simulate raw response from Perplexity API
        raw_response_dict = {
            "choices": [{"message": {"content": "Test response with citations"}}],
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
                "num_search_queries": 2,
            },
            "citations": [
                "This is the first citation with important information about the topic",
                "Another citation providing additional context for the response",
            ],
        }

        # Apply transformation to extract Perplexity-specific fields
        config._enhance_usage_with_perplexity_fields(model_response, raw_response_dict)

        # Now calculate the cost with the enhanced usage
        total_cost = completion_cost(
            completion_response=model_response, custom_llm_provider="perplexity"
        )

        # Calculate expected cost
        citation_chars = sum(
            len(citation) for citation in raw_response_dict["citations"]
        )
        citation_tokens = citation_chars // 4

        expected_prompt_cost = (100 * 2e-6) + (citation_tokens * 2e-6)
        expected_completion_cost = (
            ((50 - 10) * 8e-6) + (10 * 3e-6) + (2 * 0.005)
        )  # Output (text) + reasoning + search
        expected_total = expected_prompt_cost + expected_completion_cost

        assert math.isclose(total_cost, expected_total, rel_tol=1e-6)

    def test_cost_calculation_without_custom_fields(self):
        """Test that cost calculation works normally when custom fields are absent."""
        # Create a standard response without Perplexity-specific fields
        model_response = ModelResponse()
        model_response.model = "sonar-deep-research"
        model_response.usage = Usage(
            prompt_tokens=100, completion_tokens=50, total_tokens=150
        )

        # Calculate cost without custom fields
        total_cost = completion_cost(
            completion_response=model_response, custom_llm_provider="perplexity"
        )

        # Should only include basic input/output costs
        expected_cost = (100 * 2e-6) + (50 * 8e-6)

        assert math.isclose(total_cost, expected_cost, rel_tol=1e-6)

    def test_main_cost_calculator_integration(self):
        """Test integration with the main LiteLLM cost calculator."""
        # Create usage with all Perplexity fields
        usage = Usage(
            prompt_tokens=200,
            completion_tokens=100,
            total_tokens=300,
            reasoning_tokens=25,
            prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=3),
        )
        usage.citation_tokens = 40

        # Test main cost calculator
        prompt_cost, completion_cost_val = cost_per_token(
            model="sonar-deep-research",
            custom_llm_provider="perplexity",
            usage_object=usage,
        )

        expected_prompt_cost = (200 * 2e-6) + (40 * 2e-6)
        expected_completion_cost = (
            ((100 - 25) * 8e-6) + (25 * 3e-6) + (3 * 0.005)
        )  # Output (text) + reasoning + search

        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
        assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)

    def test_model_info_includes_custom_fields(self):
        """Test that get_model_info returns the custom Perplexity cost fields."""
        model_info = get_model_info(
            model="sonar-deep-research", custom_llm_provider="perplexity"
        )

        # Verify custom fields are included
        required_fields = [
            "citation_cost_per_token",
            "search_context_cost_per_query",
            "input_cost_per_token",
            "output_cost_per_token",
            "output_cost_per_reasoning_token",
        ]

        for field in required_fields:
            assert field in model_info, f"Missing field: {field}"
            assert model_info[field] is not None, f"Null value for field: {field}"

    def test_various_citation_sizes(self):
        """Test cost calculation with various citation sizes."""
        config = PerplexityChatConfig()

        test_cases = [
            # (citations, expected_approximate_tokens)
            (["Short"], 1),
            (["This is a medium-length citation with some content"], 12),
            (
                ["Very short", "Another citation", "Third one with more text content"],
                15,
            ),
            ([""], 0),  # Empty citation
        ]

        for citations, expected_approx_tokens in test_cases:
            model_response = ModelResponse()
            model_response.model = "sonar-deep-research"
            model_response.usage = Usage(
                prompt_tokens=100, completion_tokens=50, total_tokens=150
            )

            raw_response_dict = {
                "usage": {
                    "prompt_tokens": 100,
                    "completion_tokens": 50,
                    "total_tokens": 150,
                },
                "citations": citations,
            }

            config._enhance_usage_with_perplexity_fields(
                model_response, raw_response_dict
            )

            citation_tokens = getattr(model_response.usage, "citation_tokens", 0)

            # Allow for reasonable variance in token estimation
            if expected_approx_tokens == 0:
                assert citation_tokens == 0
            else:
                assert abs(citation_tokens - expected_approx_tokens) <= 5

    def test_cost_calculation_with_zero_values(self):
        """Test cost calculation handles zero values for custom fields correctly."""
        usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)

        # Set custom fields to zero
        usage.citation_tokens = 0
        usage.prompt_tokens_details = PromptTokensDetailsWrapper(web_search_requests=0)

        # Should not add any extra cost
        prompt_cost, completion_cost_val = cost_per_token(
            model="sonar-deep-research",
            custom_llm_provider="perplexity",
            usage_object=usage,
        )

        expected_prompt_cost = 100 * 2e-6
        expected_completion_cost = 50 * 8e-6

        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
        assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)

    def test_high_volume_cost_calculation(self):
        """Test cost calculation with high token and query counts."""
        usage = Usage(
            prompt_tokens=50000,
            completion_tokens=25000,
            total_tokens=75000,
            reasoning_tokens=10000,
        )

        usage.citation_tokens = 5000
        usage.prompt_tokens_details = PromptTokensDetailsWrapper(
            web_search_requests=100
        )

        total_cost = completion_cost(
            completion_response=ModelResponse(usage=usage, model="sonar-deep-research"),
            custom_llm_provider="perplexity",
        )

        expected_prompt_cost = (50000 * 2e-6) + (5000 * 2e-6)
        expected_completion_cost = (
            ((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 * 0.005)
        )  # $0.65
        expected_total = expected_prompt_cost + expected_completion_cost  # $0.76

        assert math.isclose(total_cost, expected_total, rel_tol=1e-6)
        assert total_cost > 0.25

    def test_transformation_preserves_existing_usage_fields(self):
        """Test that transformation doesn't overwrite existing standard usage fields."""
        config = PerplexityChatConfig()

        model_response = ModelResponse()
        model_response.usage = Usage(
            prompt_tokens=100,
            completion_tokens=50,
            total_tokens=150,
            reasoning_tokens=20,
        )

        # Store original values
        original_prompt_tokens = model_response.usage.prompt_tokens
        original_completion_tokens = model_response.usage.completion_tokens
        original_total_tokens = model_response.usage.total_tokens

        raw_response_dict = {
            "usage": {
                "prompt_tokens": 999,  # Different from original
                "completion_tokens": 999,  # Different from original
                "total_tokens": 999,  # Different from original
                "num_search_queries": 3,
            },
            "citations": ["Some citation"],
        }

        config._enhance_usage_with_perplexity_fields(model_response, raw_response_dict)

        # Original usage fields should be preserved
        assert model_response.usage.prompt_tokens == original_prompt_tokens
        assert model_response.usage.completion_tokens == original_completion_tokens
        assert model_response.usage.total_tokens == original_total_tokens

        # But custom fields should be added
        assert hasattr(model_response.usage, "prompt_tokens_details")
        assert hasattr(model_response.usage, "citation_tokens")
        assert model_response.usage.prompt_tokens_details.web_search_requests == 3

    @pytest.mark.parametrize(
        "provider_name", ["perplexity", "PERPLEXITY", "Perplexity"]
    )
    def test_case_insensitive_provider_matching(self, provider_name):
        """Test that cost calculation works with different case variations of provider name."""
        usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
        usage.citation_tokens = 10
        usage.prompt_tokens_details = PromptTokensDetailsWrapper(web_search_requests=1)

        # Should work regardless of case
        prompt_cost, completion_cost_val = cost_per_token(
            model="sonar-deep-research",
            custom_llm_provider=provider_name.lower(),  # Normalize to lowercase
            usage_object=usage,
        )

        # Should calculate costs correctly
        expected_prompt_cost = (100 * 2e-6) + (10 * 2e-6)
        expected_completion_cost = (50 * 8e-6) + (1 * 0.005)

        assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6)
        assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6)
