"""
LiteLLM callback for Risicare SDK.

Implements litellm.integrations.custom_logger.CustomLogger to create
LLM_CALL spans around all LiteLLM completion calls.

Uses a ContextVar to track the current span between log_pre_api_call
(span start) and log_success_event/log_failure_event (span end).

Provider instrumentation is suppressed during callback execution to
prevent duplicate LLM spans from the underlying SDK calls.
"""

from __future__ import annotations

import logging
import time
from contextvars import ContextVar
from typing import Any, Dict, List, Optional

from risicare.integrations._base import (
    create_framework_attributes,
    get_framework_version,
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
)
from risicare.integrations._dedup import (
    suppress_provider_instrumentation,
)

logger = logging.getLogger(__name__)

# ContextVar to track current span across callbacks
_current_span: ContextVar[Optional[Any]] = ContextVar(
    "risicare_litellm_current_span", default=None
)
# Maximum number of tracked spans (safety limit)
_MAX_TRACKED_SPANS = 10_000


class _SpanEntry:
    """Tracks a span and its suppression context."""

    __slots__ = ("span", "suppression_cm", "start_time")

    def __init__(self, span: Any, suppression_cm: Any, start_time: float) -> None:
        self.span = span
        self.suppression_cm = suppression_cm
        self.start_time = start_time


class RisicareLiteLLMLogger:
    """
    LiteLLM CustomLogger that creates Risicare spans.

    We don't inherit from litellm.integrations.custom_logger.CustomLogger
    to avoid requiring litellm as an import-time dependency. LiteLLM's
    callback system uses duck typing — it checks for method existence.
    """

    def __init__(self) -> None:
        self._version = get_framework_version("litellm")

    def log_pre_api_call(
        self, model: str, messages: Any, kwargs: Dict[str, Any]
    ) -> None:
        """Called before each LLM API call. Start a span."""
        if not is_tracing_enabled():
            return

        tracer = get_tracer()
        if tracer is None:
            return

        try:
            from risicare_core import SpanKind
        except ImportError:
            return

        provider = kwargs.get("custom_llm_provider", "unknown")
        stream = kwargs.get("stream", False)

        attrs = create_framework_attributes("litellm", self._version)
        attrs["gen_ai.system"] = provider
        attrs["gen_ai.request.model"] = model
        attrs["framework.litellm.stream"] = stream

        api_base = kwargs.get("api_base")
        if api_base:
            # Sanitize api_base: strip query params, fragments, and userinfo
            # to prevent leaking API keys or auth tokens embedded in URLs.
            try:
                from urllib.parse import urlparse, urlunparse

                parsed = urlparse(str(api_base))
                sanitized = urlunparse((
                    parsed.scheme,
                    parsed.hostname or parsed.netloc.split("@")[-1],
                    parsed.path,
                    "",  # params
                    "",  # query — may contain api_key=...
                    "",  # fragment
                ))
                attrs["framework.litellm.api_base"] = sanitized
            except Exception:
                # Fallback: just store scheme + host
                attrs["framework.litellm.api_base"] = str(api_base).split("?")[0].split("#")[0]

        if should_trace_content() and messages:
            try:
                messages_str = str(messages)
                attrs["gen_ai.request.messages"] = truncate_content(messages_str)
            except Exception:
                pass

        span = tracer.start_span_no_context(
            name=f"litellm.completion/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=attrs,
        )

        # Suppress provider instrumentation to prevent duplicate spans
        suppression_cm = suppress_provider_instrumentation()
        suppression_cm.__enter__()
        try:
            _current_span.set(
                _SpanEntry(span=span, suppression_cm=suppression_cm, start_time=time.perf_counter())
            )
        except BaseException:
            suppression_cm.__exit__(None, None, None)
            raise

    def log_success_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: Any,
        end_time: Any,
    ) -> None:
        """Called after successful API call. End the span with response data."""
        entry = _current_span.get(None)
        if entry is None:
            return

        try:
            span = entry.span

            # Token usage from response
            usage = getattr(response_obj, "usage", None)
            if usage is not None:
                safe_set_attribute(span, "gen_ai.usage.prompt_tokens", getattr(usage, "prompt_tokens", None))
                safe_set_attribute(span, "gen_ai.usage.completion_tokens", getattr(usage, "completion_tokens", None))
                total = getattr(usage, "total_tokens", None)
                if total is not None:
                    safe_set_attribute(span, "gen_ai.usage.total_tokens", total)

            # Response model
            response_model = getattr(response_obj, "model", None)
            if response_model:
                safe_set_attribute(span, "gen_ai.response.model", response_model)

            # Cost from LiteLLM's calculator
            response_cost = kwargs.get("response_cost")
            if response_cost is not None:
                safe_set_attribute(span, "gen_ai.usage.cost_usd", response_cost)

            # Latency
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Response content
            if should_trace_content():
                try:
                    choices = getattr(response_obj, "choices", None)
                    if choices and len(choices) > 0:
                        message = getattr(choices[0], "message", None)
                        content = getattr(message, "content", None) if message else None
                        if content:
                            safe_set_attribute(
                                span,
                                "gen_ai.response.content",
                                truncate_content(str(content)),
                            )
                except Exception:
                    pass

            span.end()
        except Exception as e:
            logger.debug(f"Error ending LiteLLM span: {e}")
        finally:
            self._cleanup_entry(entry)

    def log_failure_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: Any,
        end_time: Any,
    ) -> None:
        """Called after failed API call. End the span with error."""
        entry = _current_span.get(None)
        if entry is None:
            return

        try:
            span = entry.span
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            # Try to extract error info
            exception = kwargs.get("exception") or kwargs.get("original_exception")
            if exception and isinstance(exception, BaseException):
                record_error(span, exception)
            else:
                span.set_attribute("error", True)
                error_msg = str(response_obj) if response_obj else "Unknown error"
                span.set_attribute("error.message", truncate_content(error_msg, 500))

            span.end()
        except Exception as e:
            logger.debug(f"Error ending failed LiteLLM span: {e}")
        finally:
            self._cleanup_entry(entry)

    async def async_log_pre_api_call(
        self, model: str, messages: Any, kwargs: Dict[str, Any]
    ) -> None:
        """Async variant — delegates to sync version (ContextVar is async-safe)."""
        self.log_pre_api_call(model, messages, kwargs)

    async def async_log_success_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: Any,
        end_time: Any,
    ) -> None:
        """Async variant — delegates to sync version."""
        self.log_success_event(kwargs, response_obj, start_time, end_time)

    async def async_log_failure_event(
        self,
        kwargs: Dict[str, Any],
        response_obj: Any,
        start_time: Any,
        end_time: Any,
    ) -> None:
        """Async variant — delegates to sync version."""
        self.log_failure_event(kwargs, response_obj, start_time, end_time)

    def _cleanup_entry(self, entry: _SpanEntry) -> None:
        """Clean up suppression context and ContextVar."""
        try:
            if entry.suppression_cm is not None:
                entry.suppression_cm.__exit__(None, None, None)
        except Exception:
            pass
        _current_span.set(None)


# Module-level singleton
_logger_instance: Optional[RisicareLiteLLMLogger] = None


def register_callback(module: Any) -> None:
    """Register the Risicare callback with LiteLLM."""
    global _logger_instance

    if _logger_instance is not None:
        return

    _logger_instance = RisicareLiteLLMLogger()

    # LiteLLM uses litellm.callbacks list
    callbacks = getattr(module, "callbacks", None)
    if callbacks is None:
        module.callbacks = [_logger_instance]
    elif isinstance(callbacks, list):
        callbacks.append(_logger_instance)
    else:
        module.callbacks = [_logger_instance]

    logger.debug("Registered Risicare LiteLLM callback")
