"""
Tests for LlamaIndex integration.

Covers:
    - new_span creates Risicare span
    - Disabled tracer passthrough
    - prepare_to_exit_span ends span
    - prepare_to_drop_span records error
    - Nested spans have parent
    - Component name captured
    - Provider suppression during LLM spans
    - Memory guard on _open_spans
"""

from __future__ import annotations

import time
from unittest.mock import MagicMock, patch

import pytest


class TestLlamaIndexIntegration:
    """Tests for RisicareSpanHandler."""

    @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 handler(self):
        from risicare.integrations.llamaindex._handler import RisicareSpanHandler
        return RisicareSpanHandler()

    def _make_instance(self, cls_name: str, **attrs):
        """Create an instance whose type().__name__ returns cls_name."""
        # Create a real class so type(instance).__name__ works correctly
        # for the handler's _classify_component() which uses type(instance).__name__
        cls = type(cls_name, (), {})
        instance = cls()
        for key, value in attrs.items():
            setattr(instance, key, value)
        return instance

    def test_new_span_creates_risicare_span(self, mock_tracer, handler):
        """new_span should create a Risicare span."""
        instance = self._make_instance("RetrieverQueryEngine")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            result = handler.new_span(
                id_="span-001",
                instance=instance,
            )

        assert result == "span-001"
        assert "span-001" in handler._open_spans
        mock_tracer.start_span_no_context.assert_called_once()

    def test_disabled_tracer_passthrough(self, handler):
        """When tracing is disabled, no span should be created."""
        with patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=False,
        ):
            result = handler.new_span(id_="span-off", instance=MagicMock())

        assert result is None
        assert "span-off" not in handler._open_spans

    def test_prepare_to_exit_ends_span(self, mock_tracer, handler):
        """prepare_to_exit_span should end the span with result attributes."""
        instance = self._make_instance("VectorIndexRetriever")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            handler.new_span(id_="span-exit", instance=instance)
            handler.prepare_to_exit_span(
                id_="span-exit",
                instance=instance,
                result=["doc1", "doc2"],
            )

        mock_span = mock_tracer.start_span_no_context.return_value
        mock_span.end.assert_called_once()
        assert "span-exit" not in handler._open_spans

    def test_prepare_to_drop_records_error(self, mock_tracer, handler):
        """prepare_to_drop_span should record error on the span."""
        instance = self._make_instance("RetrieverQueryEngine")
        error = RuntimeError("Query failed")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            handler.new_span(id_="span-err", instance=instance)
            handler.prepare_to_drop_span(
                id_="span-err",
                instance=instance,
                err=error,
            )

        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_nested_spans_have_parent(self, mock_tracer, handler):
        """Parent-child relationships should be preserved via parent_span_id."""
        query_engine = self._make_instance("RetrieverQueryEngine")
        retriever = self._make_instance("VectorIndexRetriever")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            handler.new_span(id_="parent", instance=query_engine)
            handler.new_span(
                id_="child",
                instance=retriever,
                parent_span_id="parent",
            )

        assert "parent" in handler._open_spans
        assert "child" in handler._open_spans

        # Clean up
        handler.prepare_to_exit_span(id_="child", result=[])
        handler.prepare_to_exit_span(id_="parent", result="answer")

    def test_captures_component_name(self, mock_tracer, handler):
        """Component class name should be in span attributes."""
        instance = self._make_instance("TreeSummarize")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            handler.new_span(id_="span-comp", instance=instance)

        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.llamaindex.component") == "TreeSummarize"
        assert attrs.get("framework.span_kind") == "synthesis"

        handler.prepare_to_exit_span(id_="span-comp", result=None)

    def test_provider_suppression_during_llm(self, mock_tracer, handler):
        """Provider spans should be suppressed during LLM component spans."""
        from risicare.integrations._dedup import is_provider_suppressed

        instance = self._make_instance("OpenAI", model="gpt-4o")

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            handler.new_span(id_="llm-span", instance=instance)

            # Suppression should be active
            assert is_provider_suppressed()

            handler.prepare_to_exit_span(id_="llm-span", result=MagicMock())

            # Suppression should be released
            assert not is_provider_suppressed()

    def test_memory_guard_on_open_spans(self, mock_tracer, handler):
        """_open_spans should be bounded by _MAX_OPEN_SPANS."""
        from risicare.integrations.llamaindex._handler import _MAX_OPEN_SPANS, _SPAN_TTL_SECONDS

        with patch(
            "risicare.integrations.llamaindex._handler.get_tracer",
            return_value=mock_tracer,
        ), patch(
            "risicare.integrations.llamaindex._handler.is_tracing_enabled",
            return_value=True,
        ):
            # Fill to limit with stale entries
            from risicare.integrations.llamaindex._handler import _SpanEntry
            now = time.perf_counter()
            for i in range(_MAX_OPEN_SPANS):
                handler._open_spans[f"old-{i}"] = _SpanEntry(
                    span=MagicMock(),
                    suppression_cm=None,
                    start_time=now - _SPAN_TTL_SECONDS - 100,  # Stale
                    component_kind="component",
                )

            assert len(handler._open_spans) == _MAX_OPEN_SPANS

            # Adding a new span should trigger eviction
            instance = self._make_instance("VectorIndexRetriever")
            handler.new_span(id_="new-span", instance=instance)

            # Stale entries should have been evicted
            assert len(handler._open_spans) < _MAX_OPEN_SPANS
            assert "new-span" in handler._open_spans
