"""
LlamaIndex span handler for Risicare SDK.

Implements the BaseSpanHandler interface to create Risicare spans for
all LlamaIndex operations. Uses _open_spans dict with memory guard
(same pattern as LangChain's _run_spans).

Component type detection is defensive, using class name string matching
against known LlamaIndex component classes. Provider suppression is
applied only for LLM and embedding spans.
"""

from __future__ import annotations

import logging
import threading
import time
from typing import Any, Dict, Optional, Set

from risicare.integrations._base import (
    create_framework_attributes,
    get_framework_version,
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
)
from risicare.integrations._dedup import (
    suppress_provider_instrumentation,
)

logger = logging.getLogger(__name__)

# Memory guard constants
_MAX_OPEN_SPANS = 10_000
_SPAN_TTL_SECONDS = 300  # 5 minutes

# Component classification sets
_LLM_CLASSES: Set[str] = {
    "LLM", "OpenAI", "Anthropic", "Gemini", "Ollama", "Groq", "Mistral",
    "Cohere", "HuggingFaceLLM", "Together", "Bedrock", "VertexAI",
    "ChatOpenAI", "AzureOpenAI", "Replicate",
}
_EMBEDDING_CLASSES: Set[str] = {
    "BaseEmbedding", "OpenAIEmbedding", "HuggingFaceEmbedding",
    "CohereEmbedding", "GeminiEmbedding", "OllamaEmbedding",
    "AzureOpenAIEmbedding", "BedrockEmbedding",
}
_RETRIEVER_CLASSES: Set[str] = {
    "BaseRetriever", "VectorIndexRetriever", "ListIndexRetriever",
    "TreeSelectLeafRetriever", "KGTableRetriever", "QueryFusionRetriever",
    "RouterRetriever", "AutoMergingRetriever", "RecursiveRetriever",
}
_QUERY_CLASSES: Set[str] = {
    "BaseQueryEngine", "RetrieverQueryEngine", "CitationQueryEngine",
    "SubQuestionQueryEngine", "RouterQueryEngine", "SQLAutoVectorQueryEngine",
    "FLAREInstructQueryEngine",
}
_SYNTHESIZER_CLASSES: Set[str] = {
    "BaseSynthesizer", "ResponseSynthesizer", "TreeSummarize",
    "SimpleSummarize", "Refine", "CompactAndRefine",
}


def _classify_component(instance: Any) -> str:
    """
    Classify a LlamaIndex component by its class name.

    Returns one of: "llm", "embedding", "retrieval", "query", "synthesis", "component".
    """
    if instance is None:
        return "component"

    cls_name = type(instance).__name__

    # Check MRO class names too (for inheritance)
    all_names = {cls_name}
    try:
        for base in type(instance).__mro__:
            all_names.add(base.__name__)
    except Exception:
        pass

    for class_set, kind in [
        (_LLM_CLASSES, "llm"),
        (_EMBEDDING_CLASSES, "embedding"),
        (_RETRIEVER_CLASSES, "retrieval"),
        (_QUERY_CLASSES, "query"),
        (_SYNTHESIZER_CLASSES, "synthesis"),
    ]:
        if all_names & class_set:
            return kind

    return "component"


def _extract_model_name(instance: Any) -> Optional[str]:
    """Extract model name from LLM or embedding instance (defensive)."""
    for attr in ("model", "model_name", "model_id", "_model_name"):
        value = getattr(instance, attr, None)
        if value is not None:
            return str(value)
    return None


class _SpanEntry:
    """Tracks an open span with metadata."""

    __slots__ = ("span", "suppression_cm", "start_time", "component_kind")

    def __init__(
        self,
        span: Any,
        suppression_cm: Any,
        start_time: float,
        component_kind: str,
    ) -> None:
        self.span = span
        self.suppression_cm = suppression_cm
        self.start_time = start_time
        self.component_kind = component_kind


class RisicareSpanHandler:
    """
    LlamaIndex span handler that creates Risicare spans.

    We don't inherit from BaseSpanHandler to avoid requiring llama-index-core
    as an import-time dependency. LlamaIndex's dispatcher checks for method
    existence (duck typing).
    """

    def __init__(self) -> None:
        self._version = get_framework_version("llama-index-core")
        self._open_spans: Dict[str, _SpanEntry] = {}
        self._lock = threading.Lock()

    @classmethod
    def class_name(cls) -> str:
        return "RisicareSpanHandler"

    def _evict_stale_entries(self) -> None:
        """Evict stale entries to prevent memory leaks."""
        if len(self._open_spans) < _MAX_OPEN_SPANS:
            return

        now = time.perf_counter()
        stale_keys = [
            k for k, v in self._open_spans.items()
            if (now - v.start_time) > _SPAN_TTL_SECONDS
        ]

        if stale_keys:
            for key in stale_keys:
                entry = self._open_spans.pop(key, None)
                if entry:
                    self._cleanup_entry(entry)
        else:
            evict_count = max(1, len(self._open_spans) // 10)
            sorted_keys = sorted(
                self._open_spans.keys(),
                key=lambda k: self._open_spans[k].start_time,
            )
            for key in sorted_keys[:evict_count]:
                entry = self._open_spans.pop(key, None)
                if entry:
                    self._cleanup_entry(entry)

    def new_span(
        self,
        id_: str,
        bound_args: Any = None,
        instance: Any = None,
        parent_span_id: Optional[str] = None,
        tags: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Optional[Any]:
        """
        Called when a new LlamaIndex span starts.

        Creates a Risicare span and tracks it in _open_spans.
        """
        if not is_tracing_enabled():
            return None

        tracer = get_tracer()
        if tracer is None:
            return None

        try:
            from risicare_core import SpanKind
        except ImportError:
            return None

        with self._lock:
            self._evict_stale_entries()

        component_kind = _classify_component(instance)
        cls_name = type(instance).__name__ if instance else "unknown"

        attrs = create_framework_attributes("llamaindex", self._version)
        attrs["framework.span_kind"] = component_kind
        attrs["framework.llamaindex.component"] = cls_name

        if tags:
            try:
                attrs["framework.llamaindex.tags"] = str(tags)[:200]
            except Exception:
                pass

        # Determine span kind and name
        if component_kind == "llm":
            model = _extract_model_name(instance)
            span_name = f"llamaindex.llm/{model or cls_name}"
            kind = SpanKind.LLM_CALL
            if model:
                attrs["gen_ai.request.model"] = model
        elif component_kind == "embedding":
            model = _extract_model_name(instance)
            span_name = f"llamaindex.embedding/{model or cls_name}"
            kind = SpanKind.LLM_CALL
            if model:
                attrs["gen_ai.request.model"] = model
        elif component_kind == "retrieval":
            span_name = f"llamaindex.retrieve/{cls_name}"
            kind = SpanKind.INTERNAL
        elif component_kind == "query":
            span_name = f"llamaindex.query/{cls_name}"
            kind = SpanKind.INTERNAL
        elif component_kind == "synthesis":
            span_name = f"llamaindex.synthesize/{cls_name}"
            kind = SpanKind.INTERNAL
        else:
            span_name = f"llamaindex.component/{cls_name}"
            kind = SpanKind.INTERNAL

        span = tracer.start_span_no_context(
            name=span_name,
            kind=kind,
            attributes=attrs,
        )

        # Suppress provider spans for LLM and embedding components
        suppression_cm = None
        if component_kind in ("llm", "embedding"):
            suppression_cm = suppress_provider_instrumentation()
            suppression_cm.__enter__()

        try:
            with self._lock:
                self._open_spans[str(id_)] = _SpanEntry(
                    span=span,
                    suppression_cm=suppression_cm,
                    start_time=time.perf_counter(),
                    component_kind=component_kind,
                )
        except BaseException:
            if suppression_cm is not None:
                suppression_cm.__exit__(None, None, None)
            raise

        return id_

    def prepare_to_exit_span(
        self,
        id_: str,
        bound_args: Any = None,
        instance: Any = None,
        result: Any = None,
        **kwargs: Any,
    ) -> None:
        """Called when a LlamaIndex span ends successfully."""
        with self._lock:
            entry = self._open_spans.pop(str(id_), None)
        if entry is None:
            return

        try:
            span = entry.span
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Extract result info
            if result is not None and should_trace_content():
                try:
                    result_str = str(result)
                    safe_set_attribute(
                        span,
                        "gen_ai.response.content",
                        truncate_content(result_str, 5000),
                    )
                except Exception:
                    pass

            # Try to extract usage from LLM results
            if entry.component_kind in ("llm", "embedding") and result is not None:
                try:
                    # LlamaIndex LLM responses may have raw_output with usage
                    raw = getattr(result, "raw", None) or result
                    usage = getattr(raw, "usage", None)
                    if usage:
                        safe_set_attribute(span, "gen_ai.usage.prompt_tokens", getattr(usage, "prompt_tokens", None))
                        safe_set_attribute(span, "gen_ai.usage.completion_tokens", getattr(usage, "completion_tokens", None))
                except Exception:
                    pass

            span.end()
        except Exception as e:
            logger.debug(f"Error ending LlamaIndex span: {e}")
        finally:
            self._cleanup_entry(entry)

    def prepare_to_drop_span(
        self,
        id_: str,
        bound_args: Any = None,
        instance: Any = None,
        err: Any = None,
        **kwargs: Any,
    ) -> None:
        """Called when a LlamaIndex span ends with an error."""
        with self._lock:
            entry = self._open_spans.pop(str(id_), None)
        if entry is None:
            return

        try:
            span = entry.span
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if err is not None:
                if isinstance(err, BaseException):
                    record_error(span, err)
                else:
                    span.set_attribute("error", True)
                    span.set_attribute("error.message", str(err)[:500])
            else:
                span.set_attribute("error", True)
                span.set_attribute("error.message", "Span dropped without error details")

            span.end()
        except Exception as e:
            logger.debug(f"Error ending dropped LlamaIndex span: {e}")
        finally:
            self._cleanup_entry(entry)

    def _cleanup_entry(self, entry: _SpanEntry) -> None:
        """Clean up suppression context."""
        if entry.suppression_cm is not None:
            try:
                entry.suppression_cm.__exit__(None, None, None)
            except Exception:
                pass


# Module-level singleton
_handler_instance: Optional[RisicareSpanHandler] = None


def register_handler() -> None:
    """Register the Risicare span handler with LlamaIndex's dispatcher."""
    global _handler_instance

    if _handler_instance is not None:
        return

    _handler_instance = RisicareSpanHandler()

    try:
        from llama_index.core.instrumentation import get_dispatcher

        dispatcher = get_dispatcher()
        dispatcher.add_span_handler(_handler_instance)
        logger.debug("Registered Risicare LlamaIndex span handler")
    except ImportError:
        logger.debug("llama_index.core.instrumentation not available")
    except Exception as e:
        logger.debug(f"Failed to register LlamaIndex span handler: {e}")
