"""
LlamaIndex integration — callback handler for auto-tracing LLM and query events.

Usage:
    from syntropy import Syntropy
    from syntropy.integrations.llamaindex import SyntropyLlamaIndexHandler

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")

    # Global handler (traces everything):
    import llama_index.core
    llama_index.core.global_handler = SyntropyLlamaIndexHandler(syntropy=syn)

    # Or attach to Settings:
    from llama_index.core import Settings
    Settings.callback_manager.add_handler(SyntropyLlamaIndexHandler(syntropy=syn))
"""

from __future__ import annotations

import time
from typing import Any, Dict, List, Optional

from syntropy.client import Syntropy

# LlamaIndex callback types — import lazily to keep SDK lightweight
_CBEventType: Any = None
_BaseCallbackHandler: Any = None


def _ensure_imports() -> None:
    global _CBEventType, _BaseCallbackHandler
    if _CBEventType is not None:
        return
    try:
        from llama_index.core.callbacks import CBEventType  # pyright: ignore[reportMissingImports]
        from llama_index.core.callbacks.base_handler import BaseCallbackHandler  # pyright: ignore[reportMissingImports]
        _CBEventType = CBEventType
        _BaseCallbackHandler = BaseCallbackHandler
    except ImportError:
        try:
            from llama_index.callbacks import CBEventType  # pyright: ignore[reportMissingImports]
            from llama_index.callbacks.base_handler import BaseCallbackHandler  # pyright: ignore[reportMissingImports]
            _CBEventType = CBEventType
            _BaseCallbackHandler = BaseCallbackHandler
        except ImportError:
            raise ImportError(
                "llama-index-core is required for this integration. "
                "Install it with: pip install llama-index-core"
            )


class SyntropyLlamaIndexHandler:
    """
    LlamaIndex callback handler that traces LLM calls, queries, and retrievals.

    Works with both LlamaIndex v0.10+ (llama-index-core) and legacy versions.
    """

    def __init__(
        self,
        syntropy: Syntropy,
        agent_id: str = "",
    ):
        _ensure_imports()
        self._syntropy = syntropy
        self._agent_id = agent_id or syntropy._config.agent_id
        self._event_starts: Dict[str, dict] = {}

    # ---- Callback interface ----

    def on_event_start(
        self,
        event_type: Any,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        parent_id: str = "",
        **kwargs: Any,
    ) -> str:
        """Record the start of a callback event."""
        self._event_starts[event_id] = {
            "event_type": event_type,
            "payload": payload or {},
            "start_time": time.perf_counter(),
            "parent_id": parent_id,
        }
        return event_id

    def on_event_end(
        self,
        event_type: Any,
        payload: Optional[Dict[str, Any]] = None,
        event_id: str = "",
        **kwargs: Any,
    ) -> None:
        """Record the end of a callback event and send trace."""
        start_data = self._event_starts.pop(event_id, None)
        if not start_data:
            return

        elapsed_ms = int(
            (time.perf_counter() - start_data["start_time"]) * 1000
        )
        payload = payload or {}
        start_payload = start_data.get("payload", {})
        et = event_type

        # LLM completion events
        if _CBEventType and et == _CBEventType.LLM:
            self._trace_llm(start_payload, payload, elapsed_ms, event_id)
        # Query events
        elif _CBEventType and et == _CBEventType.QUERY:
            self._trace_query(start_payload, payload, elapsed_ms, event_id)
        # Retrieval events
        elif _CBEventType and et == _CBEventType.RETRIEVE:
            self._trace_retrieve(start_payload, payload, elapsed_ms, event_id)
        # Embedding events
        elif _CBEventType and et == _CBEventType.EMBEDDING:
            self._trace_embedding(start_payload, payload, elapsed_ms, event_id)

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        """Called when a trace group starts (e.g., a query pipeline)."""
        pass

    def end_trace(
        self,
        trace_id: Optional[str] = None,
        trace_map: Optional[Dict[str, List[str]]] = None,
    ) -> None:
        """Called when a trace group ends."""
        pass

    # ---- Internal trace helpers ----

    def _trace_llm(
        self,
        start_payload: dict,
        end_payload: dict,
        elapsed_ms: int,
        event_id: str,
    ) -> None:
        """Trace an LLM call event."""
        # Extract prompt from messages or template
        messages = start_payload.get("messages", [])
        if messages:
            prompt = messages[-1].content if hasattr(messages[-1], "content") else str(messages[-1])
        else:
            prompt = start_payload.get("template", str(start_payload))

        # Extract response
        response_obj = end_payload.get("response", None)
        output = ""
        model = "unknown"
        tokens_in = 0
        tokens_out = 0

        if response_obj:
            output = str(getattr(response_obj, "text", response_obj))
            # Try to extract model info
            raw = getattr(response_obj, "raw", None)
            if raw:
                model = getattr(raw, "model", "unknown")
                usage = getattr(raw, "usage", None)
                if usage:
                    tokens_in = getattr(usage, "prompt_tokens", 0) or getattr(
                        usage, "input_tokens", 0
                    )
                    tokens_out = getattr(usage, "completion_tokens", 0) or getattr(
                        usage, "output_tokens", 0
                    )

        # Fallback model from serialized info
        if model == "unknown":
            serialized = start_payload.get("serialized", {})
            model = serialized.get("model", serialized.get("model_name", "unknown"))

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=model,
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "llamaindex",
                "event_type": "llm",
                "event_id": event_id,
            },
        )

    def _trace_query(
        self,
        start_payload: dict,
        end_payload: dict,
        elapsed_ms: int,
        event_id: str,
    ) -> None:
        """Trace a query event."""
        query = start_payload.get("query_str", str(start_payload))
        response_obj = end_payload.get("response", None)
        output = str(response_obj) if response_obj else ""

        self._syntropy.trace(
            prompt=query,
            response=output[:2000],  # Truncate long responses
            model="query-engine",
            agent_id=self._agent_id,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "llamaindex",
                "event_type": "query",
                "event_id": event_id,
            },
        )

    def _trace_retrieve(
        self,
        start_payload: dict,
        end_payload: dict,
        elapsed_ms: int,
        event_id: str,
    ) -> None:
        """Trace a retrieval event."""
        query = start_payload.get("query_str", str(start_payload))
        nodes = end_payload.get("nodes", [])

        self._syntropy.trace(
            prompt=query,
            response=f"Retrieved {len(nodes)} nodes",
            model="retriever",
            agent_id=self._agent_id,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "llamaindex",
                "event_type": "retrieve",
                "event_id": event_id,
                "node_count": len(nodes),
                "scores": [
                    getattr(n, "score", None) for n in nodes[:5]
                ],
            },
        )

    def _trace_embedding(
        self,
        start_payload: dict,
        end_payload: dict,
        elapsed_ms: int,
        event_id: str,
    ) -> None:
        """Trace an embedding event."""
        chunks = start_payload.get("chunks", [])

        self._syntropy.trace(
            prompt=f"Embedding {len(chunks)} chunks",
            response="completed",
            model="embedding",
            agent_id=self._agent_id,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "llamaindex",
                "event_type": "embedding",
                "event_id": event_id,
                "chunk_count": len(chunks),
            },
        )
