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

from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    Dict,
    List,
    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.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)

from .base import OpenAIGuardrailBase

if TYPE_CHECKING:
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.llms.openai import AllMessageValues, OpenAIModerationResponse
    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
    from litellm.types.utils import ModelResponse, ModelResponseStream


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,
        **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"

        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 async_pre_call_hook(
        self,
        user_api_key_dict: "UserAPIKeyAuth",
        cache: Any,
        data: Dict[str, Any],
        call_type: Literal[
            "completion",
            "text_completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "pass_through_endpoint",
            "rerank",
            "mcp_call",
        ],
    ) -> Optional[Dict[str, Any]]:
        """
        Pre-call hook to scan user prompts before sending to LLM.

        Raises HTTPException if content should be blocked.
        """
        verbose_proxy_logger.debug(
            "OpenAI Moderation: Running pre-call prompt scan, on call_type: %s",
            call_type,
        )
        
        # Skip moderation calls to avoid infinite recursion
        if call_type == "moderation":
            return data

        new_messages: Optional[List["AllMessageValues"]] = data.get("messages")
        if new_messages is None:
            verbose_proxy_logger.warning(
                "OpenAI Moderation: not running guardrail. No messages in data"
            )
            return data

        user_prompt = self.get_user_prompt(new_messages)
        if user_prompt:
            verbose_proxy_logger.debug(
                f"OpenAI Moderation: User prompt: {user_prompt[:100]}..."  # Log first 100 chars for debugging
            )
            
            moderation_response = await self.async_make_request(
                input_text=user_prompt,
            )
            
            # Check if content is flagged and raise exception if needed
            self._check_moderation_result(moderation_response)
        else:
            verbose_proxy_logger.warning(
                "OpenAI Moderation: No user prompt found"
            )

        return data

    @log_guardrail_information
    async def async_moderation_hook(
        self,
        data: Dict[str, Any],
        user_api_key_dict: "UserAPIKeyAuth",
        call_type: Literal[
            "completion",
            "embeddings", 
            "image_generation",
            "moderation",
            "audio_transcription",
            "responses",
            "mcp_call",
        ],
    ) -> Optional[Dict[str, Any]]:
        """
        Moderation hook to scan user prompts during call processing.

        Raises HTTPException if content should be blocked.
        """
        verbose_proxy_logger.debug(
            "OpenAI Moderation: Running moderation hook, on call_type: %s",
            call_type,
        )
        
        # Skip moderation calls to avoid infinite recursion
        if call_type == "moderation":
            return data

        new_messages: Optional[List["AllMessageValues"]] = data.get("messages")
        if new_messages is None:
            verbose_proxy_logger.warning(
                "OpenAI Moderation: not running guardrail. No messages in data"
            )
            return data

        user_prompt = self.get_user_prompt(new_messages)
        if user_prompt:
            moderation_response = await self.async_make_request(
                input_text=user_prompt,
            )
            
            # Check if content is flagged and raise exception if needed
            self._check_moderation_result(moderation_response)

        return data

    @log_guardrail_information
    async def async_post_call_hook(
        self,
        data: Dict[str, Any],
        user_api_key_dict: "UserAPIKeyAuth",
        response: "ModelResponse",
    ) -> "ModelResponse":
        """
        Post-call hook to scan LLM responses before returning to user.

        Raises HTTPException if response should be blocked.
        """
        verbose_proxy_logger.debug(
            "OpenAI Moderation: Running post-call response scan"
        )

        # Extract response text for moderation
        response_text = self._extract_response_text(response)
        if response_text:
            verbose_proxy_logger.debug(
                f"OpenAI Moderation: Response text: {response_text[:100]}..."  # Log first 100 chars
            )
            
            moderation_response = await self.async_make_request(
                input_text=response_text,
            )
            
            # Check if content is flagged and raise exception if needed
            self._check_moderation_result(moderation_response)

        return response

    @log_guardrail_information
    async def async_post_call_streaming_iterator_hook(
        self,
        user_api_key_dict: "UserAPIKeyAuth",
        response: Any,
        request_data: Dict[str, Any],
    ) -> AsyncGenerator["ModelResponseStream", None]:
        """
        Process streaming response chunks for OpenAI moderation.

        Collects all chunks from the stream, assembles them into a complete response,
        and applies moderation check. If content violates moderation policy, raises HTTPException.
        """
        # Import here to avoid circular imports
        from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
        from litellm.main import stream_chunk_builder
        from litellm.types.utils import TextCompletionResponse

        verbose_proxy_logger.debug(
            "OpenAI Moderation: Running streaming response scan"
        )

        # Collect all chunks to process them together
        all_chunks: List["ModelResponseStream"] = []
        async for chunk in response:
            all_chunks.append(chunk)

        # Assemble the complete response from chunks
        assembled_model_response: Optional[
            Union["ModelResponse", TextCompletionResponse]
        ] = stream_chunk_builder(
            chunks=all_chunks,
        )

        if isinstance(assembled_model_response, (type(None), TextCompletionResponse)):
            # If we can't assemble a ModelResponse or it's a text completion, 
            # just yield the original chunks without moderation
            verbose_proxy_logger.warning(
                "OpenAI Moderation: Could not assemble ModelResponse from chunks, skipping moderation"
            )
            for chunk in all_chunks:
                yield chunk
            return

        # Extract response text for moderation
        response_text = self._extract_response_text(assembled_model_response)
        if response_text:
            verbose_proxy_logger.debug(
                f"OpenAI Moderation: Streaming response text: {response_text[:100]}..."  # Log first 100 chars
            )
            
            # Make moderation request - this will raise HTTPException if content is flagged
            moderation_response = await self.async_make_request(
                input_text=response_text,
            )
            
            # Check if content is flagged and raise exception if needed
            self._check_moderation_result(moderation_response)

        # If we reach here, content passed moderation - yield the original chunks
        mock_response = MockResponseIterator(
            model_response=assembled_model_response
        )

        # Return the reconstructed stream
        async for chunk in mock_response:
            yield chunk

    def _extract_response_text(self, response: "ModelResponse") -> Optional[str]:
        """
        Extract text content from the model response for moderation.
        """
        if not hasattr(response, 'choices') or not response.choices:
            return None

        response_texts = []
        for choice in response.choices:
            try:
                # Try to get content from message (chat completion)
                message = getattr(choice, 'message', None)
                if message:
                    content = getattr(message, 'content', None)
                    if content and isinstance(content, str):
                        response_texts.append(content)
                        continue
                
                # Try to get text (text completion)
                text = getattr(choice, 'text', None)
                if text and isinstance(text, str):
                    response_texts.append(text)
                    continue
                
                # Try to get content from delta (streaming)
                delta = getattr(choice, 'delta', None)
                if delta:
                    content = getattr(delta, 'content', None)
                    if content and isinstance(content, str):
                        response_texts.append(content)
                        continue
                        
            except (AttributeError, TypeError):
                # Skip choices that don't have expected attributes
                continue

        return "\n".join(response_texts) if response_texts else None

    @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
