"""
Test Vertex AI Live API Passthrough Feature

This module tests the Vertex AI Live API WebSocket passthrough functionality,
including the logging handler, cost tracking, and WebSocket message processing.
"""

import json
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from typing import Dict, List, Any, Optional

import pytest
import httpx

# Add the parent directory to the system path
sys.path.insert(0, os.path.abspath("../.."))

from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import (
    VertexAILivePassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
    PassThroughEndpointLogging,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import LlmProviders
from litellm.proxy._types import UserAPIKeyAuth


class TestVertexAILivePassthroughLoggingHandler:
    """Test the Vertex AI Live Passthrough Logging Handler"""

    @pytest.fixture
    def handler(self):
        """Create a handler instance for testing"""
        return VertexAILivePassthroughLoggingHandler()

    @pytest.fixture
    def mock_logging_obj(self):
        """Create a mock logging object"""
        mock = MagicMock(spec=LiteLLMLoggingObj)
        mock.model_call_details = {}
        return mock

    @pytest.fixture
    def sample_websocket_messages(self):
        """Sample WebSocket messages for testing"""
        return [
            {
                "type": "session.created",
                "session": {"id": "test-session-123"},
                "timestamp": "2024-01-01T00:00:00Z",
            },
            {
                "type": "response.create",
                "event_id": "event-123",
                "response": {"text": "Hello, how can I help you?"},
                "usageMetadata": {
                    "promptTokenCount": 10,
                    "candidatesTokenCount": 15,
                    "totalTokenCount": 25,
                    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 10}],
                    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 15}],
                },
            },
            {
                "type": "response.done",
                "event_id": "event-123",
                "usageMetadata": {
                    "promptTokenCount": 5,
                    "candidatesTokenCount": 8,
                    "totalTokenCount": 13,
                    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}],
                    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 8}],
                },
            },
        ]

    def test_llm_provider_name_property(self, handler):
        """Test that llm_provider_name returns the correct provider"""
        assert handler.llm_provider_name == LlmProviders.VERTEX_AI

    def test_get_provider_config(self, handler):
        """Test that get_provider_config returns a valid config"""
        config = handler.get_provider_config("gemini-1.5-pro")
        assert config is not None
        # Verify it's a Vertex AI config by checking for expected methods
        assert hasattr(config, "get_supported_openai_params")
        assert hasattr(config, "map_openai_params")

    def test_extract_usage_metadata_single_message(self, handler):
        """Test usage metadata extraction from a single message"""
        messages = [
            {
                "type": "response.create",
                "usageMetadata": {
                    "promptTokenCount": 10,
                    "candidatesTokenCount": 15,
                    "totalTokenCount": 25,
                    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 10}],
                    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 15}],
                },
            }
        ]

        result = handler._extract_usage_metadata_from_websocket_messages(messages)

        assert result is not None
        assert result["promptTokenCount"] == 10
        assert result["candidatesTokenCount"] == 15
        assert result["totalTokenCount"] == 25
        assert len(result["promptTokensDetails"]) == 1
        assert len(result["candidatesTokensDetails"]) == 1

    def test_extract_usage_metadata_multiple_messages(self, handler):
        """Test usage metadata aggregation from multiple messages"""
        messages = [
            {
                "type": "response.create",
                "usageMetadata": {
                    "promptTokenCount": 10,
                    "candidatesTokenCount": 15,
                    "totalTokenCount": 25,
                    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 10}],
                    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 15}],
                },
            },
            {
                "type": "response.done",
                "usageMetadata": {
                    "promptTokenCount": 5,
                    "candidatesTokenCount": 8,
                    "totalTokenCount": 13,
                    "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}],
                    "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 8}],
                },
            },
        ]

        result = handler._extract_usage_metadata_from_websocket_messages(messages)

        assert result is not None
        assert result["promptTokenCount"] == 15  # 10 + 5
        assert result["candidatesTokenCount"] == 23  # 15 + 8
        assert result["totalTokenCount"] == 38  # 25 + 13
        assert len(result["promptTokensDetails"]) == 1
        assert result["promptTokensDetails"][0]["tokenCount"] == 15
        assert len(result["candidatesTokensDetails"]) == 1
        assert result["candidatesTokensDetails"][0]["tokenCount"] == 23

    def test_extract_usage_metadata_no_usage(self, handler):
        """Test handling of messages without usage metadata"""
        messages = [
            {"type": "session.created", "session": {"id": "test"}},
            {"type": "response.create", "response": {"text": "Hello"}},
        ]

        result = handler._extract_usage_metadata_from_websocket_messages(messages)
        assert result is None

    def test_extract_usage_metadata_empty_list(self, handler):
        """Test handling of empty message list"""
        result = handler._extract_usage_metadata_from_websocket_messages([])
        assert result is None

    def test_extract_usage_metadata_mixed_modalities(self, handler):
        """Test usage metadata extraction with mixed modalities"""
        messages = [
            {
                "type": "response.create",
                "usageMetadata": {
                    "promptTokenCount": 20,
                    "candidatesTokenCount": 30,
                    "totalTokenCount": 50,
                    "promptTokensDetails": [
                        {"modality": "TEXT", "tokenCount": 10},
                        {"modality": "AUDIO", "tokenCount": 10},
                    ],
                    "candidatesTokensDetails": [
                        {"modality": "TEXT", "tokenCount": 20},
                        {"modality": "AUDIO", "tokenCount": 10},
                    ],
                },
            }
        ]

        result = handler._extract_usage_metadata_from_websocket_messages(messages)

        assert result is not None
        assert result["promptTokenCount"] == 20
        assert result["candidatesTokenCount"] == 30
        assert len(result["promptTokensDetails"]) == 2
        assert len(result["candidatesTokensDetails"]) == 2

        # Check modality aggregation
        text_prompt = next(
            d for d in result["promptTokensDetails"] if d["modality"] == "TEXT"
        )
        audio_prompt = next(
            d for d in result["promptTokensDetails"] if d["modality"] == "AUDIO"
        )
        assert text_prompt["tokenCount"] == 10
        assert audio_prompt["tokenCount"] == 10

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
    )
    def test_calculate_cost_basic(self, mock_get_model_info, handler):
        """Test basic cost calculation"""
        mock_get_model_info.return_value = {
            "input_cost_per_token": 0.000001,
            "output_cost_per_token": 0.000002,
        }

        usage_metadata = {
            "promptTokenCount": 100,
            "candidatesTokenCount": 50,
            "totalTokenCount": 150,
        }

        cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)

        # The cost calculation may include additional factors, so we check it's reasonable
        expected_min_cost = (100 * 0.000001) + (50 * 0.000002)
        assert cost >= expected_min_cost
        assert cost > 0

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
    )
    def test_calculate_cost_with_audio(self, mock_get_model_info, handler):
        """Test cost calculation with audio tokens"""
        mock_get_model_info.return_value = {
            "input_cost_per_token": 0.000001,
            "output_cost_per_token": 0.000002,
            "input_cost_per_audio_token": 0.0001,
            "output_cost_per_audio_token": 0.0002,
        }

        usage_metadata = {
            "promptTokenCount": 100,
            "candidatesTokenCount": 50,
            "totalTokenCount": 150,
            "promptTokensDetails": [
                {"modality": "TEXT", "tokenCount": 80},
                {"modality": "AUDIO", "tokenCount": 20},
            ],
            "candidatesTokensDetails": [
                {"modality": "TEXT", "tokenCount": 30},
                {"modality": "AUDIO", "tokenCount": 20},
            ],
        }

        cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)

        # Should include both text and audio costs
        assert cost > 0
        assert cost > (100 * 0.000001) + (
            50 * 0.000002
        )  # Should be higher due to audio

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
    )
    def test_calculate_cost_with_web_search(self, mock_get_model_info, handler):
        """Test cost calculation with web search (tool use)"""
        mock_get_model_info.return_value = {
            "input_cost_per_token": 0.000001,
            "output_cost_per_token": 0.000002,
            "web_search_cost_per_request": 0.01,
        }

        usage_metadata = {
            "promptTokenCount": 100,
            "candidatesTokenCount": 50,
            "totalTokenCount": 150,
            "toolUsePromptTokenCount": 10,
        }

        cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)

        # Should include web search cost
        expected_base_cost = (100 * 0.000001) + (50 * 0.000002)
        # The web search cost might be handled differently, so just check it's reasonable
        assert cost >= expected_base_cost
        assert cost > 0

    def test_vertex_ai_live_passthrough_handler_integration(
        self, handler, mock_logging_obj, sample_websocket_messages
    ):
        """Test the main passthrough handler method"""
        url_route = "/vertex_ai/live"
        start_time = datetime.now()
        end_time = datetime.now()
        request_body = {"messages": [{"role": "user", "content": "Hello"}]}

        result = handler.vertex_ai_live_passthrough_handler(
            websocket_messages=sample_websocket_messages,
            logging_obj=mock_logging_obj,
            url_route=url_route,
            start_time=start_time,
            end_time=end_time,
            request_body=request_body,
        )

        assert "result" in result
        assert "kwargs" in result

        # Check that the result contains expected fields
        result_data = result["result"]
        assert "model" in result_data
        assert "usage" in result_data
        assert "choices" in result_data

        # Check usage data
        usage = result_data["usage"]
        assert "prompt_tokens" in usage
        assert "completion_tokens" in usage
        assert "total_tokens" in usage

    def test_vertex_ai_live_passthrough_handler_no_usage(
        self, handler, mock_logging_obj
    ):
        """Test handler with messages that don't contain usage metadata"""
        messages = [
            {"type": "session.created", "session": {"id": "test"}},
            {"type": "response.create", "response": {"text": "Hello"}},
        ]

        url_route = "/vertex_ai/live"
        start_time = datetime.now()
        end_time = datetime.now()
        request_body = {"messages": [{"role": "user", "content": "Hello"}]}

        result = handler.vertex_ai_live_passthrough_handler(
            websocket_messages=messages,
            logging_obj=mock_logging_obj,
            url_route=url_route,
            start_time=start_time,
            end_time=end_time,
            request_body=request_body,
        )

        assert "result" in result
        assert "kwargs" in result

        # Should still return a valid result even without usage data
        result_data = result["result"]
        # When no usage metadata is found, result_data will be None
        assert result_data is None


class TestVertexAILivePassthroughIntegration:
    """Integration tests for Vertex AI Live passthrough functionality"""

    @pytest.fixture
    def mock_websocket(self):
        """Create a mock WebSocket for testing"""
        websocket = AsyncMock()
        websocket.headers = {"authorization": "Bearer test-token"}
        websocket.client_state = MagicMock()
        websocket.client_state.DISCONNECTED = "disconnected"
        return websocket

    @pytest.fixture
    def mock_user_api_key(self):
        """Create a mock user API key"""
        return UserAPIKeyAuth(
            api_key="test-key",
            user_id="test-user",
            team_id="test-team",
            user_role="customer",
        )

    @pytest.fixture
    def mock_logging_obj(self):
        """Create a mock logging object"""
        mock = MagicMock(spec=LiteLLMLoggingObj)
        mock.model_call_details = {}
        return mock

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request"
    )
    @patch(
        "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
    )
    @patch(
        "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._ensure_access_token_async"
    )
    @patch("litellm.proxy.proxy_server.proxy_logging_obj")
    @pytest.mark.asyncio
    async def test_vertex_ai_live_websocket_passthrough_route(
        self,
        mock_proxy_logging_obj,
        mock_ensure_access_token,
        mock_router,
        mock_websocket_passthrough,
        mock_websocket,
        mock_user_api_key,
        mock_logging_obj,
    ):
        """Test the Vertex AI Live WebSocket passthrough route"""
        from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
            vertex_ai_live_websocket_passthrough,
        )

        # Mock the router methods
        mock_router.get_vertex_credentials.return_value = MagicMock(
            vertex_project="test-project",
            vertex_location="us-central1",
            vertex_credentials="test-credentials",
        )
        mock_router.set_default_vertex_config.return_value = None

        # Mock the access token async call
        mock_ensure_access_token.return_value = ("test-access-token", "test-project")

        # Mock the WebSocket passthrough request - it returns None, not an AsyncMock
        mock_websocket_passthrough.return_value = None

        # Test the route
        result = await vertex_ai_live_websocket_passthrough(
            websocket=mock_websocket, user_api_key_dict=mock_user_api_key
        )

        # Verify that the WebSocket passthrough was called
        mock_websocket_passthrough.assert_called_once()

        # Check the call arguments
        call_args = mock_websocket_passthrough.call_args
        assert call_args[1]["websocket"] == mock_websocket
        assert call_args[1]["user_api_key_dict"] == mock_user_api_key
        assert call_args[1]["endpoint"] == "/vertex_ai/live"

        # The result should be None since websocket_passthrough_request returns None
        assert result is None

    def test_vertex_ai_live_route_detection(self):
        """Test that the route detection works correctly"""
        from litellm.proxy.pass_through_endpoints.success_handler import (
            PassThroughEndpointLogging,
        )

        handler = PassThroughEndpointLogging()

        # Test valid routes
        assert handler.is_vertex_ai_live_route("/vertex_ai/live") == True
        assert handler.is_vertex_ai_live_route("/vertex_ai/live/") == True
        assert handler.is_vertex_ai_live_route("/vertex_ai/live/stream") == True

        # Test invalid routes
        assert handler.is_vertex_ai_live_route("/vertex_ai") == False
        assert handler.is_vertex_ai_live_route("/vertex_ai/discovery") == False
        assert handler.is_vertex_ai_live_route("/openai/chat/completions") == False

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.VertexAILivePassthroughLoggingHandler"
    )
    @pytest.mark.asyncio
    async def test_success_handler_vertex_ai_live_integration(
        self, mock_handler_class, mock_logging_obj
    ):
        """Test the success handler integration with Vertex AI Live"""
        from litellm.proxy.pass_through_endpoints.success_handler import (
            PassThroughEndpointLogging,
        )

        # Mock the handler
        mock_handler = MagicMock()
        mock_handler.vertex_ai_live_passthrough_handler.return_value = {
            "result": {"model": "gemini-1.5-pro", "usage": {"total_tokens": 100}},
            "kwargs": {"test": "value"},
        }
        mock_handler_class.return_value = mock_handler

        # Create success handler
        success_handler = PassThroughEndpointLogging()

        # Mock the route check
        success_handler.is_vertex_ai_live_route = MagicMock(return_value=True)

        # Test data
        response_body = [{"type": "response.create", "response": {"text": "Hello"}}]
        url_route = "/vertex_ai/live"
        start_time = datetime.now()
        end_time = datetime.now()
        request_body = {"messages": [{"role": "user", "content": "Hello"}]}

        # Call the method
        result = await success_handler.pass_through_async_success_handler(
            httpx_response=MagicMock(),
            response_body=response_body,
            logging_obj=mock_logging_obj,
            url_route=url_route,
            result="test",
            start_time=start_time,
            end_time=end_time,
            cache_hit=False,
            request_body=request_body,
            passthrough_logging_payload=MagicMock(),
        )

        # Verify the handler was called
        mock_handler.vertex_ai_live_passthrough_handler.assert_called_once()

        # The method returns None (it doesn't return anything), so just verify it completed without error
        assert result is None


class TestVertexAILivePassthroughErrorHandling:
    """Test error handling in Vertex AI Live passthrough"""

    @pytest.fixture
    def mock_logging_obj(self):
        """Create a mock logging object"""
        mock = MagicMock(spec=LiteLLMLoggingObj)
        mock.model_call_details = {}
        return mock

    def test_invalid_websocket_messages_format(self):
        """Test handling of invalid WebSocket message formats"""
        handler = VertexAILivePassthroughLoggingHandler()

        # Test with invalid message format
        invalid_messages = [
            {"type": "invalid", "data": "not a proper message"},
            "not a dict at all",
            None,
        ]

        # Should not raise an exception
        result = handler._extract_usage_metadata_from_websocket_messages(
            invalid_messages
        )
        assert result is None

    def test_missing_usage_metadata(self):
        """Test handling of messages with missing usage metadata"""
        handler = VertexAILivePassthroughLoggingHandler()

        messages = [
            {"type": "response.create", "response": {"text": "Hello"}},
            {"type": "response.done", "response": {"text": "Done"}},
        ]

        result = handler._extract_usage_metadata_from_websocket_messages(messages)
        assert result is None

    @patch(
        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
    )
    def test_cost_calculation_with_missing_model_info(self, mock_get_model_info):
        """Test cost calculation when model info is missing"""
        handler = VertexAILivePassthroughLoggingHandler()

        # Mock missing model info
        mock_get_model_info.return_value = {}

        usage_metadata = {
            "promptTokenCount": 100,
            "candidatesTokenCount": 50,
            "totalTokenCount": 150,
        }

        # Should not raise an exception, should return 0 or handle gracefully
        cost = handler._calculate_live_api_cost("unknown-model", usage_metadata)
        assert cost == 0.0

    def test_handler_with_none_websocket_messages(self, mock_logging_obj):
        """Test handler with None websocket messages"""
        handler = VertexAILivePassthroughLoggingHandler()

        url_route = "/vertex_ai/live"
        start_time = datetime.now()
        end_time = datetime.now()
        request_body = {"messages": [{"role": "user", "content": "Hello"}]}

        # Should handle None gracefully
        result = handler.vertex_ai_live_passthrough_handler(
            websocket_messages=None,
            logging_obj=mock_logging_obj,
            url_route=url_route,
            start_time=start_time,
            end_time=end_time,
            request_body=request_body,
        )

        assert "result" in result
        assert "kwargs" in result


if __name__ == "__main__":
    pytest.main([__file__])
