"""LiteLLM integration using Beacon SDK Spans.

This module provides a callback handler that traces LiteLLM operations
using Beacon Span objects, exported via the BeaconClient's configured
TracerProvider (respecting isolated mode).

Usage:
    from lumenova_beacon import BeaconClient
    from lumenova_beacon import BeaconLiteLLMLogger
    import litellm

    # Initialize BeaconClient (sets up TracerProvider)
    client = BeaconClient()

    # Register logger
    litellm.callbacks = [BeaconLiteLLMLogger()]

    # All LiteLLM calls now traced via Beacon spans
    response = litellm.completion(
        model="gpt-4",
        messages=[{"role": "user", "content": "Hello"}],
        metadata={
            "generation_name": "greeting",
            "session_id": "user-123",
            "tags": ["prod"]
        }
    )

Environment Variables:
    BEACON_ENDPOINT - Beacon API endpoint (for OTLP export)
    BEACON_API_KEY - Beacon API key (for OTLP export)
    BEACON_SESSION_ID - Default session ID (optional)
"""

import json
import logging
import time
from datetime import datetime
from typing import Any

from lumenova_beacon.core.client import get_client
from lumenova_beacon.tracing.span import Span
from lumenova_beacon.types import SpanType, SpanKind, StatusCode

try:
    from litellm.integrations.custom_logger import CustomLogger
except ImportError:
    raise ImportError(
        "litellm is required for the LiteLLM integration. "
        "Install it with: pip install 'lumenova-beacon[litellm]'"
    )

logger = logging.getLogger(__name__)


class BeaconLiteLLMLogger(CustomLogger):
    """Beacon SDK logger for tracing LiteLLM operations.

    Spans are exported via the BeaconClient's configured TracerProvider,
    supporting both global and isolated modes.

    This logger captures:
    - Input messages and parameters
    - Response content
    - Token usage
    - Request timing
    - Errors and exceptions
    - Custom metadata (generation names, trace IDs, tags, etc.)

    Example:
        >>> from lumenova_beacon import BeaconClient
        >>> from lumenova_beacon import BeaconLiteLLMLogger
        >>> import litellm
        >>>
        >>> # Initialize BeaconClient (sets up TracerProvider)
        >>> client = BeaconClient()
        >>>
        >>> # All calls automatically traced
        >>> response = litellm.completion(
        ...     model="gpt-4",
        ...     messages=[{"role": "user", "content": "Hello"}]
        ... )
    """

    # GenAI semantic conventions
    # See: https://opentelemetry.io/docs/specs/semconv/gen-ai/
    ATTR_LLM_MODEL = "gen_ai.request.model"
    ATTR_LLM_TEMPERATURE = "gen_ai.request.temperature"
    ATTR_LLM_MAX_TOKENS = "gen_ai.request.max_tokens"
    ATTR_LLM_TOP_P = "gen_ai.request.top_p"
    ATTR_LLM_TOP_K = "gen_ai.request.top_k"
    ATTR_LLM_FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty"
    ATTR_LLM_PRESENCE_PENALTY = "gen_ai.request.presence_penalty"
    ATTR_LLM_INPUT_MESSAGES = "gen_ai.input.messages"
    ATTR_LLM_OUTPUT_MESSAGES = "gen_ai.output.messages"
    ATTR_LLM_PROVIDER_NAME = "gen_ai.provider.name"
    ATTR_LLM_INPUT_TOKENS = "gen_ai.usage.input_tokens"
    ATTR_LLM_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
    ATTR_OPERATION_NAME = "gen_ai.operation.name"
    ATTR_RESPONSE_MODEL = "gen_ai.response.model"
    ATTR_RESPONSE_ID = "gen_ai.response.id"
    ATTR_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
    ATTR_REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences"
    ATTR_REQUEST_SEED = "gen_ai.request.seed"
    ATTR_CONVERSATION_ID = "gen_ai.conversation.id"
    ATTR_TTFT_MS = "gen_ai.response.time_to_first_token_ms"
    ATTR_STREAMING = "gen_ai.request.streaming"

    def __init__(self):
        """Initialize the Beacon LiteLLM logger."""
        super().__init__()
        self._client = get_client()

        # TTFT tracking for streaming calls
        self._first_chunk_times: dict[str, float] = {}
        self._start_perf_times: dict[str, float] = {}

        logger.debug("Initialized BeaconLiteLLMLogger")

    def _extract_provider_from_model(self, model: str) -> str:
        """Extract provider from model name (e.g., 'openai/gpt-4' -> 'openai')."""
        if "/" in model:
            return model.split("/")[0].lower()
        model_lower = model.lower()
        if model_lower.startswith(("gpt-", "o1-", "o3-", "chatgpt")):
            return "openai"
        if model_lower.startswith(("claude",)):
            return "anthropic"
        if model_lower.startswith(("gemini",)):
            return "google"
        if model_lower.startswith(("mistral", "mixtral")):
            return "mistral"
        if model_lower.startswith(("command",)):
            return "cohere"
        return "litellm"

    def _extract_metadata(self, kwargs: dict) -> dict:
        """Extract Beacon-compatible metadata from LiteLLM kwargs."""
        litellm_params = kwargs.get("litellm_params", {})
        metadata = litellm_params.get("metadata", {})

        return {
            "generation_name": metadata.get("generation_name"),
            "trace_id": metadata.get("trace_id"),
            "session_id": metadata.get("session_id"),
            "user_id": metadata.get("user_id") or metadata.get("trace_user_id"),
            "user_email": metadata.get("user_email"),
            "user_name": metadata.get("user_name"),
            "tags": metadata.get("tags", []),
            "custom": {
                k: v for k, v in metadata.items()
                if k not in ["generation_name", "trace_id", "session_id",
                             "user_id", "trace_user_id", "user_email",
                             "user_name", "tags"]
            }
        }

    def _extract_messages(self, kwargs: dict) -> list[dict]:
        """Extract messages in a clean format."""
        messages = kwargs.get("messages", [])

        clean_messages = []
        for msg in messages:
            if isinstance(msg, dict):
                clean_messages.append(msg)
            else:
                clean_messages.append({
                    "role": getattr(msg, "role", "unknown"),
                    "content": getattr(msg, "content", str(msg))
                })

        return clean_messages

    def _extract_response_content(self, response_obj: Any) -> dict:
        """Extract response content in a clean format."""
        try:
            if hasattr(response_obj, "choices") and response_obj.choices:
                choice = response_obj.choices[0]

                if hasattr(choice, "message"):
                    message = choice.message
                    result = {
                        "role": "assistant",
                        "content": getattr(message, "content", None)
                    }

                    if hasattr(message, "tool_calls") and message.tool_calls:
                        result["tool_calls"] = [
                            {
                                "id": tc.id,
                                "type": getattr(tc, "type", "function"),
                                "function": {
                                    "name": tc.function.name,
                                    "arguments": tc.function.arguments
                                }
                            }
                            for tc in message.tool_calls
                        ]

                    if hasattr(message, "function_call") and message.function_call:
                        result["function_call"] = {
                            "name": message.function_call.name,
                            "arguments": message.function_call.arguments
                        }

                    return result

                elif hasattr(choice, "text"):
                    return {"text": choice.text}

            if hasattr(response_obj, "data"):
                return {"data": "embedding_response"}

        except Exception as e:
            logger.debug(f"Error extracting response content: {e}")

        return {"content": str(response_obj)[:500]}

    def _set_span_attributes(
        self,
        span: Span,
        kwargs: dict,
        response_obj: Any,
        metadata: dict,
        is_error: bool = False
    ) -> None:
        """Set all attributes on the Beacon span."""
        model = kwargs.get("model", "unknown")
        provider = self._extract_provider_from_model(model)
        span.set_attribute(self.ATTR_OPERATION_NAME, "chat")
        span.set_attribute(self.ATTR_LLM_PROVIDER_NAME, provider)
        span.set_attribute(self.ATTR_LLM_MODEL, model)

        # Input messages
        messages = self._extract_messages(kwargs)
        span.set_input(messages)
        span.set_attribute(self.ATTR_LLM_INPUT_MESSAGES, json.dumps(messages, default=str))

        # Model parameters
        param_mapping = {
            "temperature": self.ATTR_LLM_TEMPERATURE,
            "max_tokens": self.ATTR_LLM_MAX_TOKENS,
            "top_p": self.ATTR_LLM_TOP_P,
            "top_k": self.ATTR_LLM_TOP_K,
            "frequency_penalty": self.ATTR_LLM_FREQUENCY_PENALTY,
            "presence_penalty": self.ATTR_LLM_PRESENCE_PENALTY,
        }

        for param_name, attr_name in param_mapping.items():
            if param_name in kwargs and kwargs[param_name] is not None:
                span.set_attribute(attr_name, kwargs[param_name])

        # Stop sequences and seed
        stop = kwargs.get("stop")
        if stop is not None:
            span.set_attribute(
                self.ATTR_REQUEST_STOP_SEQUENCES,
                json.dumps(stop) if isinstance(stop, (list, tuple)) else stop,
            )
        seed = kwargs.get("seed")
        if seed is not None:
            span.set_attribute(self.ATTR_REQUEST_SEED, seed)

        # Output and usage (only for successful calls)
        if not is_error and response_obj:
            output = self._extract_response_content(response_obj)
            span.set_output(output)
            span.set_attribute(self.ATTR_LLM_OUTPUT_MESSAGES, json.dumps(output, default=str))

            actual_model = getattr(response_obj, "model", None)
            if actual_model:
                span.set_attribute(self.ATTR_RESPONSE_MODEL, actual_model)

            response_id = getattr(response_obj, "id", None)
            if response_id:
                span.set_attribute(self.ATTR_RESPONSE_ID, response_id)

            if hasattr(response_obj, "choices") and response_obj.choices:
                finish_reasons = [
                    getattr(choice, "finish_reason", None)
                    for choice in response_obj.choices
                    if getattr(choice, "finish_reason", None)
                ]
                if finish_reasons:
                    span.set_attribute(self.ATTR_RESPONSE_FINISH_REASONS, json.dumps(finish_reasons))

            if hasattr(response_obj, "usage") and response_obj.usage:
                usage = response_obj.usage
                input_tokens = getattr(usage, "prompt_tokens", 0)
                output_tokens = getattr(usage, "completion_tokens", 0)
                span.set_attribute(self.ATTR_LLM_INPUT_TOKENS, input_tokens)
                span.set_attribute(self.ATTR_LLM_OUTPUT_TOKENS, output_tokens)

        # Session ID
        session_id = metadata.get("session_id") or self._client.config.session_id
        if session_id:
            span.set_attribute(self.ATTR_CONVERSATION_ID, session_id)

        # Custom metadata
        for key, value in metadata.get("custom", {}).items():
            if isinstance(value, (str, int, float, bool)):
                attr_value = value
            else:
                attr_value = json.dumps(value, default=str)
            span.set_attribute(f"beacon.metadata.{key}", attr_value)

    def _get_call_id(self, kwargs: dict) -> str:
        """Get a unique call identifier for TTFT tracking."""
        call_id = kwargs.get("litellm_call_id")
        if call_id is not None:
            return str(call_id)
        return str(id(kwargs))

    def log_pre_api_call(
        self,
        model: str,
        messages: list,
        kwargs: dict,
    ) -> None:
        """Record high-precision start time before the API call."""
        try:
            call_id = self._get_call_id(kwargs)
            self._start_perf_times[call_id] = time.perf_counter()
        except Exception as e:
            logger.debug(f"Error in log_pre_api_call: {e}")

    def log_stream_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Log streaming chunk event for TTFT calculation."""
        try:
            call_id = self._get_call_id(kwargs)
            if call_id not in self._first_chunk_times:
                self._first_chunk_times[call_id] = time.perf_counter()
        except Exception as e:
            logger.debug(f"Error in log_stream_event: {e}")

    async def async_log_stream_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Async version of log_stream_event."""
        self.log_stream_event(kwargs, response_obj, start_time, end_time)

    def log_success_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Log successful LiteLLM completion as a Beacon span."""
        try:
            metadata = self._extract_metadata(kwargs)
            span_name = metadata.get("generation_name") or f"{kwargs.get('model', 'litellm')}.completion"
            call_id = self._get_call_id(kwargs)
            session_id = metadata.get("session_id") or self._client.config.session_id
            user_email = metadata.get("user_email") or self._client.config.user_email
            user_name = metadata.get("user_name") or self._client.config.user_name
            user_id = metadata.get("user_id") or self._client.config.user_id

            # Create Beacon span with explicit timing
            span = Span(
                name=span_name,
                span_type=SpanType.GENERATION,
                kind=SpanKind.CLIENT,
                session_id=session_id,
                user_email=user_email,
                user_name=user_name,
                user_id=user_id,
            )
            span.start(start_time=start_time)

            self._set_span_attributes(span, kwargs, response_obj, metadata)
            span.set_status(StatusCode.OK)

            # Calculate and set TTFT for streaming calls
            is_streaming = kwargs.get("stream", False)
            first_chunk_time = self._first_chunk_times.pop(call_id, None)
            start_perf_time = self._start_perf_times.pop(call_id, None)

            if is_streaming and first_chunk_time is not None and start_perf_time is not None:
                ttft_ms = (first_chunk_time - start_perf_time) * 1000
                span.set_attribute(self.ATTR_TTFT_MS, ttft_ms)
                span.set_attribute(self.ATTR_STREAMING, True)
            elif is_streaming:
                span.set_attribute(self.ATTR_STREAMING, True)

            span.end(end_time=end_time)
            self._client.export_span(span)

            logger.debug(f"Logged LiteLLM call: {span_name}")

        except Exception as e:
            logger.error(f"Error logging success event: {e}", exc_info=True)

    def log_failure_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Log failed LiteLLM completion as a Beacon span."""
        try:
            metadata = self._extract_metadata(kwargs)
            span_name = metadata.get("generation_name") or f"{kwargs.get('model', 'litellm')}.completion"
            call_id = self._get_call_id(kwargs)
            session_id = metadata.get("session_id") or self._client.config.session_id
            user_email = metadata.get("user_email") or self._client.config.user_email
            user_name = metadata.get("user_name") or self._client.config.user_name
            user_id = metadata.get("user_id") or self._client.config.user_id

            # Clean up TTFT tracking state
            self._first_chunk_times.pop(call_id, None)
            self._start_perf_times.pop(call_id, None)

            span = Span(
                name=span_name,
                span_type=SpanType.GENERATION,
                kind=SpanKind.CLIENT,
                session_id=session_id,
                user_email=user_email,
                user_name=user_name,
                user_id=user_id,
            )
            span.start(start_time=start_time)

            self._set_span_attributes(span, kwargs, response_obj, metadata, is_error=True)

            # Record error
            exception = kwargs.get("exception")
            if exception:
                span.set_attribute("error.type", type(exception).__qualname__)
                span.set_status(StatusCode.ERROR, str(exception))
            else:
                error_msg = str(response_obj)[:500] if response_obj else "Unknown error"
                span.set_attribute("error.type", type(response_obj).__qualname__ if response_obj else "Unknown")
                span.set_status(StatusCode.ERROR, error_msg)

            span.end(end_time=end_time)
            self._client.export_span(span)

            logger.debug(f"Logged failed LiteLLM call: {span_name}")

        except Exception as e:
            logger.error(f"Error logging failure event: {e}", exc_info=True)

    async def async_log_success_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Async version of log_success_event."""
        self.log_success_event(kwargs, response_obj, start_time, end_time)

    async def async_log_failure_event(
        self,
        kwargs: dict,
        response_obj: Any,
        start_time: datetime,
        end_time: datetime,
    ) -> None:
        """Async version of log_failure_event."""
        self.log_failure_event(kwargs, response_obj, start_time, end_time)
