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

from typing import TYPE_CHECKING, Final, Literal, Optional

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: str | None = None,
        api_base: str | None = None,
        model: Literal["omni-moderation-latest", "text-moderation-latest"] | None = None,
        streaming_end_of_stream_only: bool | None = None,
        streaming_sampling_rate: int | None = None,
        **kwargs,
    ):
        """Initialize OpenAI Moderation guardrail handler."""
        super().__init__(
            guardrail_name=guardrail_name,
            supported_event_hooks=list(self.get_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(
            "Initialized OpenAI Moderation Guardrail: %s with model: %s", guardrail_name, self.model
        )

    def _get_api_key(self) -> str | None:
        """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: Final = {"model": self.model, "input": input_text}

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

        response: Final = 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: Final = moderation_response.results[0]
        if result.flagged:
            # Build detailed violation information
            violated_categories: Final = []
            if result.categories:
                for category, is_violated in result.categories.items():
                    if is_violated:
                        violated_categories.append(category)

            violation_details: Final = {
                "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: str | None = 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: Final = 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: Final = 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: dict | None,
        request_data: dict,
        start_time: float | None = None,
        end_time: float | None = None,
        duration: float | None = None,
        event_type: GuardrailEventHooks | None = None,
        original_inputs: dict | None = 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: Final = 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: float | None = None,
        end_time: float | None = None,
        duration: float | None = None,
        event_type: GuardrailEventHooks | None = None,
    ):
        """
        Override to log the full OpenAI Moderation API response on error
        instead of the stringified exception.
        """
        guardrail_status: Final[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: Final[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: dict | str | Exception,
    ) -> GuardrailTracingDetail | None:
        """
        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: Final = guardrail_response.get("results") or []
        violation_categories: Final = [
            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() -> type["GuardrailConfigModel"] | None:
        """
        Get the config model for the OpenAI Moderation guardrail.
        """
        from litellm.types.proxy.guardrails.guardrail_hooks.openai.openai_moderation import (
            OpenAIModerationGuardrailConfigModel,
        )

        return OpenAIModerationGuardrailConfigModel

    @classmethod
    def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
        return [
            GuardrailEventHooks.pre_call,
            GuardrailEventHooks.during_call,
            GuardrailEventHooks.post_call,
        ]
