#!/usr/bin/env python3
"""
OpenAI Moderation Guardrail Integration for LiteLLM
"""

from typing import (
    TYPE_CHECKING,
    Dict,
    Literal,
    Optional,
    Type,
    Union,
)

from fastapi import HTTPException

from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
    CustomGuardrail,
    log_guardrail_information,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
    GenericGuardrailAPIInputs,
    GuardrailStatus,
    GuardrailTracingDetail,
)

from .base import OpenAIGuardrailBase

if TYPE_CHECKING:
    from litellm.types.llms.openai import OpenAIModerationResponse
    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel


class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
    """
    LiteLLM Built-in Guardrail for OpenAI Content Moderation.

    This guardrail scans prompts and responses using the OpenAI Moderation API to detect
    harmful content, including violence, hate, harassment, self-harm, sexual content, etc.

    Configuration:
        guardrail_name: Name of the guardrail instance
        api_key: OpenAI API key
        api_base: OpenAI API endpoint
        model: OpenAI moderation model to use
        default_on: Whether to enable by default
    """

    def __init__(
        self,
        guardrail_name: str,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        model: Optional[
            Literal["omni-moderation-latest", "text-moderation-latest"]
        ] = None,
        streaming_end_of_stream_only: Optional[bool] = None,
        streaming_sampling_rate: Optional[int] = None,
        **kwargs,
    ):
        """Initialize OpenAI Moderation guardrail handler."""
        from litellm.types.guardrails import GuardrailEventHooks

        # Initialize parent CustomGuardrail
        supported_event_hooks = [
            GuardrailEventHooks.pre_call,
            GuardrailEventHooks.during_call,
            GuardrailEventHooks.post_call,
        ]
        super().__init__(
            guardrail_name=guardrail_name,
            supported_event_hooks=supported_event_hooks,
            **kwargs,
        )

        self.async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )

        # Store configuration
        self.api_key = api_key or self._get_api_key()
        self.api_base = api_base or "https://api.openai.com/v1"
        self.model: Literal["omni-moderation-latest", "text-moderation-latest"] = (
            model or "omni-moderation-latest"
        )

        # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook
        # via getattr(guardrail_to_apply, "streaming_*", default).
        self.streaming_end_of_stream_only: bool = (
            False
            if streaming_end_of_stream_only is None
            else streaming_end_of_stream_only
        )
        self.streaming_sampling_rate: int = (
            5 if streaming_sampling_rate is None else streaming_sampling_rate
        )

        if not self.api_key:
            raise ValueError(
                "OpenAI Moderation: api_key is required. Set OPENAI_API_KEY environment variable or pass it in configuration."
            )

        verbose_proxy_logger.debug(
            f"Initialized OpenAI Moderation Guardrail: {guardrail_name} with model: {self.model}"
        )

    def _get_api_key(self) -> Optional[str]:
        """Get API key from environment variables or litellm configuration"""
        import os

        import litellm
        from litellm.secret_managers.main import get_secret_str

        return (
            os.environ.get("OPENAI_API_KEY")
            or litellm.api_key
            or litellm.openai_key
            or get_secret_str("OPENAI_API_KEY")
        )

    async def async_make_request(self, input_text: str) -> "OpenAIModerationResponse":
        """
        Make a request to the OpenAI Moderation API.
        """
        request_body = {"model": self.model, "input": input_text}

        verbose_proxy_logger.debug("OpenAI Moderation guard request: %s", request_body)

        response = await self.async_handler.post(
            url=f"{self.api_base}/moderations",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            json=request_body,
        )

        verbose_proxy_logger.debug(
            "OpenAI Moderation guard response: %s", response.json()
        )

        if response.status_code != 200:
            raise HTTPException(
                status_code=response.status_code,
                detail={
                    "error": "OpenAI Moderation API request failed",
                    "details": response.text,
                },
            )

        from litellm.types.llms.openai import OpenAIModerationResponse

        return OpenAIModerationResponse(**response.json())

    def _check_moderation_result(
        self, moderation_response: "OpenAIModerationResponse"
    ) -> None:
        """
        Check if the moderation response indicates harmful content and raise exception if needed.
        """
        if not moderation_response.results:
            return

        result = moderation_response.results[0]
        if result.flagged:
            # Build detailed violation information
            violated_categories = []
            if result.categories:
                for category, is_violated in result.categories.items():
                    if is_violated:
                        violated_categories.append(category)

            violation_details = {
                "violated_categories": violated_categories,
                "category_scores": result.category_scores or {},
            }

            verbose_proxy_logger.warning(
                "OpenAI Moderation: Content flagged for violations: %s",
                violation_details,
            )

            raise HTTPException(
                status_code=400,
                detail={
                    "error": "Violated OpenAI moderation policy",
                    "moderation_result": violation_details,
                },
            )

    @log_guardrail_information
    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional["LiteLLMLoggingObj"] = None,
    ) -> GenericGuardrailAPIInputs:
        """
        Apply OpenAI moderation guardrail using the unified guardrail interface.

        This method is called by the UnifiedLLMGuardrails system for all endpoint types
        (chat completions, embeddings, responses API, etc.).

        Args:
            inputs: GenericGuardrailAPIInputs containing texts and/or structured_messages
            request_data: The original request data
            input_type: Whether this is a "request" (pre-call) or "response" (post-call)
            logging_obj: Optional logging object

        Returns:
            The inputs unchanged (moderation doesn't modify content, only blocks)

        Raises:
            HTTPException: If content violates moderation policy
        """
        # Extract text to moderate from inputs
        text_to_moderate: Optional[str] = None

        # Prefer structured_messages if available (has role context)
        if structured_messages := inputs.get("structured_messages"):
            text_to_moderate = self.get_user_prompt(structured_messages)

        # Fall back to texts
        if not text_to_moderate:
            if texts := inputs.get("texts"):
                # Join all texts for moderation
                text_to_moderate = "\n".join(texts)

        if not text_to_moderate:
            verbose_proxy_logger.debug(
                "OpenAI Moderation: No text content to moderate in inputs"
            )
            return inputs

        # Make moderation request
        moderation_response = await self.async_make_request(input_text=text_to_moderate)

        # Stash full moderation response in request_data for logging
        # (Model Armor pattern — per-request dict avoids race conditions)
        if isinstance(request_data, dict):
            metadata = request_data.get("metadata") or {}
            request_data["metadata"] = metadata
            metadata["_openai_moderation_response"] = moderation_response.model_dump()

        # Check if content is flagged and raise exception if needed
        self._check_moderation_result(moderation_response)

        # Moderation doesn't modify content, just blocks - return inputs unchanged
        return inputs

    def _process_response(
        self,
        response: Optional[Dict],
        request_data: dict,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        duration: Optional[float] = None,
        event_type: Optional[GuardrailEventHooks] = None,
        original_inputs: Optional[Dict] = None,
    ):
        """
        Override to log the full OpenAI Moderation API response instead of
        the decorator's simplified "allow"/"mask" string.

        Follows the Model Armor pattern (model_armor.py:325-360).
        """
        if isinstance(request_data, dict):
            metadata = request_data.get("metadata") or {}
            request_data["metadata"] = metadata  # anchor so pop() mutates the real dict
        else:
            metadata = {}

        # .pop() cleans up the internal key so it doesn't leak to downstream
        # loggers. Falls back to "allow" when no moderation call was made
        # (e.g. no text to moderate — early return in apply_guardrail).
        guardrail_response = metadata.pop("_openai_moderation_response", "allow")

        self.add_standard_logging_guardrail_information_to_request_data(
            guardrail_json_response=guardrail_response,
            request_data=request_data,
            guardrail_status="success",
            duration=duration,
            start_time=start_time,
            end_time=end_time,
            event_type=event_type,
            tracing_detail=self._build_tracing_detail(guardrail_response),
        )
        return response

    def _process_error(
        self,
        e: Exception,
        request_data: dict,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        duration: Optional[float] = None,
        event_type: Optional[GuardrailEventHooks] = None,
    ):
        """
        Override to log the full OpenAI Moderation API response on error
        instead of the stringified exception.
        """
        guardrail_status: GuardrailStatus = (
            "guardrail_intervened"
            if self._is_guardrail_intervention(e)
            else "guardrail_failed_to_respond"
        )

        if isinstance(request_data, dict):
            metadata = request_data.get("metadata") or {}
            request_data["metadata"] = metadata  # anchor so pop() mutates the real dict
        else:
            metadata = {}

        # Use the stashed moderation response if available, fall back to exception
        guardrail_response: Union[dict, Exception, str] = metadata.pop(
            "_openai_moderation_response", e
        )

        self.add_standard_logging_guardrail_information_to_request_data(
            guardrail_json_response=guardrail_response,
            request_data=request_data,
            guardrail_status=guardrail_status,
            duration=duration,
            start_time=start_time,
            end_time=end_time,
            event_type=event_type,
            tracing_detail=self._build_tracing_detail(guardrail_response),
        )
        raise e

    @staticmethod
    def _build_tracing_detail(
        guardrail_response: Union[dict, str, Exception],
    ) -> Optional[GuardrailTracingDetail]:
        """
        Pull the flagged category names out of the moderation response so trace
        backends can index a short, queryable ``guardrail_violation_categories``
        attribute instead of the full ``guardrail_response`` blob, whose
        ``category_scores`` map (one float per category) blows past indexed-field
        length limits on backends like ELK (1024 chars).
        """
        if not isinstance(guardrail_response, dict):
            return None

        results = guardrail_response.get("results") or []
        violation_categories = [
            category
            for result in results
            if isinstance(result, dict)
            for category, is_flagged in (result.get("categories") or {}).items()
            if is_flagged
        ]
        if not violation_categories:
            return None
        return GuardrailTracingDetail(violation_categories=violation_categories)

    @staticmethod
    def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
        """
        Get the config model for the OpenAI Moderation guardrail.
        """
        from litellm.types.proxy.guardrails.guardrail_hooks.openai.openai_moderation import (
            OpenAIModerationGuardrailConfigModel,
        )

        return OpenAIModerationGuardrailConfigModel
