# +-------------------------------------------------------------+
#
#           Use Aim Security Guardrails for your LLM calls
#                   https://www.aim.security/
#
# +-------------------------------------------------------------+
import asyncio
import json
import os
from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeAlias

from pydantic import BaseModel
from typing_extensions import NotRequired, ReadOnly, TypedDict
from websockets.asyncio.client import ClientConnection, connect

from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import (
    apply_redacted_messages_back,
    build_inspection_messages,
    has_non_string_content,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
    CallTypesLiteral,
    Choices,
    LLMResponseTypes,
    ModelResponse,
    ModelResponseStream,
)

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


class AimGuardrailMissingSecrets(Exception):
    pass


class AimRequiredAction(TypedDict):
    """The ``required_action`` block of an Aim ``/fw/v1/analyze`` response."""

    action_type: ReadOnly[NotRequired[str]]
    detection_message: ReadOnly[str]


class AimAnalysisResult(TypedDict):
    """The ``analysis_result`` block of an Aim ``/fw/v1/analyze`` response."""

    policy_drill_down: ReadOnly[Mapping[str, object]]


class AimRedactedMessage(TypedDict):
    """One entry of Aim's ``redacted_chat.all_redacted_messages``."""

    role: ReadOnly[str]
    content: ReadOnly[str]


class AimRedactedChat(TypedDict):
    """The ``redacted_chat`` block of an Aim ``/fw/v1/analyze`` response."""

    all_redacted_messages: ReadOnly[Sequence[AimRedactedMessage]]


class AimAnalyzeResponse(TypedDict):
    """Body returned by Aim's ``POST /fw/v1/analyze``."""

    required_action: ReadOnly[AimRequiredAction]
    analysis_result: ReadOnly[AimAnalysisResult]
    redacted_chat: ReadOnly[NotRequired[AimRedactedChat]]


class AimOutputGuardrailResult(TypedDict, total=False):
    """Outcome of inspecting one model completion with Aim."""

    detection_message: ReadOnly[str]
    redacted_output: ReadOnly[str]


class AimStreamMessage(TypedDict, total=False):
    """One frame of Aim's ``/fw/v1/analyze/stream`` websocket protocol."""

    verified_chunk: ReadOnly[Mapping[str, object]]
    done: ReadOnly[bool]
    blocking_message: ReadOnly[str]


AimStreamChunk: TypeAlias = BaseModel | Mapping[str, object] | str | bytes


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

    def __init__(self, api_key: str | None = None, api_base: str | None = None, **kwargs):
        kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
        ssl_verify: Final = kwargs.pop("ssl_verify", None)
        self.async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback,
            params={"ssl_verify": ssl_verify} if ssl_verify is not None else None,
        )
        self.api_key = api_key or os.environ.get("AIM_API_KEY")
        if not self.api_key:
            msg: Final = (
                "Couldn't get Aim api key, either set the `AIM_API_KEY` in the environment or "
                "pass it as a parameter to the guardrail in the config file"
            )
            raise AimGuardrailMissingSecrets(msg)
        self.api_base = api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security"
        self.ws_api_base = self.api_base.replace("http://", "ws://").replace("https://", "wss://")
        self.dlp_entities: list[dict] = []
        self._max_dlp_entities = 100
        super().__init__(**kwargs)

    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache: DualCache,
        data: dict,
        call_type: CallTypesLiteral,
    ) -> Exception | str | dict | None:
        verbose_proxy_logger.debug("Inside AIM Pre-Call Hook")
        return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias)

    async def async_moderation_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        call_type: CallTypesLiteral,
    ) -> Exception | str | dict | None:
        verbose_proxy_logger.debug("Inside AIM Moderation Hook")

        await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias)
        return data

    async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict:
        user_email: Final = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
        call_id: Final = data.get("litellm_call_id")
        headers: Final = self._build_aim_headers(
            hook=hook,
            key_alias=key_alias,
            user_email=user_email,
            litellm_call_id=call_id,
        )
        response: Final = await self.async_handler.post(
            f"{self.api_base}/fw/v1/analyze",
            headers=headers,
            json={"messages": self._build_aim_inspection_messages(data)},
        )
        response.raise_for_status()
        res: Final[AimAnalyzeResponse] = response.json()
        required_action: Final = res.get("required_action")
        action_type: Final = required_action and required_action.get("action_type", None)
        if action_type is None:
            verbose_proxy_logger.debug("Aim: No required action specified")
            return data
        if action_type == "monitor_action":
            verbose_proxy_logger.info("Aim: monitor action")
        elif action_type == "block_action":
            self._handle_block_action(res["analysis_result"], required_action)
        elif action_type == "anonymize_action":
            return self._anonymize_request(res, data)
        else:
            verbose_proxy_logger.error("Aim: %s action", action_type)
        return data

    @staticmethod
    def _build_aim_inspection_messages(data: dict) -> list[dict[str, str]]:
        """AIM validates against the OpenAI chat schema. Bare ``role: "tool"``
        without ``tool_call_id`` and bare ``role: "function"`` without ``name``
        are rejected; the flatten drops those fields, so any role outside
        ``{system, user, assistant}`` collapses to ``user`` for the AIM POST."""
        safe_roles: Final = {"system", "user", "assistant"}
        return [{**m, "role": "user"} if m["role"] not in safe_roles else m for m in build_inspection_messages(data)]

    @staticmethod
    def _rejection(message: str, *, openai_code: str | None = None) -> ProxyException:
        return ProxyException(
            message=message,
            type="invalid_request_error",
            param=None,
            code=400,
            openai_code=openai_code,
        )

    def _handle_block_action(self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction) -> None:
        detection_message: Final = required_action.get("detection_message", None)
        verbose_proxy_logger.info(
            "Aim: Violation detected enabled policies: {policies}".format(
                policies=list(analysis_result["policy_drill_down"].keys()),
            ),
        )
        raise self._rejection(detection_message, openai_code="content_policy_violation")

    def _anonymize_request(self, res: AimAnalyzeResponse, data: dict) -> dict:
        verbose_proxy_logger.info("Aim: anonymize action")
        redacted_chat: Final = res.get("redacted_chat")
        if not redacted_chat:
            return data
        # Aim returns text-only redacted messages. Overwriting
        # ``data["messages"]`` with that would silently strip image/audio
        # parts from a multimodal request — degrade to block so the
        # multimodal payload is never silently rewritten.
        if has_non_string_content(data):
            raise self._rejection(
                "Aim: anonymize action requested for multimodal input "
                "but mask-in-place would drop non-text parts. Send the "
                "request with plain string content to use anonymize, "
                "or rely on block-mode policies."
            )
        redacted_messages: Final = [
            {
                "role": message["role"],
                "content": message["content"],
            }
            for message in redacted_chat["all_redacted_messages"]
        ]
        # Write back to ``messages`` AND ``input``. The Responses-API
        # backend reads ``input``; writing only to ``messages`` would let
        # unredacted text reach the LLM for ``/v1/responses`` calls.
        apply_redacted_messages_back(data, redacted_messages)
        return data

    async def call_aim_guardrail_on_output(
        self, request_data: dict, output: str, hook: str, key_alias: str | None
    ) -> AimOutputGuardrailResult | None:
        user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
        call_id: Final = request_data.get("litellm_call_id")
        response: Final = await self.async_handler.post(
            f"{self.api_base}/fw/v1/analyze",
            headers=self._build_aim_headers(
                hook=hook,
                key_alias=key_alias,
                user_email=user_email,
                litellm_call_id=call_id,
            ),
            json={
                "messages": self._build_aim_inspection_messages(request_data)
                + [{"role": "assistant", "content": output}]
            },
        )
        response.raise_for_status()
        res: Final[AimAnalyzeResponse] = response.json()
        required_action: Final = res.get("required_action")
        action_type: Final = required_action and required_action.get("action_type", None)
        if action_type and action_type == "block_action":
            return self._handle_block_action_on_output(res["analysis_result"], required_action)
        redacted_chat: Final = res.get("redacted_chat", None)

        if action_type and action_type == "anonymize_action" and redacted_chat:
            return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]}
        return {"redacted_output": output}

    def _handle_block_action_on_output(
        self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction
    ) -> AimOutputGuardrailResult | None:
        detection_message: Final = required_action.get("detection_message", None)
        verbose_proxy_logger.info(
            "Aim: detected: {detected}, enabled policies: {policies}".format(
                detected=True,
                policies=list(analysis_result["policy_drill_down"].keys()),
            ),
        )
        return {"detection_message": detection_message}

    def _build_aim_headers(
        self,
        *,
        hook: str,
        key_alias: str | None,
        user_email: str | None,
        litellm_call_id: str | None,
    ):
        """
        A helper function to build the http headers that are required by AIM guardrails.
        """
        return (
            {
                "Authorization": f"Bearer {self.api_key}",
                # Used by Aim to apply only the guardrails that should be applied in a specific request phase.
                "x-aim-litellm-hook": hook,
                # Used by Aim to track LiteLLM version and provide backward compatibility.
                "x-aim-litellm-version": litellm_version,
            }
            # Used by Aim to track together single call input and output
            | ({"x-aim-call-id": litellm_call_id} if litellm_call_id else {})
            # Used by Aim to track guardrails violations by user.
            | ({"x-aim-user-email": user_email} if user_email else {})
            | (
                {
                    # Used by Aim apply only the guardrails that are associated with the key alias.
                    "x-aim-gateway-key-alias": key_alias,
                }
                if key_alias
                else {}
            )
        )

    async def async_post_call_success_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        response: LLMResponseTypes,
    ) -> LLMResponseTypes:
        if not (isinstance(response, ModelResponse) and response.choices):
            return response
        # Inspect every choice — when ``n>1`` the additional completions
        # used to bypass Aim entirely because the hook only inspected
        # ``choices[0]``. Run inspections concurrently so multi-completion
        # responses don't pay an n× latency penalty.
        choices_to_inspect: Final = [c for c in response.choices if isinstance(c, Choices)]
        if not choices_to_inspect:
            return response
        # ``return_exceptions=True`` lets every inspection finish even if
        # one fails — without it, the first exception would propagate and
        # leave the remaining tasks running in the background.
        results: Final = await asyncio.gather(
            *(
                self.call_aim_guardrail_on_output(
                    data,
                    choice.message.content or "",
                    hook="output",
                    key_alias=user_api_key_dict.key_alias,
                )
                for choice in choices_to_inspect
            ),
            return_exceptions=True,
        )
        for choice, aim_output_guardrail_result in zip(choices_to_inspect, results):
            if isinstance(aim_output_guardrail_result, BaseException):
                raise aim_output_guardrail_result
            if aim_output_guardrail_result and (
                detection_message := aim_output_guardrail_result.get("detection_message")
            ):
                raise self._rejection(
                    detection_message,
                    openai_code="content_policy_violation",
                )
            if aim_output_guardrail_result and aim_output_guardrail_result.get("redacted_output"):
                choice.message.content = aim_output_guardrail_result.get("redacted_output")
        return response

    async def async_post_call_streaming_iterator_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        response: AsyncIterator[AimStreamChunk],
        request_data: dict,
    ) -> AsyncGenerator[ModelResponseStream, None]:
        user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
        call_id: Final = request_data.get("litellm_call_id")
        async with connect(
            f"{self.ws_api_base}/fw/v1/analyze/stream",
            additional_headers=self._build_aim_headers(
                hook="output",
                key_alias=user_api_key_dict.key_alias,
                user_email=user_email,
                litellm_call_id=call_id,
            ),
        ) as websocket:
            sender: Final = asyncio.create_task(self.forward_the_stream_to_aim(websocket, response))
            while True:
                result: AimStreamMessage = json.loads(await websocket.recv())
                if verified_chunk := result.get("verified_chunk"):
                    yield ModelResponseStream.model_validate(verified_chunk)
                else:
                    sender.cancel()
                    if result.get("done"):
                        return
                    if blocking_message := result.get("blocking_message"):
                        from litellm.proxy.proxy_server import StreamingCallbackError

                        raise StreamingCallbackError(blocking_message)
                    verbose_proxy_logger.error("Unknown message received from AIM: %s", result)
                    return

    async def forward_the_stream_to_aim(
        self,
        websocket: ClientConnection,
        response_iter: AsyncIterator[AimStreamChunk],
    ) -> None:
        async for chunk in response_iter:
            if isinstance(chunk, BaseModel):
                chunk = chunk.model_dump_json()
            if isinstance(chunk, dict):
                chunk = json.dumps(chunk)
            await websocket.send(chunk)
        await websocket.send(json.dumps({"done": True}))

    @staticmethod
    def get_config_model() -> type["GuardrailConfigModel"] | None:
        from litellm.types.proxy.guardrails.guardrail_hooks.aim import (
            AimGuardrailConfigModel,
        )

        return AimGuardrailConfigModel
