"""
Test Arize health check functionality and proxy integration.
"""
import json
import os
import sys
from unittest.mock import patch, MagicMock

# Adds the grandparent directory to sys.path to allow importing project modules
sys.path.insert(0, os.path.abspath("../.."))

import asyncio
import pytest

import litellm
from litellm.integrations.arize.arize import ArizeLogger
from litellm.types.utils import StandardCallbackDynamicParams


class TestArizeHealthCheck:
    """Test Arize health check functionality."""

    @pytest.mark.asyncio
    async def test_arize_health_check_with_credentials(self):
        """Test Arize health check returns healthy when credentials are available."""
        
        with patch.dict(os.environ, {
            "ARIZE_SPACE_KEY": "test-space-key",
            "ARIZE_API_KEY": "test-api-key",
            "ARIZE_ENDPOINT": "https://otlp.arize.com/v1"
        }):
            arize_logger = ArizeLogger()
            response = await arize_logger.async_health_check()
            
            assert response["status"] == "healthy"
            assert "configured properly" in response["message"]

    @pytest.mark.asyncio
    async def test_arize_health_check_missing_space_key(self):
        """Test Arize health check returns unhealthy when space key is missing."""
        
        with patch.dict(os.environ, {
            "ARIZE_API_KEY": "test-api-key"
        }, clear=True):
            arize_logger = ArizeLogger()
            response = await arize_logger.async_health_check()
            
            assert response["status"] == "unhealthy"
            assert "ARIZE_SPACE_KEY" in response["error_message"]

    @pytest.mark.asyncio
    async def test_arize_health_check_missing_api_key(self):
        """Test Arize health check returns unhealthy when API key is missing."""
        
        with patch.dict(os.environ, {
            "ARIZE_SPACE_KEY": "test-space-key"
        }, clear=True):
            arize_logger = ArizeLogger()
            response = await arize_logger.async_health_check()
            
            assert response["status"] == "unhealthy"
            assert "ARIZE_API_KEY" in response["error_message"]

    @pytest.mark.asyncio
    async def test_arize_health_check_missing_both_keys(self):
        """Test Arize health check when both keys are missing."""
        
        with patch.dict(os.environ, {}, clear=True):
            arize_logger = ArizeLogger()
            response = await arize_logger.async_health_check()
            
            assert response["status"] == "unhealthy"
            assert "ARIZE_SPACE_KEY" in response["error_message"]


class TestArizeIntegrationWithProxy:
    """Test Arize integration with LiteLLM completion requests."""

    @pytest.mark.asyncio
    async def test_arize_logging_with_completion(self):
        """Test that Arize logging works with actual completion requests."""
        
        with patch.dict(os.environ, {
            "ARIZE_SPACE_KEY": "test-space-key", 
            "ARIZE_API_KEY": "test-api-key",
            "ARIZE_ENDPOINT": "https://otlp.arize.com/v1"
        }):
            # Create ArizeLogger instance
            arize_logger = ArizeLogger()
            
            # Store original callbacks
            original_callbacks = litellm.success_callback.copy() if litellm.success_callback else []
            
            try:
                # Add ArizeLogger to callbacks
                litellm.success_callback = [arize_logger]
                
                # Make completion request
                response = await litellm.acompletion(
                    model="openai/litellm-mock-response-model",
                    messages=[{"role": "user", "content": "Test message for Arize health check"}],
                    mock_response="This is a test response that validates Arize integration.",
                    user="test-arize-health"
                )
                
                # Verify response is valid
                assert response is not None
                print(f"Response type: {type(response)}")
                print("✅ Arize completion request completed successfully")
                
                # Give time for async logging
                await asyncio.sleep(0.1)
                
                print("✅ Arize completion logging test successful")
                
            finally:
                # Restore original callbacks
                litellm.success_callback = original_callbacks

    def test_arize_get_config(self):
        """Test ArizeLogger.get_arize_config() method."""
        
        with patch.dict(os.environ, {
            "ARIZE_SPACE_KEY": "test-space-123",
            "ARIZE_API_KEY": "test-api-456", 
            "ARIZE_ENDPOINT": "https://custom.arize.com/v1",
            "ARIZE_PROJECT_NAME": "custom-project",
        }):
            config = ArizeLogger.get_arize_config()
            
            assert config.space_key == "test-space-123"
            assert config.api_key == "test-api-456"
            assert config.endpoint == "https://custom.arize.com/v1"
            assert config.protocol == "otlp_grpc"
            assert config.project_name == "custom-project"

    def test_arize_get_config_defaults(self):
        """Test ArizeLogger.get_arize_config() with default endpoint."""
        
        with patch.dict(os.environ, {
            "ARIZE_SPACE_KEY": "test-space-default",
            "ARIZE_API_KEY": "test-api-default",
            "ARIZE_PROJECT_NAME": "default-project",
        }, clear=True):
            config = ArizeLogger.get_arize_config()
            
            assert config.space_key == "test-space-default"
            assert config.api_key == "test-api-default"
            assert config.endpoint == "https://otlp.arize.com/v1"  # Default endpoint
            assert config.protocol == "otlp_grpc"  # Default protocol
            assert config.project_name == "default-project"

    def test_arize_construct_dynamic_headers(self):
        """Test dynamic OTEL headers construction for team/key logging."""
        
        arize_logger = ArizeLogger()
        
        dynamic_params = StandardCallbackDynamicParams(
            arize_space_key="dynamic-space-123",
            arize_api_key="dynamic-api-456"
        )
        
        headers = arize_logger.construct_dynamic_otel_headers(dynamic_params)
        
        assert headers is not None
        assert headers["arize-space-id"] == "dynamic-space-123"
        assert headers["api_key"] == "dynamic-api-456"

    def test_arize_construct_dynamic_headers_space_id_fallback(self):
        """Test dynamic headers with arize_space_id parameter (fallback)."""
        
        arize_logger = ArizeLogger()
        
        dynamic_params = StandardCallbackDynamicParams(
            arize_space_id="fallback-space-789",  # Using space_id instead of space_key
            arize_api_key="fallback-api-999"
        )
        
        headers = arize_logger.construct_dynamic_otel_headers(dynamic_params)
        
        assert headers is not None
        assert headers["arize-space-id"] == "fallback-space-789"
        assert headers["api_key"] == "fallback-api-999"


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