"""LiteLLM callback integration for MCA SDK.

This module provides MCALiteLLMCallback, a CustomLogger implementation
that automatically captures GenAI telemetry when using LiteLLM.

SECURITY WARNING - Global State:
    This callback uses global registration via litellm.callbacks = [callback].
    This is a fundamental limitation of LiteLLM's callback system and may cause:
    - Metric mis-attribution in multi-tenant applications
    - Configuration conflicts with other libraries using LiteLLM callbacks
    - Unpredictable behavior when multiple MCAClient instances exist

    For production use:
    - Only register ONE callback instance per application process
    - Register callback early in application startup
    - Be aware ALL litellm.completion() calls will use this callback
    - Avoid using multiple MCAClient configurations in the same process

    Note: This is a LiteLLM design constraint that cannot be fully mitigated
    without modifications to LiteLLM itself.

LOGGING STRATEGY:
    - Normal telemetry operations: Logged via mca_client.logger (configured logger)
    - Critical callback failures: Logged to stderr (ensures visibility even if logger fails)
    - Cost calculation failures: Logged as warnings via mca_client.logger
"""

import logging
import sys
import time
from collections import OrderedDict
from datetime import datetime
from typing import Any, Dict, Optional

try:
    import litellm
    from litellm.integrations.custom_logger import CustomLogger

    _LITELLM_AVAILABLE = True
except ImportError as e:
    _LITELLM_AVAILABLE = False
    _LITELLM_IMPORT_ERROR = str(e)
    # Create a dummy CustomLogger for type checking
    CustomLogger = object  # type: ignore

from opentelemetry.trace import Status, StatusCode

from ..core.client import MCAClient

logger = logging.getLogger(__name__)

# Maximum number of active spans to prevent memory leaks
_MAX_ACTIVE_SPANS = 1000


class MCALiteLLMCallback(CustomLogger):
    """LiteLLM callback handler for MCA SDK telemetry.

    Automatically captures metrics, traces, and logs for LiteLLM
    API calls including token usage, cost, latency, and errors.

    Metrics recorded:
        - emms_genai_tokens_total (token_type=prompt|completion|total): Token usage by type (counter).
          WARNING: Aggregating without a token_type filter triple-counts — always filter by token_type.
        - emms_genai_cost_usd: Request cost in USD (histogram)
        - emms_genai_requests_total: Total requests (counter)
        - emms_genai_requests_failed_total: Failed requests (counter)
        - emms_genai_latency_seconds: Request latency (histogram)

    Example:
        >>> from mca_sdk.integrations import create_genai_client, MCALiteLLMCallback
        >>> import litellm
        >>>
        >>> client = create_genai_client(
        ...     service_name="clinical-summarizer",
        ...     llm_provider="openai",
        ...     llm_model="gpt-4",
        ...     team_name="ai-team"
        ... )
        >>>
        >>> callback = MCALiteLLMCallback(client)
        >>> litellm.callbacks = [callback]
        >>>
        >>> # Now all litellm.completion() calls are automatically instrumented
        >>> response = litellm.completion(
        ...     model="gpt-4",
        ...     messages=[{"role": "user", "content": "Hello"}]
        ... )
    """

    def __init__(
        self,
        client: MCAClient,
        record_prompts: bool = False,
        record_completions: bool = False,
        record_error_messages: bool = False,
        max_active_spans: int = 1000,
    ):
        """Initialize the LiteLLM callback handler.

        Args:
            client: MCAClient instance configured for GenAI
            record_prompts: If True, log prompts (HIPAA WARNING: may contain PHI)
            record_completions: If True, log completions (HIPAA WARNING: may contain PHI)
            record_error_messages: If True, log full error messages (SECURITY WARNING: may contain PHI)
            max_active_spans: Maximum number of active spans to prevent memory leaks (default: 1000)

        Security Notes:
            - By default, NO sensitive data is logged (prompts, completions, error messages)
            - Error messages often contain user input and should not be logged in HIPAA environments
            - Only enable recording flags in non-production or when you have explicit approval

        Raises:
            ImportError: If litellm package is not installed
        """
        if not _LITELLM_AVAILABLE:
            raise ImportError(
                f"MCALiteLLMCallback requires the 'litellm' package. "
                f"Install it with: pip install mca-sdk[genai] or pip install litellm>=1.61.15. "
                f"Import error: {_LITELLM_IMPORT_ERROR}"
            )

        super().__init__()
        self._client = client
        self._record_prompts = record_prompts
        self._record_completions = record_completions
        self._record_error_messages = record_error_messages
        self._max_active_spans = max_active_spans

        # Create metrics instruments (reused across calls)
        self._requests_counter = client.meter.create_counter(
            "emms_genai_requests_total", description="Total number of GenAI requests"
        )
        self._requests_failed_counter = client.meter.create_counter(
            "emms_genai_requests_failed_total",
            description="Total number of failed GenAI requests",
        )
        self._tokens_counter = client.meter.create_counter(
            "emms_genai_tokens_total",
            description=(
                "GenAI token usage by type (prompt, completion, total). "
                "Always filter by token_type when aggregating to avoid double-counting total."
            ),
        )
        self._cost_histogram = client.meter.create_histogram(
            "emms_genai_cost_usd", description="Request cost in USD", unit="usd"
        )
        self._latency_histogram = client.meter.create_histogram(
            "emms_genai_latency_seconds", description="Request latency in seconds", unit="s"
        )

        # Track active spans for correlation (OrderedDict for FIFO behavior)
        self._active_spans: OrderedDict[str, Any] = OrderedDict()

        # Runtime safety check for global state conflicts
        self._check_global_state_safety()

    def log_pre_api_call(self, model: str, messages: list, kwargs: dict) -> None:
        """Called before the LLM API call.

        Creates a trace span for the LLM call with initial attributes.

        Note: This method MUST NOT raise exceptions that could crash the user's application.
        All errors are caught, logged to stderr, and suppressed.
        """
        try:
            call_id = kwargs.get("litellm_call_id", str(id(kwargs)))

            # Prevent memory leak: enforce max active spans (FIFO eviction)
            if len(self._active_spans) >= self._max_active_spans:
                # Remove oldest span (FIFO from OrderedDict)
                oldest_id, oldest_data = self._active_spans.popitem(last=False)
                # End the orphaned span to prevent resource leak
                if "span" in oldest_data and oldest_data["span"]:
                    try:
                        oldest_data["span"].set_status(
                            Status(StatusCode.ERROR, "Span evicted due to max limit")
                        )
                        oldest_data["span"].end()
                    except Exception:
                        pass  # Suppress errors from cleanup

            # Extract span attributes (no PHI - just metadata)
            span_attributes = self._extract_span_attributes(model, kwargs)

            # Start trace span
            span = self._client.tracer.start_span(name=f"llm.{model}", attributes=span_attributes)

            # Store span for later completion
            self._active_spans[call_id] = {
                "span": span,
                "start_time": time.time(),
                "model": model,
            }
        except Exception as e:
            # CRITICAL: Never crash the user's application
            print(f"MCALiteLLMCallback.log_pre_api_call failed: {e}", file=sys.stderr)
            # Suppress exception - monitoring failure should not break the app

    def log_post_api_call(
        self, kwargs: dict, response_obj: Any, start_time: datetime, end_time: datetime
    ) -> None:
        """Called after the LLM API call (regardless of success/failure).

        This is called before log_success_event or log_failure_event.
        """
        # Post-call processing handled in success/failure handlers
        pass

    def log_success_event(
        self, kwargs: dict, response_obj: Any, start_time: datetime, end_time: datetime
    ) -> None:
        """Called on successful LLM API call.

        Records token metrics, cost, latency, and completes the trace span.

        Note: This method MUST NOT raise exceptions that could crash the user's application.
        All errors are caught, logged to stderr, and suppressed.
        """
        try:
            call_id = kwargs.get("litellm_call_id", str(id(kwargs)))
            model = kwargs.get("model", "unknown")

            # Calculate latency
            latency = (end_time - start_time).total_seconds()

            # Extract token usage from response
            token_usage = self._extract_token_usage(response_obj)
            prompt_tokens = token_usage.get("prompt_tokens", 0)
            completion_tokens = token_usage.get("completion_tokens", 0)
            total_tokens = token_usage.get("total_tokens", 0)

            # Calculate cost using litellm.completion_cost()
            cost = self._calculate_cost(response_obj, kwargs)

            # Common attributes for all metrics
            attributes = {
                "model": model,
                "provider": self._extract_provider(model),
                "status": "success",
            }

            # Record metrics
            self._requests_counter.add(1, attributes=attributes)
            self._tokens_counter.add(total_tokens, attributes={**attributes, "token_type": "total"})
            self._tokens_counter.add(
                prompt_tokens, attributes={**attributes, "token_type": "prompt"}
            )
            self._tokens_counter.add(
                completion_tokens, attributes={**attributes, "token_type": "completion"}
            )
            self._latency_histogram.record(latency, attributes=attributes)

            if cost is not None and cost > 0:
                self._cost_histogram.record(cost, attributes=attributes)

            # Complete trace span
            span_data = self._active_spans.pop(call_id, None)
            if span_data:
                span = span_data["span"]
                span.set_attribute("emms_genai_tokens_prompt", prompt_tokens)
                span.set_attribute("emms_genai_tokens_completion", completion_tokens)
                span.set_attribute("emms_genai_tokens_total", total_tokens)
                span.set_attribute("emms_genai_latency_seconds", latency)
                if cost is not None:
                    span.set_attribute("emms_genai_cost_usd", cost)
                span.set_status(Status(StatusCode.OK))
                span.end()

            # Log success (no PHI)
            self._client.logger.info(
                "LLM request completed",
                extra={
                    "model": model,
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens,
                    "latency_seconds": latency,
                    "cost_usd": cost,
                    "status": "success",
                },
            )
        except (AttributeError, KeyError) as e:
            # Specific: Data extraction/access errors (likely malformed response)
            print(
                f"MCALiteLLMCallback.log_success_event - data extraction error: {e}",
                file=sys.stderr,
            )
        except (ConnectionError, TimeoutError) as e:
            # Specific: Network errors during metric export
            print(
                f"MCALiteLLMCallback.log_success_event - telemetry backend unavailable: {e}",
                file=sys.stderr,
            )
        except Exception as e:
            # Catch-all: Unexpected errors (don't crash user's app)
            print(
                f"MCALiteLLMCallback.log_success_event - unexpected error: {e}",
                file=sys.stderr,
            )
            # Suppress exception - monitoring failure should not break the app

    def log_failure_event(
        self, kwargs: dict, response_obj: Any, start_time: datetime, end_time: datetime
    ) -> None:
        """Called on failed LLM API call.

        Records failure metrics and completes the trace span with error status.

        Note: This method MUST NOT raise exceptions that could crash the user's application.
        All errors are caught, logged to stderr, and suppressed.

        Security: Error messages are NOT logged by default (may contain PHI).
        Only error_type is recorded unless record_error_messages=True.
        """
        try:
            call_id = kwargs.get("litellm_call_id", str(id(kwargs)))
            model = kwargs.get("model", "unknown")

            # Calculate latency
            latency = (end_time - start_time).total_seconds()

            # Extract error information
            error_type, error_message = self._extract_error_info(response_obj)

            # Common attributes
            attributes = {
                "model": model,
                "provider": self._extract_provider(model),
                "status": "error",
                "error_type": error_type,
            }

            # Record failure metrics
            self._requests_counter.add(1, attributes=attributes)
            self._requests_failed_counter.add(1, attributes=attributes)
            self._latency_histogram.record(latency, attributes=attributes)

            # Complete trace span with error
            span_data = self._active_spans.pop(call_id, None)
            if span_data:
                span = span_data["span"]
                span.set_attribute("error", True)
                span.set_attribute("error.type", error_type)
                # SECURITY: Only add error message if explicitly enabled
                if self._record_error_messages:
                    span.set_attribute("error.message", error_message[:200])
                span.set_attribute("emms_genai_latency_seconds", latency)
                # Don't include error message in status (may contain PHI)
                span.set_status(Status(StatusCode.ERROR, f"LLM request failed: {error_type}"))
                span.end()

            # Log error (error_type only by default - no PHI)
            log_extra = {
                "model": model,
                "error_type": error_type,
                "latency_seconds": latency,
                "status": "error",
            }
            # SECURITY: Only add error message if explicitly enabled
            if self._record_error_messages:
                log_extra["error_message"] = error_message[:200]

            self._client.logger.error(f"LLM request failed: {error_type}", extra=log_extra)
        except (AttributeError, KeyError) as e:
            # Specific: Data extraction/access errors
            print(
                f"MCALiteLLMCallback.log_failure_event - data extraction error: {e}",
                file=sys.stderr,
            )
        except (ConnectionError, TimeoutError) as e:
            # Specific: Network errors during metric export
            print(
                f"MCALiteLLMCallback.log_failure_event - telemetry backend unavailable: {e}",
                file=sys.stderr,
            )
        except Exception as e:
            # Catch-all: Unexpected errors
            print(
                f"MCALiteLLMCallback.log_failure_event - unexpected error: {e}",
                file=sys.stderr,
            )
            # Suppress exception - monitoring failure should not break the app

    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 for acompletion calls.

        Note: This method MUST NOT raise exceptions that could crash the user's application.
        """
        try:
            # Delegate to sync implementation (metrics/tracing are thread-safe)
            self.log_success_event(kwargs, response_obj, start_time, end_time)
        except Exception as e:
            # CRITICAL: Never crash the user's application
            print(f"MCALiteLLMCallback.async_log_success_event failed: {e}", file=sys.stderr)

    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 for acompletion calls.

        Note: This method MUST NOT raise exceptions that could crash the user's application.
        """
        try:
            # Delegate to sync implementation (metrics/tracing are thread-safe)
            self.log_failure_event(kwargs, response_obj, start_time, end_time)
        except Exception as e:
            # CRITICAL: Never crash the user's application
            print(f"MCALiteLLMCallback.async_log_failure_event failed: {e}", file=sys.stderr)

    def _extract_token_usage(self, response_obj: Any) -> Dict[str, int]:
        """Extract token usage from LiteLLM response.

        Handles both ModelResponse objects and dict responses.
        """
        try:
            # Try ModelResponse.usage attribute first
            if hasattr(response_obj, "usage") and response_obj.usage:
                usage = response_obj.usage
                return {
                    "prompt_tokens": getattr(usage, "prompt_tokens", 0) or 0,
                    "completion_tokens": getattr(usage, "completion_tokens", 0) or 0,
                    "total_tokens": getattr(usage, "total_tokens", 0) or 0,
                }

            # Try dict access
            if isinstance(response_obj, dict) and "usage" in response_obj:
                usage = response_obj["usage"]
                return {
                    "prompt_tokens": usage.get("prompt_tokens", 0) or 0,
                    "completion_tokens": usage.get("completion_tokens", 0) or 0,
                    "total_tokens": usage.get("total_tokens", 0) or 0,
                }
        except Exception as e:
            logger.debug(f"Could not extract token usage: {e}")

        return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    def _calculate_cost(self, response_obj: Any, kwargs: dict) -> Optional[float]:
        """Calculate request cost using litellm.completion_cost().

        Returns None if cost cannot be calculated.

        Note: This method attempts to use litellm's public completion_cost() API.
        If cost calculation fails (e.g., no pricing data for custom models,
        local models like Ollama, or if the API changes), None is returned
        and the cost metric is simply not recorded for that request.

        Limitations:
            - Cost data unavailable for local/custom models without pricing
            - Relies on LiteLLM maintaining pricing data for commercial models
            - If litellm.completion_cost() fails, cost metric is skipped
        """
        try:
            # Use public API: litellm.completion_cost()
            cost = litellm.completion_cost(completion_response=response_obj)
            return float(cost) if cost else None

        except Exception as e:
            # Cost calculation failures are expected for custom/local models
            # Log as warning to make metric loss visible to operators
            model_name = kwargs.get("model", "unknown")
            logger.warning(
                f"Failed to calculate cost for model '{model_name}': {e}. "
                f"Cost metric will not be recorded. This is expected for custom/local models."
            )
            return None

    def _extract_provider(self, model: str) -> str:
        """Extract provider from model string (e.g., 'openai/gpt-4' -> 'openai')."""
        if "/" in model:
            return model.split("/")[0]

        # Common model prefix mappings
        model_lower = model.lower()
        if model_lower.startswith("gpt"):
            return "openai"
        elif model_lower.startswith("claude"):
            return "anthropic"
        elif model_lower.startswith("gemini"):
            return "google"
        elif model_lower.startswith("llama"):
            return "meta"
        elif model_lower.startswith("mistral"):
            return "mistral"

        return "unknown"

    def _extract_span_attributes(self, model: str, kwargs: dict) -> Dict[str, Any]:
        """Extract span attributes from request kwargs (no PHI)."""
        attributes = {
            "llm.model": model,
            "llm.provider": self._extract_provider(model),
        }

        # Extract common parameters (no prompt content)
        if "temperature" in kwargs:
            attributes["llm.temperature"] = kwargs["temperature"]
        if "max_tokens" in kwargs:
            attributes["llm.max_tokens"] = kwargs["max_tokens"]
        if "top_p" in kwargs:
            attributes["llm.top_p"] = kwargs["top_p"]
        if "stream" in kwargs:
            attributes["llm.stream"] = kwargs["stream"]

        return attributes

    def _extract_error_info(self, response_obj: Any) -> tuple:
        """Extract error type and message from failure response.

        Returns (error_type, error_message) tuple.

        Note: Error message is extracted but only used when record_error_messages=True.
        Truncation is applied as defense-in-depth against PHI leakage.
        """
        error_type = "UnknownError"
        error_message = "Unknown error occurred"

        try:
            if isinstance(response_obj, Exception):
                error_type = type(response_obj).__name__
                error_message = str(response_obj)
            elif hasattr(response_obj, "error"):
                error = response_obj.error
                if isinstance(error, dict):
                    error_type = error.get("type", "APIError")
                    error_message = error.get("message", str(error))
                else:
                    error_message = str(error)
        except Exception:
            pass

        # Sanitize error message (truncate as defense-in-depth)
        # Note: Message is only logged when record_error_messages=True
        if len(error_message) > 200:
            error_message = error_message[:200] + "...[truncated]"

        return error_type, error_message

    def _check_global_state_safety(self) -> None:
        """Runtime check for global state conflicts.

        Detects if multiple MCALiteLLMCallback instances may be registered,
        which would cause metric mis-attribution due to LiteLLM's global callback system.

        This is a defensive measure but cannot prevent all conflicts.
        """
        try:
            # Check if callbacks are already registered
            if hasattr(litellm, "callbacks") and litellm.callbacks:
                # Check for MCALiteLLMCallback instances
                mca_callbacks = [
                    cb for cb in litellm.callbacks if isinstance(cb, MCALiteLLMCallback)
                ]

                if len(mca_callbacks) > 0:
                    # Another MCALiteLLMCallback is already registered
                    logger.warning(
                        "Multiple MCALiteLLMCallback instances detected in litellm.callbacks. "
                        "This may cause metric mis-attribution due to LiteLLM's global callback registry. "
                        "Only ONE callback instance should be registered per process. "
                        "This callback will still function but data integrity may be affected. "
                        "See module docstring for mitigation guidance."
                    )
        except Exception as e:
            # Don't fail initialization if check fails
            logger.debug(f"Could not perform global state safety check: {e}")
