"""LlamaIndex integration for AgentTrace.

Usage::

    from agenttrace.llamaindex import AgentTraceLlamaIndexHandler
    from llama_index.core import Settings
    from llama_index.core.callbacks import CallbackManager

    handler = AgentTraceLlamaIndexHandler()
    Settings.callback_manager = CallbackManager([handler])
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

from agenttrace.langchain import _safe_json
from agenttrace.trace import trace

# LlamaIndex CBEventType string values we map to agenttrace SpanType
_EVENT_TYPE_MAP: dict[str, str] = {
    "llm": "llm",
    "embedding": "llm",
    "query": "chain",
    "retrieve": "retrieval",
    "synthesize": "chain",
    "sub_question": "chain",
    "templating": "chain",
    "function_call": "tool",
    "agent_step": "agent",
}


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _safe_payload(value: Any) -> Any:
    if value is None:
        return {}
    return _safe_json(value)


class _PendingEvent:
    __slots__ = ("span_type", "name", "start_time", "payload")

    def __init__(self, span_type: str, name: str, start_time: datetime, payload: Any) -> None:
        self.span_type = span_type
        self.name = name
        self.start_time = start_time
        self.payload = payload


class AgentTraceLlamaIndexHandler:
    """LlamaIndex BaseCallbackHandler-compatible handler for AgentTrace.

    Implements the on_event_start / on_event_end protocol without requiring
    llama_index as a hard dependency.
    """

    event_starts_to_ignore: list[str] = []
    event_ends_to_ignore: list[str] = []

    def __init__(self) -> None:
        self._pending: dict[str, _PendingEvent] = {}

    def on_event_start(
        self,
        event_type: str,
        payload: dict[str, Any] | None = None,
        event_id: str = "",
        parent_id: str = "",
        **kwargs: Any,
    ) -> str:
        """Called when a LlamaIndex event starts."""
        if not event_id:
            event_id = str(uuid4())

        event_key = event_type.lower() if isinstance(event_type, str) else str(event_type).split(".")[-1].lower()
        span_type = _EVENT_TYPE_MAP.get(event_key, "chain")

        self._pending[event_id] = _PendingEvent(
            span_type=span_type,
            name=f"llamaindex_{event_key}",
            start_time=_now_utc(),
            payload=_safe_payload(payload),
        )
        return event_id

    def on_event_end(
        self,
        event_type: str,
        payload: dict[str, Any] | None = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> None:
        """Called when a LlamaIndex event ends."""
        ctx = trace.current()
        if ctx is None:
            return

        pending = self._pending.pop(event_id, None)
        end_time = _now_utc()

        if pending is None:
            event_key = event_type.lower() if isinstance(event_type, str) else str(event_type).split(".")[-1].lower()
            span_type = _EVENT_TYPE_MAP.get(event_key, "chain")
            pending = _PendingEvent(
                span_type=span_type,
                name=f"llamaindex_{event_key}",
                start_time=end_time,
                payload={},
            )

        ctx.add_span(
            span_type=pending.span_type,  # type: ignore[arg-type]
            name=pending.name,
            start_time=pending.start_time,
            end_time=end_time,
            input_json=pending.payload,
            output_json=_safe_payload(payload),
            metadata_json={"source": "llamaindex"},
        )

    def start_trace(self, trace_id: str | None = None) -> None:
        """Called by LlamaIndex CallbackManager at trace start. No-op."""

    def end_trace(
        self,
        trace_id: str | None = None,
        trace_map: dict[str, list[str]] | None = None,
    ) -> None:
        """Called by LlamaIndex CallbackManager at trace end. No-op."""
