"""
Tests for LiteLLM integration.

Covers:
    - log_success_event creates span with correct attributes
    - Disabled tracer passthrough
    - log_failure_event records error
    - Token usage capture
    - Cost capture
    - Provider suppression during callback
    - Streaming completion creates span
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest


class TestLiteLLMIntegration:
    """Tests for RisicareLiteLLMLogger callback."""

    @pytest.fixture
    def mock_tracer(self):
        tracer = MagicMock()
        tracer.is_enabled = True
        mock_span = MagicMock()
        mock_span.span_id = "abcdef0123456789"
        mock_span.trace_id = "a" * 32
        tracer.start_span_no_context.return_value = mock_span
        return tracer

    @pytest.fixture
    def logger_instance(self):
        from risicare.integrations.litellm._callback import RisicareLiteLLMLogger
        return RisicareLiteLLMLogger()

    @pytest.fixture
    def mock_response(self):
        response = MagicMock()
        response.model = "gpt-4o-2024-08-06"
        response.usage = MagicMock()
        response.usage.prompt_tokens = 100
        response.usage.completion_tokens = 50
        response.usage.total_tokens = 150
        response.choices = [MagicMock()]
        response.choices[0].message = MagicMock()
        response.choices[0].message.content = "Hello, world!"
        return response

    def test_log_success_creates_span(self, mock_tracer, logger_instance, mock_response):
        """Successful completion should create a span with correct attributes."""
        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            # Pre-call
            logger_instance.log_pre_api_call(
                model="gpt-4o",
                messages=[{"role": "user", "content": "Hi"}],
                kwargs={"custom_llm_provider": "openai", "stream": False},
            )

            # Success
            logger_instance.log_success_event(
                kwargs={"custom_llm_provider": "openai", "response_cost": 0.001},
                response_obj=mock_response,
                start_time=None,
                end_time=None,
            )

        # Verify span was created
        mock_tracer.start_span_no_context.assert_called_once()
        call_kwargs = mock_tracer.start_span_no_context.call_args
        assert "litellm.completion/gpt-4o" in str(call_kwargs)

        # Verify span was ended
        mock_span = mock_tracer.start_span_no_context.return_value
        mock_span.end.assert_called_once()

    def test_disabled_tracer_passthrough(self, logger_instance):
        """When tracing is disabled, no span should be created."""
        with patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=False,
        ):
            logger_instance.log_pre_api_call(
                model="gpt-4o",
                messages=[],
                kwargs={},
            )

        # No tracer calls should have been made

    def test_log_failure_records_error(self, mock_tracer, logger_instance):
        """Failed completion should record error on the span."""
        error = RuntimeError("API rate limit exceeded")

        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            # Pre-call
            logger_instance.log_pre_api_call(
                model="gpt-4o",
                messages=[],
                kwargs={"custom_llm_provider": "openai"},
            )

            # Failure
            logger_instance.log_failure_event(
                kwargs={"exception": error},
                response_obj=None,
                start_time=None,
                end_time=None,
            )

        mock_span = mock_tracer.start_span_no_context.return_value
        mock_span.record_exception.assert_called_once_with(error)
        mock_span.set_attribute.assert_any_call("error", True)
        mock_span.end.assert_called_once()

    def test_captures_token_usage(self, mock_tracer, logger_instance, mock_response):
        """Token usage should be captured from response_obj.usage."""
        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            logger_instance.log_pre_api_call(
                model="gpt-4o", messages=[], kwargs={}
            )
            logger_instance.log_success_event(
                kwargs={}, response_obj=mock_response, start_time=None, end_time=None
            )

        mock_span = mock_tracer.start_span_no_context.return_value
        mock_span.set_attribute.assert_any_call("gen_ai.usage.prompt_tokens", 100)
        mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 50)
        mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150)

    def test_captures_cost(self, mock_tracer, logger_instance, mock_response):
        """response_cost should be captured as gen_ai.usage.cost_usd."""
        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            logger_instance.log_pre_api_call(
                model="claude-3-5-sonnet",
                messages=[],
                kwargs={"custom_llm_provider": "anthropic"},
            )
            logger_instance.log_success_event(
                kwargs={"response_cost": 0.0045},
                response_obj=mock_response,
                start_time=None,
                end_time=None,
            )

        mock_span = mock_tracer.start_span_no_context.return_value
        mock_span.set_attribute.assert_any_call("gen_ai.usage.cost_usd", 0.0045)

    def test_provider_suppression_active(self, mock_tracer, logger_instance):
        """Provider suppression should be active during callback execution."""
        from risicare.integrations._dedup import is_provider_suppressed

        suppression_was_active = False

        def check_suppression(*args, **kwargs):
            nonlocal suppression_was_active
            suppression_was_active = is_provider_suppressed()
            return MagicMock()

        # Replace the actual LLM call with our check
        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            logger_instance.log_pre_api_call(
                model="gpt-4o", messages=[], kwargs={}
            )

            # At this point, suppression should be active
            assert is_provider_suppressed()

            # Clean up by calling success
            logger_instance.log_success_event(
                kwargs={}, response_obj=MagicMock(), start_time=None, end_time=None
            )

            # After cleanup, suppression should be released
            assert not is_provider_suppressed()

    def test_stream_completion_creates_span(self, mock_tracer, logger_instance, mock_response):
        """Streaming completions should also create spans."""
        with patch(
            "risicare.integrations.litellm._callback.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.litellm._callback.is_tracing_enabled",
            return_value=True,
        ):
            logger_instance.log_pre_api_call(
                model="gpt-4o",
                messages=[],
                kwargs={"stream": True, "custom_llm_provider": "openai"},
            )
            logger_instance.log_success_event(
                kwargs={}, response_obj=mock_response, start_time=None, end_time=None
            )

        mock_tracer.start_span_no_context.assert_called_once()
        call_kwargs = mock_tracer.start_span_no_context.call_args
        attrs = call_kwargs.kwargs.get("attributes", call_kwargs[1].get("attributes", {}))
        assert attrs.get("framework.litellm.stream") is True
