"""
LlamaIndex integration for Risicare SDK.

Registers a BaseSpanHandler with LlamaIndex's instrumentation dispatcher
to create spans for all LlamaIndex operations (queries, retrieval,
LLM calls, embeddings).

Span Hierarchy:
    llamaindex.query/{engine}              [INTERNAL, span_kind="query"]
      llamaindex.retrieve/{retriever}       [INTERNAL, span_kind="retrieval"]
        llamaindex.embedding/{model}        [LLM_CALL]
          (provider span suppressed)
      llamaindex.synthesize/{synthesizer}   [INTERNAL, span_kind="synthesis"]
        llamaindex.llm/{model}              [LLM_CALL]
          (provider span suppressed)

Suppresses provider instrumentation for LLM/embedding spans to avoid duplicates.

Usage (automatic — zero config):
    import risicare
    risicare.init(api_key="rsk-...")
    from llama_index.core import VectorStoreIndex
    index = VectorStoreIndex.from_documents(documents)
    engine = index.as_query_engine()
    response = engine.query("What is...?")  # Traced
"""

from __future__ import annotations

import logging
import threading
from typing import Any

logger = logging.getLogger(__name__)

_instrumented = False
_lock = threading.RLock()


def instrument_llamaindex(module: Any) -> None:
    """
    Apply instrumentation to LlamaIndex module.

    Called by the import hook when `llama_index` or `llama_index.core` is imported.
    Registers a RisicareSpanHandler with the LlamaIndex dispatcher.
    """
    global _instrumented
    if _instrumented:
        return

    with _lock:
        if _instrumented:
            return
        _instrumented = True
        try:
            from risicare.integrations._base import check_version_compatibility

            check_version_compatibility("llama-index-core")

            from risicare.integrations.llamaindex._handler import register_handler

            register_handler()
            logger.debug("Instrumented LlamaIndex")
        except Exception as e:
            _instrumented = False
            logger.debug(f"Failed to instrument LlamaIndex: {e}")
