"""
LlamaIndex integration for Amplitude AI SDK.

Provides :class:`AmplitudeLlamaIndexHandler` — a LlamaIndex callback
handler that automatically tracks LLM calls, tool/function calls, and
embeddings in Amplitude.

Usage::

    from llama_index.core import Settings
    from llama_index.core.callbacks import CallbackManager
    from amplitude import Amplitude
    from amplitude_ai.integrations.llamaindex import AmplitudeLlamaIndexHandler

    amplitude = Amplitude("YOUR_API_KEY")
    handler = AmplitudeLlamaIndexHandler(
        amplitude=amplitude,
        user_id="user-1",
    )

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

from __future__ import annotations

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

try:
    from llama_index.core.callbacks import CBEventType, EventPayload
    from llama_index.core.callbacks.base_handler import BaseCallbackHandler

    LLAMAINDEX_AVAILABLE = True
except ImportError:
    try:
        # Older LlamaIndex (< 0.10)
        from llama_index.callbacks import CBEventType, EventPayload  # type: ignore[no-redef]
        from llama_index.callbacks.base import BaseCallbackHandler  # type: ignore[no-redef]

        LLAMAINDEX_AVAILABLE = True
    except ImportError:
        LLAMAINDEX_AVAILABLE = False
        BaseCallbackHandler = object  # type: ignore[assignment,misc]
        CBEventType = None  # type: ignore[assignment,misc]
        EventPayload = None  # type: ignore[assignment,misc]

from amplitude import Amplitude

from ..context import get_active_context
from ..core.privacy import PrivacyConfig
from ..core.tracking import (
    track_ai_message,
    track_embedding,
    track_tool_call,
    track_user_message,
)
from ..utils.costs import calculate_cost, infer_provider
from ..utils.logger import get_logger


class AmplitudeLlamaIndexHandler(BaseCallbackHandler):  # type: ignore[misc]
    """LlamaIndex callback handler that tracks LLM usage in Amplitude.

    Captures ``LLM``, ``FUNCTION_CALL``, and ``EMBEDDING`` events and
    emits the standard Amplitude AI event taxonomy.
    """

    def __init__(
        self,
        amplitude: Amplitude,
        user_id: str,
        privacy_config: Optional[PrivacyConfig] = None,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        parent_agent_id: Optional[str] = None,
        customer_org_id: Optional[str] = None,
        agent_version: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
        env: Optional[str] = None,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
    ) -> None:
        # Auto-inherit from SessionContext for any unset fields
        ctx = get_active_context()

        self.amplitude = amplitude
        self.user_id = user_id or (ctx.user_id if ctx else None) or user_id
        self.privacy_config = privacy_config or PrivacyConfig()
        self.session_id = session_id or (ctx.session_id if ctx else None) or str(uuid.uuid4())
        self.trace_id = trace_id or (ctx.trace_id if ctx else None) or str(uuid.uuid4())
        self.agent_id = agent_id or (ctx.agent_id if ctx else None)
        self.parent_agent_id = parent_agent_id or (ctx.parent_agent_id if ctx else None)
        self.customer_org_id = customer_org_id or (ctx.customer_org_id if ctx else None)
        self.agent_version = agent_version or (ctx.agent_version if ctx else None)
        self.context = context or (ctx.context if ctx else None)
        self.env = env or (ctx.env if ctx else None)
        self.event_properties = event_properties
        self.user_properties = user_properties
        self.groups = groups or (ctx.groups if ctx else None)

        self._turn_counter = 1

        # In-flight event state keyed by event_id
        self._event_starts: Dict[str, Dict[str, Any]] = {}

        # Required by BaseCallbackHandler (when it's available).
        # When LlamaIndex isn't installed, BaseCallbackHandler is `object`
        # and object.__init__() doesn't accept kwargs.
        if LLAMAINDEX_AVAILABLE:
            super().__init__(
                event_starts_to_ignore=[],
                event_ends_to_ignore=[],
            )

    # ------------------------------------------------------------------
    # BaseCallbackHandler 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 time and payload of an event."""
        if not event_id:
            event_id = str(uuid.uuid4())

        self._event_starts[event_id] = {
            "event_type": event_type,
            "payload": payload or {},
            "start_time": time.time(),
            "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:
        """Process the completed event and emit Amplitude tracking."""
        start_info = self._event_starts.pop(event_id, None)
        start_time = start_info["start_time"] if start_info else time.time()
        latency_ms = (time.time() - start_time) * 1000
        start_payload = start_info.get("payload", {}) if start_info else {}
        end_payload = payload or {}

        _logger = get_logger(self.amplitude)

        try:
            if not LLAMAINDEX_AVAILABLE:
                return

            if event_type == CBEventType.LLM:
                self._handle_llm_event(start_payload, end_payload, latency_ms)
            elif event_type == CBEventType.FUNCTION_CALL:
                self._handle_function_call(start_payload, end_payload, latency_ms)
            elif event_type == CBEventType.EMBEDDING:
                self._handle_embedding(start_payload, end_payload, latency_ms)
        except Exception as exc:
            _logger.warning(f"LlamaIndex tracking failed: {exc}")

    def start_trace(self, trace_id: Optional[str] = None) -> None:
        """Called when a LlamaIndex trace starts."""

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

    # ------------------------------------------------------------------
    # Internal handlers
    # ------------------------------------------------------------------

    def _handle_llm_event(
        self,
        start_payload: Dict[str, Any],
        end_payload: Dict[str, Any],
        latency_ms: float,
    ) -> None:
        """Track an LLM completion event."""
        # Extract model name
        serialized = end_payload.get(EventPayload.SERIALIZED, {}) if LLAMAINDEX_AVAILABLE else {}
        model_name = serialized.get("model", serialized.get("model_name", "unknown"))

        # Extract user messages from the start payload
        messages = start_payload.get(EventPayload.MESSAGES, []) if LLAMAINDEX_AVAILABLE else []
        for msg in messages:
            role = getattr(msg, "role", "user") if hasattr(msg, "role") else "user"
            content = getattr(msg, "content", str(msg)) if hasattr(msg, "content") else str(msg)
            if role == "user" and content:
                track_user_message(
                    amplitude=self.amplitude,
                    user_id=self.user_id,
                    message_content=content,
                    session_id=self.session_id,
                    trace_id=self.trace_id,
                    turn_id=self._turn_counter,
                    agent_id=self.agent_id,
                    parent_agent_id=self.parent_agent_id,
                    customer_org_id=self.customer_org_id,
                    agent_version=self.agent_version,
                    context=self.context,
                    env=self.env,
                    event_properties=self.event_properties,
                    user_properties=self.user_properties,
                    groups=self.groups,
                    privacy_config=self.privacy_config,
                )
                self._turn_counter += 1

        # Extract response
        response = end_payload.get(EventPayload.RESPONSE, None) if LLAMAINDEX_AVAILABLE else None
        response_text = ""
        if response is not None:
            response_text = getattr(response, "text", str(response)) if hasattr(response, "text") else str(response)

        # Extract token usage
        raw_usage = end_payload.get("token_usage", {})
        if not raw_usage and hasattr(response, "raw"):
            raw_resp = response.raw
            if hasattr(raw_resp, "usage"):
                usage_obj = raw_resp.usage
                raw_usage = {
                    "prompt_tokens": getattr(usage_obj, "prompt_tokens", None),
                    "completion_tokens": getattr(usage_obj, "completion_tokens", None),
                    "total_tokens": getattr(usage_obj, "total_tokens", None),
                }
                # Best-effort extraction of reasoning & cache tokens from
                # the underlying provider response (e.g. OpenAI, Anthropic).
                details = getattr(usage_obj, "completion_tokens_details", None)
                if details:
                    raw_usage["reasoning_tokens"] = getattr(details, "reasoning_tokens", None)
                raw_usage.setdefault(
                    "cache_read_input_tokens",
                    getattr(usage_obj, "cache_read_input_tokens", None),
                )
                raw_usage.setdefault(
                    "cache_creation_input_tokens",
                    getattr(usage_obj, "cache_creation_input_tokens", None),
                )

        raw_input = raw_usage.get("prompt_tokens") or raw_usage.get("input_tokens")
        output_tokens = raw_usage.get("completion_tokens") or raw_usage.get("output_tokens")
        total_tokens = raw_usage.get("total_tokens")
        reasoning_tokens = raw_usage.get("reasoning_tokens")
        cache_read = raw_usage.get("cache_read_input_tokens")
        cache_creation = raw_usage.get("cache_creation_input_tokens")

        # Normalize input_tokens to total (including cache).
        # OpenAI's prompt_tokens is already total. Anthropic's input_tokens
        # is non-cached only — add cache tokens when present.
        input_tokens = raw_input
        if raw_input is not None and (cache_read or cache_creation):
            is_anthropic_format = raw_usage.get("prompt_tokens") is None
            if is_anthropic_format:
                input_tokens = raw_input + (cache_read or 0) + (cache_creation or 0)

        # Cost
        total_cost_usd = None
        if input_tokens is not None and output_tokens is not None:
            try:
                total_cost_usd = calculate_cost(
                    model_name=model_name,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    reasoning_tokens=reasoning_tokens or 0,
                    cache_read_input_tokens=cache_read or 0,
                    cache_creation_input_tokens=cache_creation or 0,
                    default_provider=infer_provider(model_name),
                )
            except Exception:
                pass

        track_ai_message(
            amplitude=self.amplitude,
            user_id=self.user_id,
            model_name=model_name,
            provider="llamaindex",
            response_content=response_text,
            latency_ms=latency_ms,
            session_id=self.session_id,
            trace_id=self.trace_id,
            turn_id=self._turn_counter,
            total_cost_usd=total_cost_usd,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            event_properties=self.event_properties,
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
            reasoning_tokens=reasoning_tokens,
            cache_read_input_tokens=cache_read,
            cache_creation_input_tokens=cache_creation,
        )
        self._turn_counter += 1

    def _handle_function_call(
        self,
        start_payload: Dict[str, Any],
        end_payload: Dict[str, Any],
        latency_ms: float,
    ) -> None:
        """Track a tool / function call event."""
        tool_name = start_payload.get("tool", {}).get("name", "unknown")
        tool_input = start_payload.get("tool", {}).get("input", None)
        tool_output = end_payload.get("tool_output", end_payload.get("function_call_response"))

        # Detect tool failures: LlamaIndex may signal errors via an "error"
        # key or an exception stored in the end payload.
        error = end_payload.get("error")
        success = error is None

        track_tool_call(
            amplitude=self.amplitude,
            user_id=self.user_id,
            tool_name=tool_name,
            success=success,
            latency_ms=latency_ms,
            session_id=self.session_id,
            trace_id=self.trace_id,
            turn_id=self._turn_counter,
            tool_input=tool_input,
            tool_output=tool_output,
            error_message=str(error) if error else None,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            event_properties=self.event_properties,
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

    def _handle_embedding(
        self,
        start_payload: Dict[str, Any],
        end_payload: Dict[str, Any],
        latency_ms: float,
    ) -> None:
        """Track an embedding event."""
        serialized = end_payload.get("serialized", start_payload.get("serialized", {}))
        model_name = serialized.get("model_name", serialized.get("model", "unknown"))

        track_embedding(
            amplitude=self.amplitude,
            user_id=self.user_id,
            model=model_name,
            provider="llamaindex",
            latency_ms=latency_ms,
            session_id=self.session_id,
            agent_id=self.agent_id,
            agent_version=self.agent_version,
            context=self.context,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )


def create_amplitude_llamaindex_handler(
    amplitude: Amplitude,
    user_id: str,
    agent_id: Optional[str] = None,
    parent_agent_id: Optional[str] = None,
    customer_org_id: Optional[str] = None,
    agent_version: Optional[str] = None,
    context: Optional[Dict[str, Any]] = None,
    env: Optional[str] = None,
    **kwargs: Any,
) -> AmplitudeLlamaIndexHandler:
    """Convenience factory for :class:`AmplitudeLlamaIndexHandler`."""
    return AmplitudeLlamaIndexHandler(
        amplitude=amplitude,
        user_id=user_id,
        agent_id=agent_id,
        parent_agent_id=parent_agent_id,
        customer_org_id=customer_org_id,
        agent_version=agent_version,
        context=context,
        env=env,
        **kwargs,
    )
