"""LlamaIndex auto-instrumentation.

Patches ``llama_index.core.llms.LLM.chat`` / ``achat`` and
``llama_index.core.base.query_engine.BaseQueryEngine.query`` / ``aquery``
to capture model_call and retrieval spans.
"""
from __future__ import annotations

import functools
import logging
from datetime import datetime, timezone
from typing import Any

from fibinaci.instrumentation._base import BaseInstrumentor
from fibinaci.instrumentation._wrapper import ProviderConfig, _make_error_span

logger = logging.getLogger("fibinaci.instrumentation")

_CONFIG = ProviderConfig(
    name="llamaindex",
    input_tokens_field="prompt_tokens",
    output_tokens_field="completion_tokens",
    finish_reason_getter=lambda r: None,
)


def _extract_llm_usage(response: Any, self_ref: Any) -> tuple[int | None, int | None, str | None]:
    """Extract token counts and model from a LlamaIndex ChatResponse."""
    raw = getattr(response, "raw", None) or {}
    usage = raw.get("usage") if isinstance(raw, dict) else getattr(raw, "usage", None)

    inp = out = None
    if usage:
        inp = getattr(usage, "prompt_tokens", None) or (usage.get("prompt_tokens") if isinstance(usage, dict) else None)
        out = getattr(usage, "completion_tokens", None) or (usage.get("completion_tokens") if isinstance(usage, dict) else None)

    # Additional metadata fallback
    additional = getattr(response, "additional_kwargs", {}) or {}
    if inp is None:
        inp = additional.get("prompt_tokens")
    if out is None:
        out = additional.get("completion_tokens")

    model = getattr(self_ref, "model", None) or getattr(self_ref, "model_name", None)
    return inp, out, model


def _wrap_chat(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                inp, out, model = _extract_llm_usage(response, self)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="llamaindex",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model, inp, out),
                )
                sd.set_model_call(model_name=model)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture llamaindex span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model", None) or getattr(self, "model_name", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture llamaindex error span", exc_info=True)
            raise

    return wrapper


def _wrap_achat(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData, estimate_cost

                inp, out, model = _extract_llm_usage(response, self)
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="model_call",
                    provider="llamaindex",
                    model_type="llm",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                    tokens_input=inp,
                    tokens_output=out,
                    tokens_total=(inp or 0) + (out or 0) if inp else None,
                    cost_usd=estimate_cost(model, inp, out),
                )
                sd.set_model_call(model_name=model)
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async llamaindex span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                model = getattr(self, "model", None) or getattr(self, "model_name", None)
                get_buffer().push(_make_error_span(_CONFIG, {"model": model}, exc, started_at, ended_at))
            except Exception:
                logger.debug("Failed to capture async llamaindex error span", exc_info=True)
            raise

    return wrapper


def _wrap_query(original: Any) -> Any:
    @functools.wraps(original)
    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                query_str = args[0] if args else kwargs.get("str_or_query_bundle", "")
                source_nodes = getattr(response, "source_nodes", []) or []
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="retrieval",
                    provider="llamaindex",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                )
                sd.retrieval = {
                    "query": str(query_str),
                    "results_count": len(source_nodes),
                }
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture llamaindex query span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="retrieval",
                    provider="llamaindex",
                    started_at=started_at,
                    ended_at=datetime.now(timezone.utc),
                    status="error",
                    error_type=type(exc).__name__,
                    error_message=str(exc),
                )
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture llamaindex query error span", exc_info=True)
            raise

    return wrapper


def _wrap_aquery(original: Any) -> Any:
    @functools.wraps(original)
    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        started_at = datetime.now(timezone.utc)
        try:
            response = await original(self, *args, **kwargs)
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                query_str = args[0] if args else kwargs.get("str_or_query_bundle", "")
                source_nodes = getattr(response, "source_nodes", []) or []
                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="retrieval",
                    provider="llamaindex",
                    started_at=started_at,
                    ended_at=ended_at,
                    duration_ms=(ended_at - started_at).total_seconds() * 1000,
                    status="ok",
                )
                sd.retrieval = {
                    "query": str(query_str),
                    "results_count": len(source_nodes),
                }
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async llamaindex query span", exc_info=True)
            return response
        except Exception as exc:
            ended_at = datetime.now(timezone.utc)
            try:
                from fibinaci._globals import get_buffer
                from fibinaci.tracing.context import _new_span_id, _new_trace_id, get_current_trace_id
                from fibinaci.tracing.span_builder import SpanData

                sd = SpanData(
                    trace_id=get_current_trace_id() or _new_trace_id(),
                    span_id=_new_span_id(),
                    span_kind="retrieval",
                    provider="llamaindex",
                    started_at=started_at,
                    ended_at=datetime.now(timezone.utc),
                    status="error",
                    error_type=type(exc).__name__,
                    error_message=str(exc),
                )
                get_buffer().push(sd.to_dict())
            except Exception:
                logger.debug("Failed to capture async llamaindex query error span", exc_info=True)
            raise

    return wrapper


class LlamaIndexInstrumentor(BaseInstrumentor):
    def __init__(self) -> None:
        self._originals: dict[str, Any] = {}

    def instrument(self) -> None:
        from llama_index.core.llms import LLM

        self._originals["llm_chat"] = LLM.chat
        LLM.chat = _wrap_chat(LLM.chat)

        self._originals["llm_achat"] = LLM.achat
        LLM.achat = _wrap_achat(LLM.achat)

        try:
            from llama_index.core.base.base_query_engine import BaseQueryEngine

            self._originals["query"] = BaseQueryEngine.query
            BaseQueryEngine.query = _wrap_query(BaseQueryEngine.query)

            self._originals["aquery"] = BaseQueryEngine.aquery
            BaseQueryEngine.aquery = _wrap_aquery(BaseQueryEngine.aquery)
        except ImportError:
            logger.debug("llama_index query engine not available, skipping")

    def uninstrument(self) -> None:
        if not self._originals:
            return
        from llama_index.core.llms import LLM

        LLM.chat = self._originals["llm_chat"]
        LLM.achat = self._originals["llm_achat"]

        try:
            from llama_index.core.base.base_query_engine import BaseQueryEngine

            if "query" in self._originals:
                BaseQueryEngine.query = self._originals["query"]
                BaseQueryEngine.aquery = self._originals["aquery"]
        except ImportError:
            pass
        self._originals.clear()
