import contextvars
import logging
from typing import Optional, Dict, Any, List, Union
import uuid

from llama_index.core.callbacks import schema as llama_index_schema
from llama_index.core.callbacks import base_handler

from opik import context_storage, tracing_runtime_config
from opik.decorator import arguments_helpers, span_creation_handler
from opik.api_objects import opik_client, span, trace

from . import event_parsing_utils

LOGGER = logging.getLogger(__name__)

INDEX_CONSTRUCTION_TRACE_NAME = "index_construction"
LLAMA_INDEX_METADATA = {"created_from": "llama_index"}

# Context variable for root trace/span created by LlamaIndex
_llama_root: contextvars.ContextVar[Optional[Union[span.SpanData, trace.TraceData]]] = (
    contextvars.ContextVar("_llama_root", default=None)
)


def _get_last_event(trace_map: Dict[str, List[str]]) -> str:
    def dfs(key: str) -> str:
        if key not in trace_map or not trace_map[key]:
            return key
        return dfs(trace_map[key][-1])

    start_key = next(iter(trace_map))
    return dfs(start_key)


class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
    def __init__(
        self,
        event_starts_to_ignore: Optional[List[llama_index_schema.CBEventType]] = None,
        event_ends_to_ignore: Optional[List[llama_index_schema.CBEventType]] = None,
        project_name: Optional[str] = None,
        skip_index_construction_trace: bool = False,
    ):
        """Initialize LlamaIndex callback handler for Opik tracing.

        Args:
            event_starts_to_ignore: Event start types to be ignored during processing.
            event_ends_to_ignore: Event end types to be ignored during processing.
            project_name: Project name for trace/span context.
            skip_index_construction_trace: Whether to skip index construction traces.
        """
        event_starts_to_ignore = (
            event_starts_to_ignore if event_starts_to_ignore else []
        )
        event_ends_to_ignore = event_ends_to_ignore if event_ends_to_ignore else []
        super().__init__(
            event_starts_to_ignore=event_starts_to_ignore,
            event_ends_to_ignore=event_ends_to_ignore,
        )

        self._skip_index_construction_trace = skip_index_construction_trace
        self._project_name = project_name
        self._opik_client = opik_client.get_client_cached()
        self._opik_context_storage = context_storage.get_current_context_instance()

        # Event tracking - shared across contexts, but events have unique IDs
        self._map_event_id_to_span_data: Dict[str, span.SpanData] = {}
        self._map_event_id_to_output: Dict[str, Any] = {}

        # For streaming: end_trace may be called before event_end, so we need to
        # defer the trace output update until the event output is available
        self._pending_root_output_updates: Dict[
            str, Union[span.SpanData, trace.TraceData]
        ] = {}

    def _send_root_to_backend(
        self, root: Union[span.SpanData, trace.TraceData]
    ) -> None:
        """Send root trace or span data to the backend."""
        if isinstance(root, span.SpanData):
            self._opik_client.__internal_api__span__(**root.as_parameters)
        elif isinstance(root, trace.TraceData):
            self._opik_client.__internal_api__trace__(**root.as_parameters)

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        if (
            self._skip_index_construction_trace
            and trace_id == INDEX_CONSTRUCTION_TRACE_NAME
        ):
            return

        trace_name = trace_id if trace_id else "llama_index_operation"

        span_creation_result = span_creation_handler.create_span_respecting_context(
            start_span_arguments=arguments_helpers.StartSpanParameters(
                name=trace_name,
                type="general",
                project_name=self._project_name,
                metadata=LLAMA_INDEX_METADATA,
            ),
            distributed_trace_headers=None,
            opik_context_storage=self._opik_context_storage,
        )

        if span_creation_result.trace_data is not None:
            self._opik_context_storage.set_trace_data(span_creation_result.trace_data)
            self._opik_client.__internal_api__trace__(
                **span_creation_result.trace_data.as_start_parameters
            )
            _llama_root.set(span_creation_result.trace_data)
        else:
            self._opik_context_storage.add_span_data(span_creation_result.span_data)
            self._opik_client.__internal_api__span__(
                **span_creation_result.span_data.as_start_parameters
            )
            _llama_root.set(span_creation_result.span_data)

    def end_trace(
        self,
        trace_id: Optional[str] = None,
        trace_map: Optional[Dict[str, List[str]]] = None,
    ) -> None:
        if not trace_map:
            return

        root = _llama_root.get()
        if root is None:
            return

        last_event = _get_last_event(trace_map)

        # Check if the output for the last event is already available.
        # For streaming calls, LlamaIndex calls end_trace() BEFORE event_end(),
        # so the output won't be stored yet.
        if last_event in self._map_event_id_to_output:
            last_event_output = self._map_event_id_to_output.get(last_event)
            root.init_end_time().update(output=last_event_output)

            # Send the trace/span with output
            self._send_root_to_backend(root)
        else:
            # Output not available yet (streaming scenario).
            # Store the root so we can update it when event_end is called.
            # Don't send the trace/span yet - it will be sent in on_event_end
            # with the output and correct end_time to avoid race conditions.
            # Note: We don't set end_time here because the actual end is when
            # the last event ends, not when LlamaIndex calls end_trace().
            self._pending_root_output_updates[last_event] = root

        # Clean up context storage
        if isinstance(root, span.SpanData):
            self._opik_context_storage.pop_span_data(ensure_id=root.id)
        elif isinstance(root, trace.TraceData):
            self._opik_context_storage.pop_trace_data(ensure_id=root.id)

        # Clean up
        _llama_root.set(None)

    def on_event_start(
        self,
        event_type: llama_index_schema.CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: Optional[str] = None,
        parent_id: Optional[str] = None,
        **kwargs: Any,
    ) -> str:
        if not event_id:
            event_id = str(uuid.uuid4())

        root_span_or_trace = _llama_root.get()

        if root_span_or_trace is None:
            if not self._skip_index_construction_trace:
                LOGGER.warning(
                    "No active LlamaIndex trace/span found in context. "
                    "parent_id=%s, event_type=%s, event_id=%s",
                    parent_id,
                    event_type,
                    event_id,
                )
            return event_id

        span_input = event_parsing_utils.get_span_input_from_events(event_type, payload)

        # Skip creating span if event duplicates root operation name
        root_name = root_span_or_trace.name if root_span_or_trace else None
        event_duplicates_root = (
            parent_id == llama_index_schema.BASE_TRACE_EVENT
            and event_type.value == root_name
        )
        if event_duplicates_root:
            if span_input:
                root_span_or_trace.update(input=span_input)
            return event_id

        span_creation_result = span_creation_handler.create_span_respecting_context(
            start_span_arguments=arguments_helpers.StartSpanParameters(
                name=event_type.value,
                input=span_input,
                type=(
                    "llm"
                    if event_type == llama_index_schema.CBEventType.LLM
                    else "general"
                ),
                project_name=self._project_name,
                metadata=LLAMA_INDEX_METADATA,
            ),
            distributed_trace_headers=None,
            opik_context_storage=self._opik_context_storage,
        )
        span_data = span_creation_result.span_data
        self._map_event_id_to_span_data[event_id] = span_data
        self._opik_context_storage.add_span_data(span_data)

        if (
            self._opik_client.config.log_start_trace_span
            and tracing_runtime_config.is_tracing_active()
        ):
            self._opik_client.__internal_api__span__(**span_data.as_start_parameters)

        # Update root input from first child event
        if parent_id == llama_index_schema.BASE_TRACE_EVENT and span_input is not None:
            root_span_or_trace.update(input=span_input)

        return event_id

    def on_event_end(
        self,
        event_type: llama_index_schema.CBEventType,
        payload: Optional[Dict[str, Any]] = None,
        event_id: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        span_output = event_parsing_utils.get_span_output_from_event(
            event_type, payload
        )
        error_info = event_parsing_utils.get_span_error_info(payload)

        if not event_id:
            return

        # Store output for end_trace
        self._map_event_id_to_output[event_id] = span_output

        # Check if there's a pending root trace/span output update for this event.
        # This happens when end_trace() was called before event_end() (streaming scenario).
        if event_id in self._pending_root_output_updates:
            root = self._pending_root_output_updates.pop(event_id)
            # Set end_time now (the actual end) and update with output
            root.init_end_time().update(output=span_output)

            # Send the trace/span to the backend with correct end_time and output
            self._send_root_to_backend(root)

        # Finalize span if it exists
        if event_id in self._map_event_id_to_span_data:
            span_data = self._map_event_id_to_span_data[event_id]

            llm_usage_info = event_parsing_utils.get_usage_data(payload)
            span_data.update(**llm_usage_info.__dict__)
            span_data.update(output=span_output, error_info=error_info).init_end_time()

            if tracing_runtime_config.is_tracing_active():
                self._opik_client.__internal_api__span__(**span_data.as_parameters)

            self._opik_context_storage.pop_span_data(ensure_id=span_data.id)
            del self._map_event_id_to_span_data[event_id]

    def flush(self) -> None:
        """Flush pending Opik data to backend."""
        self._opik_client.flush()
