#!/usr/bin/env python3
"""
Azure Text Moderation Native Guardrail Integrationfor LiteLLM
"""

from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast

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 litellm.proxy._types import UserAPIKeyAuth

from .base import AzureGuardrailBase

if TYPE_CHECKING:
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.llms.openai import AllMessageValues
    from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
        AzureTextModerationGuardrailResponse,
    )
    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
    from litellm.types.utils import EmbeddingResponse, ImageResponse, ModelResponse


class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardrail):
    """
    LiteLLM Built-in Guardrail for Azure Content Safety Guardrail (Prompt Shield).

    This guardrail scans prompts and responses using the Azure Prompt Shield API to detect
    malicious content, injection attempts, and policy violations.

    Configuration:
        guardrail_name: Name of the guardrail instance
        api_key: Azure Prompt Shield API key
        api_base: Azure Prompt Shield API endpoint
        default_on: Whether to enable by default
    """

    default_severity_threshold: int = 2

    def __init__(
        self,
        guardrail_name: str,
        api_key: str,
        api_base: str,
        severity_threshold: Optional[int] = None,
        severity_threshold_by_category: Optional[Dict[str, int]] = None,
        **kwargs,
    ):
        """Initialize Azure Text Moderation guardrail handler."""
        # Initialize parent CustomGuardrail
        from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
            AzureTextModerationRequestBodyOptionalParams,
        )

        super().__init__(
            guardrail_name=guardrail_name,
            **kwargs,
        )
        self.async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )

        # Store configuration
        self.api_key = api_key
        self.api_base = api_base
        self.api_version = kwargs.get("api_version") or "2024-09-01"
        self.optional_params_request_body: (
            AzureTextModerationRequestBodyOptionalParams
        ) = {
            "categories": kwargs.get("categories")
            or [
                "Hate",
                "Sexual",
                "SelfHarm",
                "Violence",
            ],
            "blocklistNames": cast(
                Optional[List[str]], kwargs.get("blocklistNames") or None
            ),
            "haltOnBlocklistHit": kwargs.get("haltOnBlocklistHit") or False,
            "outputType": kwargs.get("outputType") or "FourSeverityLevels",
        }

        self.severity_threshold = (
            int(severity_threshold) if severity_threshold else None
        )
        self.severity_threshold_by_category = severity_threshold_by_category

        verbose_proxy_logger.info(
            f"Initialized Azure Prompt Shield Guardrail: {guardrail_name}"
        )

    @staticmethod
    def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
        from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
            AzureContentSafetyTextModerationConfigModel,
        )

        return AzureContentSafetyTextModerationConfigModel

    async def async_make_request(
        self, text: str
    ) -> "AzureTextModerationGuardrailResponse":
        """
        Make a request to the Azure Prompt Shield API.
        """
        from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
            AzureTextModerationGuardrailRequestBody,
            AzureTextModerationGuardrailResponse,
        )

        request_body = AzureTextModerationGuardrailRequestBody(
            text=text,
            **self.optional_params_request_body,
        )
        verbose_proxy_logger.debug(
            "Azure Text Moderation guard request: %s", request_body
        )

        response = await self.async_handler.post(
            url=f"{self.api_base}/contentsafety/text:analyze?api-version={self.api_version}",
            headers={
                "Ocp-Apim-Subscription-Key": self.api_key,
                "Content-Type": "application/json",
            },
            json=cast(dict, request_body),
        )

        verbose_proxy_logger.debug(
            "Azure Text Moderation guard response: %s", response.json()
        )
        return AzureTextModerationGuardrailResponse(**response.json())  # type: ignore

    def check_severity_threshold(
        self, response: "AzureTextModerationGuardrailResponse"
    ) -> Literal[True]:
        """
        - Check if threshold set by category
        - Check if general severity threshold set
        - If both none, use default_severity_threshold
        """

        if self.severity_threshold_by_category:
            for category in response["categoriesAnalysis"]:
                severity_category_threshold_item = (
                    self.severity_threshold_by_category.get(category["category"])
                )
                if (
                    severity_category_threshold_item is not None
                    and category["severity"] >= severity_category_threshold_item
                ):
                    raise HTTPException(
                        status_code=400,
                        detail={
                            "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format(
                                category["category"],
                                self.severity_threshold_by_category.get(
                                    category["category"]
                                ),
                                category["severity"],
                            )
                        },
                    )
        if self.severity_threshold:
            for category in response["categoriesAnalysis"]:
                if category["severity"] >= self.severity_threshold:
                    raise HTTPException(
                        status_code=400,
                        detail={
                            "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format(
                                category["category"],
                                self.severity_threshold,
                                category["severity"],
                            )
                        },
                    )
        if (
            self.severity_threshold is None
            and self.severity_threshold_by_category is None
        ):
            for category in response["categoriesAnalysis"]:
                if category["severity"] >= self.default_severity_threshold:
                    raise HTTPException(
                        status_code=400,
                        detail={
                            "error": "Azure Content Safety Guardrail: {} crossed severity {}, Got severity: {}".format(
                                category["category"],
                                self.default_severity_threshold,
                                category["severity"],
                            )
                        },
                    )
        return True

    @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.info(
            "Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
            call_type,
        )
        new_messages: Optional[List[AllMessageValues]] = data.get("messages")
        if new_messages is None:
            verbose_proxy_logger.warning(
                "Lakera AI: not running guardrail. No messages in data"
            )
            return data
        user_prompt = self.get_user_prompt(new_messages)

        if user_prompt:
            verbose_proxy_logger.info(
                f"Azure Text Moderation: User prompt: {user_prompt}"
            )
            azure_text_moderation_response = await self.async_make_request(
                text=user_prompt,
            )
            self.check_severity_threshold(response=azure_text_moderation_response)
        else:
            verbose_proxy_logger.warning("Azure Text Moderation: No text found")
        return None

    async def async_post_call_success_hook(
        self,
        data: dict,
        user_api_key_dict: "UserAPIKeyAuth",
        response: Union[Any, "ModelResponse", "EmbeddingResponse", "ImageResponse"],
    ) -> Any:
        from litellm.types.utils import Choices, ModelResponse

        if (
            isinstance(response, ModelResponse)
            and response.choices
            and isinstance(response.choices[0], Choices)
        ):
            content = response.choices[0].message.content or ""
            azure_text_moderation_response = await self.async_make_request(
                text=content,
            )
            self.check_severity_threshold(response=azure_text_moderation_response)
        return response

    async def async_post_call_streaming_hook(
        self, user_api_key_dict: UserAPIKeyAuth, response: str
    ) -> Any:
        try:
            if response is not None and len(response) > 0:
                azure_text_moderation_response = await self.async_make_request(
                    text=response,
                )
                self.check_severity_threshold(response=azure_text_moderation_response)
            return response
        except HTTPException as e:
            import json

            error_returned = json.dumps({"error": e.detail})
            return f"data: {error_returned}\n\n"
