"""
Unit tests for SDK instrumentation

Tests:
1. Instrumentation can be applied and removed
2. Instrumented SDK calls create spans
3. Span attributes are correct
4. Cost calculation is accurate
5. Error handling works correctly
"""

import pytest
from unittest.mock import Mock, patch, MagicMock
import sys
import os

# Setup path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../sdk/python'))

from kalibr.instrumentation import auto_instrument, get_instrumented_providers
from kalibr.collector import setup_collector, shutdown_collector
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult


class InMemorySpanExporter(SpanExporter):
    """In-memory span exporter for testing"""
    
    def __init__(self):
        self.spans = []
    
    def export(self, spans):
        self.spans.extend(spans)
        return SpanExportResult.SUCCESS
    
    def shutdown(self):
        pass
    
    def clear(self):
        self.spans.clear()


@pytest.fixture
def tracer_provider():
    """Setup tracer provider with in-memory exporter"""
    provider = TracerProvider()
    exporter = InMemorySpanExporter()
    processor = SimpleSpanProcessor(exporter)
    provider.add_span_processor(processor)
    trace.set_tracer_provider(provider)
    
    yield provider, exporter
    
    # Cleanup
    provider.shutdown()
    exporter.clear()


class TestOpenAIInstrumentation:
    """Tests for OpenAI SDK instrumentation"""
    
    def test_can_instrument_openai(self):
        """Test that OpenAI SDK can be instrumented"""
        from kalibr.instrumentation.openai_instr import get_instrumentation
        
        instr = get_instrumentation()
        
        # Mock the OpenAI SDK
        with patch('openai.resources.chat.completions.Completions') as mock:
            mock.create = Mock()
            success = instr.instrument()
            assert success or not success  # May fail if OpenAI not installed
    
    def test_openai_cost_calculation(self):
        """Test OpenAI cost calculation"""
        from kalibr.instrumentation.openai_instr import OpenAICostAdapter
        
        adapter = OpenAICostAdapter()
        
        # Test GPT-4o-mini pricing
        cost = adapter.calculate_cost(
            "gpt-4o-mini",
            {"prompt_tokens": 1000, "completion_tokens": 500}
        )
        
        # GPT-4o-mini: $0.15 input, $0.60 output per 1M tokens
        expected = (1000 / 1_000_000 * 0.15) + (500 / 1_000_000 * 0.60)
        assert abs(cost - expected) < 0.000001
    
    def test_openai_span_attributes(self, tracer_provider):
        """Test that OpenAI instrumentation creates correct span attributes"""
        provider, exporter = tracer_provider
        
        try:
            from kalibr.instrumentation.openai_instr import get_instrumentation
            import openai
            
            # Instrument
            instr = get_instrumentation()
            instr.instrument()
            
            # Mock response
            mock_response = Mock()
            mock_response.model = "gpt-4o-mini"
            mock_response.id = "test-123"
            mock_response.usage = Mock(
                prompt_tokens=100,
                completion_tokens=50,
                total_tokens=150
            )
            mock_response.choices = [Mock(message=Mock(content="Test response"))]
            
            # Mock the create method
            with patch.object(openai.OpenAI, 'chat') as mock_chat:
                mock_completions = Mock()
                mock_completions.completions = Mock()
                mock_completions.completions.create = Mock(return_value=mock_response)
                mock_chat.return_value = mock_completions
                
                client = openai.OpenAI(api_key="test-key")
                response = client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[{"role": "user", "content": "test"}]
                )
                
                # Check spans
                assert len(exporter.spans) >= 1
                span = exporter.spans[-1]
                
                # Check attributes
                attrs = dict(span.attributes)
                assert attrs.get("llm.vendor") == "openai"
                assert attrs.get("llm.request.model") == "gpt-4o-mini"
                
        except ImportError:
            pytest.skip("OpenAI SDK not installed")
        finally:
            # Cleanup
            instr.uninstrument()


class TestAnthropicInstrumentation:
    """Tests for Anthropic SDK instrumentation"""
    
    def test_anthropic_cost_calculation(self):
        """Test Anthropic cost calculation"""
        from kalibr.instrumentation.anthropic_instr import AnthropicCostAdapter
        
        adapter = AnthropicCostAdapter()
        
        # Test Claude 3 Haiku pricing
        cost = adapter.calculate_cost(
            "claude-3-haiku",
            {"input_tokens": 1000, "output_tokens": 500}
        )
        
        # Claude 3 Haiku: $0.25 input, $1.25 output per 1M tokens
        expected = (1000 / 1_000_000 * 0.25) + (500 / 1_000_000 * 1.25)
        assert abs(cost - expected) < 0.000001


class TestGoogleInstrumentation:
    """Tests for Google Generative AI SDK instrumentation"""
    
    def test_google_cost_calculation(self):
        """Test Google cost calculation"""
        from kalibr.instrumentation.google_instr import GoogleCostAdapter
        
        adapter = GoogleCostAdapter()
        
        # Test Gemini 1.5 Flash pricing
        cost = adapter.calculate_cost(
            "gemini-1.5-flash",
            {"prompt_tokens": 1000, "completion_tokens": 500}
        )
        
        # Gemini 1.5 Flash: $0.075 input, $0.30 output per 1M tokens
        expected = (1000 / 1_000_000 * 0.075) + (500 / 1_000_000 * 0.30)
        assert abs(cost - expected) < 0.000001


class TestAutoInstrumentation:
    """Tests for auto-instrumentation registry"""
    
    def test_auto_instrument_returns_results(self):
        """Test that auto_instrument returns status dict"""
        results = auto_instrument(["openai", "anthropic", "google"])
        
        assert isinstance(results, dict)
        assert "openai" in results or "anthropic" in results or "google" in results
    
    def test_get_instrumented_providers(self):
        """Test getting list of instrumented providers"""
        providers = get_instrumented_providers()
        assert isinstance(providers, list)
    
    def test_auto_instrument_unknown_provider(self):
        """Test that unknown providers are handled gracefully"""
        results = auto_instrument(["unknown_provider"])
        assert results.get("unknown_provider") == False


class TestCollector:
    """Tests for OpenTelemetry collector setup"""
    
    def test_setup_collector(self):
        """Test that collector can be setup"""
        provider = setup_collector(
            service_name="test",
            file_export=False,
            console_export=False
        )
        
        assert provider is not None
        assert isinstance(provider, TracerProvider)
        
        # Cleanup
        shutdown_collector()
    
    def test_collector_configured_state(self):
        """Test collector configured state tracking"""
        from kalibr.collector import is_configured
        
        # Should be configured after setup
        setup_collector(service_name="test", file_export=False)
        assert is_configured()
        
        # Should not be configured after shutdown
        shutdown_collector()
        # Note: is_configured may still return True due to global state


class TestCostConsistency:
    """Test consistency between instrumentation adapters and centralized pricing"""
    
    def test_openai_adapter_consistency(self):
        """Test that OpenAI instrumentation adapter matches pricing module"""
        from kalibr.instrumentation.openai_instr import OpenAICostAdapter
        from kalibr.pricing import compute_cost as pricing_compute_cost
        
        adapter = OpenAICostAdapter()
        
        # Test several models
        test_cases = [
            ("gpt-4o", 1000, 500),
            ("gpt-4", 2000, 1000),
            ("gpt-4o-mini", 5000, 2500),
        ]
        
        for model, input_tokens, output_tokens in test_cases:
            adapter_cost = adapter.calculate_cost(
                model, 
                {"prompt_tokens": input_tokens, "completion_tokens": output_tokens}
            )
            pricing_cost = pricing_compute_cost("openai", model, input_tokens, output_tokens)
            assert adapter_cost == pricing_cost, f"Mismatch for {model}"
    
    def test_anthropic_adapter_consistency(self):
        """Test that Anthropic instrumentation adapter matches pricing module"""
        from kalibr.instrumentation.anthropic_instr import AnthropicCostAdapter
        from kalibr.pricing import compute_cost as pricing_compute_cost
        
        adapter = AnthropicCostAdapter()
        
        # Test several models
        test_cases = [
            ("claude-3-opus", 1000, 500),
            ("claude-3-sonnet", 2000, 1000),
            ("claude-3-haiku", 5000, 2500),
        ]
        
        for model, input_tokens, output_tokens in test_cases:
            adapter_cost = adapter.calculate_cost(
                model,
                {"input_tokens": input_tokens, "output_tokens": output_tokens}
            )
            pricing_cost = pricing_compute_cost("anthropic", model, input_tokens, output_tokens)
            assert adapter_cost == pricing_cost, f"Mismatch for {model}"
    
    def test_google_adapter_consistency(self):
        """Test that Google instrumentation adapter matches pricing module"""
        from kalibr.instrumentation.google_instr import GoogleCostAdapter
        from kalibr.pricing import compute_cost as pricing_compute_cost
        
        adapter = GoogleCostAdapter()
        
        # Test several models
        test_cases = [
            ("gemini-1.5-pro", 1000, 500),
            ("gemini-1.5-flash", 2000, 1000),
            ("gemini-pro", 5000, 2500),
        ]
        
        for model, input_tokens, output_tokens in test_cases:
            adapter_cost = adapter.calculate_cost(
                model,
                {"prompt_tokens": input_tokens, "completion_tokens": output_tokens}
            )
            pricing_cost = pricing_compute_cost("google", model, input_tokens, output_tokens)
            assert adapter_cost == pricing_cost, f"Mismatch for {model}"
    
    def test_all_adapters_same_cost_for_same_model(self):
        """Test that using different adapters for same vendor produces same cost"""
        from kalibr.cost_adapter import OpenAICostAdapter as CoreOpenAIAdapter
        from kalibr.instrumentation.openai_instr import OpenAICostAdapter as InstrOpenAIAdapter
        
        core_adapter = CoreOpenAIAdapter()
        instr_adapter = InstrOpenAIAdapter()
        
        # Core adapter uses compute_cost(model, tokens_in, tokens_out)
        # Instrumentation adapter uses calculate_cost(model, usage_dict)
        core_cost = core_adapter.compute_cost("gpt-4o", 1000, 500)
        instr_cost = instr_adapter.calculate_cost(
            "gpt-4o",
            {"prompt_tokens": 1000, "completion_tokens": 500}
        )
        
        assert core_cost == instr_cost


class TestOpenAIResponsesInstrumentation:
    """Tests for OpenAI Responses API instrumentation"""

    def test_can_instrument_responses_api(self):
        """Test that OpenAI Responses API can be instrumented"""
        from kalibr.instrumentation.openai_responses_instr import get_instrumentation

        instr = get_instrumentation()
        success = instr.instrument()
        assert success is True
        assert instr.is_instrumented is True
        instr.uninstrument()

    def test_can_uninstrument_responses_api(self):
        """Test that OpenAI Responses API can be uninstrumented"""
        from kalibr.instrumentation.openai_responses_instr import get_instrumentation

        instr = get_instrumentation()
        instr.instrument()
        success = instr.uninstrument()
        assert success is True
        assert instr.is_instrumented is False

    def test_responses_cost_calculation(self):
        """Test cost calculation for OpenAI Responses API"""
        from kalibr.instrumentation.openai_responses_instr import OpenAIResponsesCostAdapter

        adapter = OpenAIResponsesCostAdapter()

        cost = adapter.calculate_cost(
            "gpt-4o-mini",
            {"input_tokens": 1000, "output_tokens": 500}
        )

        # GPT-4o-mini: $0.15 input, $0.60 output per 1M tokens
        expected = (1000 / 1_000_000 * 0.15) + (500 / 1_000_000 * 0.60)
        assert abs(cost - expected) < 0.000001

    def test_responses_cost_adapter_vendor_name(self):
        """Test that vendor name is openai"""
        from kalibr.instrumentation.openai_responses_instr import OpenAIResponsesCostAdapter

        adapter = OpenAIResponsesCostAdapter()
        assert adapter.get_vendor_name() == "openai"

    def test_responses_cost_consistency_with_pricing_module(self):
        """Test that Responses cost adapter matches centralized pricing"""
        from kalibr.instrumentation.openai_responses_instr import OpenAIResponsesCostAdapter
        from kalibr.pricing import compute_cost as pricing_compute_cost

        adapter = OpenAIResponsesCostAdapter()

        test_cases = [
            ("gpt-4o", 1000, 500),
            ("gpt-4", 2000, 1000),
            ("gpt-4o-mini", 5000, 2500),
        ]

        for model, input_tokens, output_tokens in test_cases:
            adapter_cost = adapter.calculate_cost(
                model,
                {"input_tokens": input_tokens, "output_tokens": output_tokens}
            )
            pricing_cost = pricing_compute_cost("openai", model, input_tokens, output_tokens)
            assert adapter_cost == pricing_cost, f"Mismatch for {model}"

    def test_responses_module_level_functions(self):
        """Test module-level instrument/uninstrument functions"""
        from kalibr.instrumentation import openai_responses_instr

        success = openai_responses_instr.instrument()
        assert success is True

        success = openai_responses_instr.uninstrument()
        assert success is True

    def test_responses_singleton_pattern(self):
        """Test that get_instrumentation returns same instance"""
        from kalibr.instrumentation.openai_responses_instr import get_instrumentation

        instr1 = get_instrumentation()
        instr2 = get_instrumentation()
        assert instr1 is instr2

    def test_auto_instrument_includes_openai_responses(self):
        """Test that auto_instrument with defaults includes openai_responses"""
        results = auto_instrument()
        assert "openai_responses" in results

    def test_responses_patches_create_method(self):
        """Test that instrumentation monkey-patches Responses.create"""
        from kalibr.instrumentation.openai_responses_instr import get_instrumentation
        from openai.resources.responses import Responses

        instr = get_instrumentation()
        original = Responses.create

        instr.instrument()
        assert Responses.create is not original

        instr.uninstrument()
        assert Responses.create is original


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