"""Adapter for patching LiteLLM completion calls.

LiteLLM is used internally by OpenAI Agents SDK, so this adapter is critical
for capturing LLM spans when using the Agents SDK. Our existing OpenAI/Anthropic
patches don't capture LiteLLM calls because LiteLLM abstracts the underlying API.
"""

import copy
import functools
import logging
from typing import Any, Dict, Optional

from klira.sdk.adapters.llm_base_adapter import BaseLLMAdapter

# Try to import LiteLLM
try:
    import litellm

    LITELLM_AVAILABLE = True
except ImportError:
    litellm = None  # type: ignore[assignment]
    LITELLM_AVAILABLE = False

# Try to import OTel context
try:
    from opentelemetry import context as otel_context
except ImportError:
    otel_context = None  # type: ignore[assignment]

try:
    from opentelemetry import trace as otel_trace
    from opentelemetry.trace import StatusCode

    OTEL_TRACE_AVAILABLE = True
except ImportError:
    otel_trace = None  # type: ignore[assignment]
    StatusCode = None  # type: ignore[assignment, misc]
    OTEL_TRACE_AVAILABLE = False

logger = logging.getLogger("klira.adapters.litellm")


class LiteLLMAdapter(BaseLLMAdapter):
    """Patches LiteLLM completion calls for guideline injection and tracing.

    LiteLLM provides a unified interface for multiple LLM providers. This adapter
    ensures that all LiteLLM calls (regardless of underlying provider) are traced
    with consistent span attributes following the Unified Trace Architecture.
    """

    is_available = LITELLM_AVAILABLE

    def __init__(self) -> None:
        super().__init__()
        self.original_completion: Optional[Any] = None
        self.original_acompletion: Optional[Any] = None
        self._patched: bool = False

    def patch(self) -> None:
        """Patch LiteLLM completion methods."""
        if self._patched:
            logger.debug("LiteLLM adapter already patched")
            return

        if not LITELLM_AVAILABLE or litellm is None:
            logger.debug("LiteLLM not available. Skipping patch.")
            return

        logger.debug("LiteLLMAdapter: Attempting to patch litellm.completion...")

        self._patch_sync_completion()
        self._patch_async_completion()

        self._patched = True
        logger.info("LiteLLM adapter patching complete.")

    def _patch_sync_completion(self) -> None:
        """Patch litellm.completion()"""
        if not hasattr(litellm, "completion"):
            logger.debug("litellm.completion not found, skipping sync patch")
            return

        original_completion = litellm.completion
        self.original_completion = original_completion

        # Store reference to self for use in nested function
        adapter_self = self

        @functools.wraps(original_completion)
        def patched_completion(*args: Any, **kwargs: Any) -> Any:
            from klira.sdk.tracing.tracing import get_tracer as get_traceloop_tracer

            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            original_messages = copy.deepcopy(messages) if messages else None

            # Check if we're in a unified trace context (matches Anthropic adapter pattern)
            if adapter_self._is_in_unified_trace():
                tracer = otel_trace.get_tracer("klira")
            elif OTEL_TRACE_AVAILABLE and otel_trace:
                tracer = get_traceloop_tracer()
            else:
                # No tracing available, just execute with guidelines
                try:
                    modified_kwargs = adapter_self._inject_guidelines_into_kwargs(kwargs)
                except Exception as e:
                    logger.error(f"Policy injection error: {str(e)}")
                    modified_kwargs = kwargs
                result = original_completion(*args, **modified_kwargs)
                return adapter_self._apply_outbound_guardrails(result)

            with tracer.start_as_current_span("klira.llm.litellm") as llm_span:
                try:
                    # Set request attributes using base class method
                    adapter_self._create_llm_request_span(
                        llm_span,
                        provider="litellm",
                        model=model,
                        messages=messages,
                        temperature=kwargs.get("temperature"),
                        max_tokens=kwargs.get("max_tokens"),
                        top_p=kwargs.get("top_p"),
                    )

                    # Add augmentation tracking
                    adapter_self._add_augmentation_attributes(
                        llm_span, original_messages=original_messages
                    )

                    # Inject guidelines
                    try:
                        modified_kwargs = adapter_self._inject_guidelines_into_kwargs(
                            kwargs
                        )
                    except Exception as e:
                        logger.error(f"Policy injection error: {str(e)}")
                        modified_kwargs = kwargs

                    # Call original
                    result = original_completion(*args, **modified_kwargs)

                    # Add response attributes using base class method
                    adapter_self._add_llm_response_attributes(llm_span, result, "litellm")

                    if OTEL_TRACE_AVAILABLE:
                        llm_span.set_status(StatusCode.OK)

                except Exception as e:
                    if OTEL_TRACE_AVAILABLE:
                        llm_span.set_status(StatusCode.ERROR, str(e))
                    llm_span.record_exception(e)
                    raise

                # Apply outbound guardrails
                result = adapter_self._apply_outbound_guardrails(result)
                return result

        litellm.completion = patched_completion
        logger.info("Successfully patched litellm.completion")

    def _patch_async_completion(self) -> None:
        """Patch litellm.acompletion()"""
        if not hasattr(litellm, "acompletion"):
            logger.debug("litellm.acompletion not found, skipping async patch")
            return

        original_acompletion = litellm.acompletion
        self.original_acompletion = original_acompletion

        # Store reference to self for use in nested function
        adapter_self = self

        @functools.wraps(original_acompletion)
        async def patched_acompletion(*args: Any, **kwargs: Any) -> Any:
            from klira.sdk.tracing.tracing import get_tracer as get_traceloop_tracer

            model = kwargs.get("model", "unknown")
            messages = kwargs.get("messages", [])
            original_messages = copy.deepcopy(messages) if messages else None

            # Check if we're in a unified trace context
            if adapter_self._is_in_unified_trace():
                tracer = otel_trace.get_tracer("klira")
            elif OTEL_TRACE_AVAILABLE and otel_trace:
                tracer = get_traceloop_tracer()
            else:
                # No tracing available, just execute with guidelines
                try:
                    modified_kwargs = adapter_self._inject_guidelines_into_kwargs(kwargs)
                except Exception as e:
                    logger.error(f"Policy injection error: {str(e)}")
                    modified_kwargs = kwargs
                result = await original_acompletion(*args, **modified_kwargs)
                return await adapter_self._apply_outbound_guardrails_async(result)

            with tracer.start_as_current_span("klira.llm.litellm") as llm_span:
                try:
                    # Set request attributes using base class method
                    adapter_self._create_llm_request_span(
                        llm_span,
                        provider="litellm",
                        model=model,
                        messages=messages,
                        temperature=kwargs.get("temperature"),
                        max_tokens=kwargs.get("max_tokens"),
                        top_p=kwargs.get("top_p"),
                    )

                    # Add augmentation tracking
                    adapter_self._add_augmentation_attributes(
                        llm_span, original_messages=original_messages
                    )

                    # Inject guidelines
                    try:
                        modified_kwargs = adapter_self._inject_guidelines_into_kwargs(
                            kwargs
                        )
                    except Exception as e:
                        logger.error(f"Policy injection error: {str(e)}")
                        modified_kwargs = kwargs

                    # Call original
                    result = await original_acompletion(*args, **modified_kwargs)

                    # Add response attributes using base class method
                    adapter_self._add_llm_response_attributes(llm_span, result, "litellm")

                    if OTEL_TRACE_AVAILABLE:
                        llm_span.set_status(StatusCode.OK)

                except Exception as e:
                    if OTEL_TRACE_AVAILABLE:
                        llm_span.set_status(StatusCode.ERROR, str(e))
                    llm_span.record_exception(e)
                    raise

                # Apply outbound guardrails
                result = await adapter_self._apply_outbound_guardrails_async(result)
                return result

        litellm.acompletion = patched_acompletion
        logger.info("Successfully patched litellm.acompletion")

    def _inject_guidelines_into_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Inject guidelines into LiteLLM kwargs (OpenAI-compatible format)."""
        guidelines = None

        # Get guidelines from GuardrailsEngine
        try:
            from klira.sdk.guardrails.engine import GuardrailsEngine

            engine = GuardrailsEngine.get_instance()
            if engine:
                guidelines = engine.get_current_guidelines()
                if guidelines:
                    logger.debug(
                        f"Found {len(guidelines)} guidelines to inject at LiteLLM call time"
                    )

            # Fallback to OTel context if no guidelines from engine
            if not guidelines and otel_context is not None:
                try:
                    current_ctx = otel_context.get_current()
                    otel_guidelines = otel_context.get_value(
                        "klira.augmentation.guidelines", context=current_ctx
                    )
                    if otel_guidelines and isinstance(otel_guidelines, list):
                        guidelines = otel_guidelines
                        logger.debug(
                            f"[LiteLLMAdapter] Retrieved {len(guidelines)} guidelines from OTel context."
                        )
                except Exception as e:
                    logger.debug(
                        f"[LiteLLMAdapter] Could not retrieve guidelines from OTel context: {e}"
                    )

            if not guidelines:
                logger.debug("[LiteLLMAdapter] No guidelines found to inject.")
                return kwargs

            # Create a copy of kwargs to avoid modifying the original
            modified_kwargs = kwargs.copy()

            # LiteLLM uses OpenAI-compatible messages format
            messages = modified_kwargs.get("messages", [])
            if messages:
                messages_copy = [msg.copy() for msg in messages]

                # Find the last user message to augment
                user_msg_index = None
                for i in range(len(messages_copy) - 1, -1, -1):
                    if messages_copy[i].get("role") == "user":
                        user_msg_index = i
                        break

                if user_msg_index is not None:
                    user_msg = messages_copy[user_msg_index]
                    content = user_msg.get("content", "")

                    if isinstance(content, str):
                        augmentation_text = "\n\nIMPORTANT POLICY DIRECTIVES:\n"
                        augmentation_text += "\n".join([f"• {g}" for g in guidelines])
                        user_msg["content"] = content + augmentation_text
                        messages_copy[user_msg_index] = user_msg
                        modified_kwargs["messages"] = messages_copy

                        logger.info(
                            f"Injected {len(guidelines)} policy guidelines into LiteLLM user message"
                        )
                else:
                    # No user message, try system message or prepend one
                    system_found = False
                    for i, msg in enumerate(messages_copy):
                        if msg.get("role") == "system":
                            content = msg.get("content", "")
                            if isinstance(content, str):
                                augmentation_text = "\n\nIMPORTANT POLICY DIRECTIVES:\n"
                                augmentation_text += "\n".join(
                                    [f"• {g}" for g in guidelines]
                                )
                                msg["content"] = content + augmentation_text
                                messages_copy[i] = msg
                                modified_kwargs["messages"] = messages_copy
                                system_found = True
                                logger.info(
                                    f"Injected {len(guidelines)} policy guidelines into LiteLLM system message"
                                )
                                break

                    if not system_found:
                        # Prepend a system message with guidelines
                        augmentation_text = "IMPORTANT POLICY DIRECTIVES:\n"
                        augmentation_text += "\n".join([f"• {g}" for g in guidelines])
                        messages_copy.insert(
                            0, {"role": "system", "content": augmentation_text}
                        )
                        modified_kwargs["messages"] = messages_copy
                        logger.info(
                            f"Created system message with {len(guidelines)} policy guidelines for LiteLLM"
                        )

            # Clear guidelines from context to prevent double-injection
            if engine:
                engine.clear_current_guidelines()

            return modified_kwargs

        except Exception as e:
            logger.error(
                f"Error injecting guidelines into LiteLLM request: {e}", exc_info=True
            )
            return kwargs

    def _apply_outbound_guardrails(self, result: Any) -> Any:
        """Apply outbound guardrails to LiteLLM results (synchronous)."""
        try:
            from klira.sdk.guardrails.engine import GuardrailsEngine

            response_text = self._extract_response_text(result, "litellm")
            if not response_text.strip():
                return result

            context = {
                "llm_client": "litellm",
                "function_name": "litellm.completion",
            }

            engine = GuardrailsEngine.get_instance()

            import asyncio

            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    logger.warning(
                        "Cannot run outbound guardrails for LiteLLM in sync context within async loop."
                    )
                    return result
                else:
                    decision = loop.run_until_complete(
                        engine.evaluate(response_text, context, direction="outbound")
                    )
            except RuntimeError:
                decision = asyncio.run(
                    engine.evaluate(response_text, context, direction="outbound")
                )

            if not decision.allowed:
                logger.warning(
                    f"LiteLLM outbound guardrails blocked response: {decision.reason}"
                )
                return self._create_blocked_response(
                    result, decision.reason or "Content policy violation detected"
                )

            return result

        except Exception as e:
            logger.error(
                f"Error applying outbound guardrails for LiteLLM: {e}",
                exc_info=True,
            )
            return result

    async def _apply_outbound_guardrails_async(self, result: Any) -> Any:
        """Apply outbound guardrails to LiteLLM results (asynchronous)."""
        try:
            from klira.sdk.guardrails.engine import GuardrailsEngine

            response_text = self._extract_response_text(result, "litellm")
            if not response_text.strip():
                return result

            context = {
                "llm_client": "litellm",
                "function_name": "litellm.acompletion",
            }

            engine = GuardrailsEngine.get_instance()
            decision = await engine.evaluate(
                response_text, context, direction="outbound"
            )

            if not decision.allowed:
                logger.warning(
                    f"LiteLLM async outbound guardrails blocked response: {decision.reason}"
                )
                return self._create_blocked_response(
                    result, decision.reason or "Content policy violation detected"
                )

            return result

        except Exception as e:
            logger.error(
                f"Error applying outbound guardrails for LiteLLM async: {e}",
                exc_info=True,
            )
            return result

    def _create_blocked_response(self, original_result: Any, reason: str) -> Any:
        """Create a blocked response that matches the original response format."""
        try:
            blocked_message = (
                f"[BLOCKED BY GUARDRAILS] - {reason}"
                if reason
                else "[BLOCKED BY GUARDRAILS] - Content policy violation detected"
            )

            # LiteLLM uses OpenAI-compatible response format
            if hasattr(original_result, "choices") and original_result.choices:
                choice = original_result.choices[0]
                if hasattr(choice, "message") and hasattr(choice.message, "content"):
                    choice.message.content = blocked_message
                    return original_result

            return blocked_message

        except Exception as e:
            logger.error(f"Error creating blocked response: {e}")
            return "[BLOCKED BY GUARDRAILS] - Content policy violation detected"

    # Provider-specific extraction methods for LLM span attributes

    def _extract_prompt_text(self, messages: Any, provider: str) -> str:
        """Extract prompt text from LiteLLM messages format (OpenAI-compatible)."""
        if isinstance(messages, list):
            parts = []
            for msg in messages:
                if isinstance(msg, dict):
                    role = msg.get("role", "user")
                    content = msg.get("content", "")
                    if isinstance(content, str):
                        parts.append(f"{role}: {content}")
                    elif isinstance(content, list):
                        # Handle content blocks
                        for block in content:
                            if isinstance(block, dict) and "text" in block:
                                parts.append(f"{role}: {block['text']}")
            return "\n".join(parts) if parts else str(messages)
        return str(messages)

    def _extract_response_text(self, response: Any, provider: str) -> str:
        """Extract response text from LiteLLM response (OpenAI-compatible)."""
        try:
            if hasattr(response, "choices") and response.choices:
                choice = response.choices[0]
                if hasattr(choice, "message") and hasattr(choice.message, "content"):
                    return str(choice.message.content or "")
                elif hasattr(choice, "text"):
                    return str(choice.text or "")
            return str(response) if response else ""
        except Exception as e:
            logger.debug(f"Error extracting LiteLLM response text: {e}")
            return ""

    def _extract_token_usage(
        self, response: Any, provider: str
    ) -> Optional[Dict[str, int]]:
        """Extract token usage from LiteLLM response."""
        try:
            if hasattr(response, "usage") and response.usage:
                return {
                    "input_tokens": getattr(response.usage, "prompt_tokens", 0),
                    "output_tokens": getattr(response.usage, "completion_tokens", 0),
                }
            return None
        except Exception as e:
            logger.debug(f"Error extracting LiteLLM token usage: {e}")
            return None

    def _extract_finish_reason(self, response: Any, provider: str) -> Optional[str]:
        """Extract finish reason from LiteLLM response."""
        try:
            if hasattr(response, "choices") and response.choices:
                choice = response.choices[0]
                if hasattr(choice, "finish_reason") and choice.finish_reason:
                    return str(choice.finish_reason)
            return None
        except Exception as e:
            logger.debug(f"Error extracting LiteLLM finish reason: {e}")
            return None

    def _extract_response_model(self, response: Any, provider: str) -> Optional[str]:
        """Extract model name from LiteLLM response."""
        try:
            if hasattr(response, "model") and response.model is not None:
                return str(response.model)
            return None
        except Exception as e:
            logger.debug(f"Error extracting LiteLLM response model: {e}")
            return None

    def unpatch_llm_client(self) -> bool:
        """Restore original LiteLLM methods."""
        if not LITELLM_AVAILABLE or not self._patched or litellm is None:
            return False

        try:
            if self.original_completion:
                litellm.completion = self.original_completion
                self.original_completion = None

            if self.original_acompletion:
                litellm.acompletion = self.original_acompletion
                self.original_acompletion = None

            self._patched = False
            logger.info("Successfully unpatched LiteLLM methods")
            return True

        except Exception as e:
            logger.error(f"Failed to unpatch LiteLLM: {e}")
            return False

    def get_client_info(self) -> Dict[str, Any]:
        """Get information about the LiteLLM client."""
        if not LITELLM_AVAILABLE or litellm is None:
            return {"available": False, "reason": "litellm not installed"}

        try:
            return {
                "available": True,
                "library": "litellm",
                "version": getattr(litellm, "__version__", "unknown"),
                "patched": self._patched,
            }
        except Exception as e:
            return {"available": False, "reason": str(e)}
