"""Gray Swan Cygnal guardrail integration."""

import os
from typing import Any, Dict, Literal, Optional, 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.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
    add_guardrail_to_applied_guardrails_header,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import Choices, LLMResponseTypes, ModelResponse


class GraySwanGuardrailMissingSecrets(Exception):
    """Raised when the Gray Swan API key is missing."""


class GraySwanGuardrailAPIError(Exception):
    """Raised when the Gray Swan API returns an error."""


class GraySwanGuardrail(CustomGuardrail):
    """
    Guardrail that calls Gray Swan's Cygnal monitoring endpoint.

    see: https://docs.grayswan.ai/cygnal/monitor-requests
    """

    SUPPORTED_ON_FLAGGED_ACTIONS = {"block", "monitor", "passthrough"}
    DEFAULT_ON_FLAGGED_ACTION = "monitor"
    BASE_API_URL = "https://api.grayswan.ai"
    MONITOR_PATH = "/cygnal/monitor"
    SUPPORTED_REASONING_MODES = {"off", "hybrid", "thinking"}

    def __init__(
        self,
        guardrail_name: Optional[str] = "grayswan",
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        on_flagged_action: Optional[str] = None,
        violation_threshold: Optional[float] = None,
        reasoning_mode: Optional[str] = None,
        categories: Optional[Dict[str, str]] = None,
        policy_id: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        self.async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )

        api_key_value = api_key or os.getenv("GRAYSWAN_API_KEY")
        if not api_key_value:
            raise GraySwanGuardrailMissingSecrets(
                "Gray Swan API key missing. Set `GRAYSWAN_API_KEY` or pass `api_key`."
            )
        self.api_key: str = api_key_value

        base = api_base or os.getenv("GRAYSWAN_API_BASE") or self.BASE_API_URL
        self.api_base = base.rstrip("/")
        self.monitor_url = f"{self.api_base}{self.MONITOR_PATH}"

        action = on_flagged_action
        if action and action.lower() in self.SUPPORTED_ON_FLAGGED_ACTIONS:
            self.on_flagged_action = action.lower()
        else:
            if action:
                verbose_proxy_logger.warning(
                    "Gray Swan Guardrail: Unsupported on_flagged_action '%s', defaulting to '%s'.",
                    action,
                    self.DEFAULT_ON_FLAGGED_ACTION,
                )
            self.on_flagged_action = self.DEFAULT_ON_FLAGGED_ACTION

        self.violation_threshold = self._resolve_threshold(violation_threshold)
        self.reasoning_mode = self._resolve_reasoning_mode(reasoning_mode)
        self.categories = categories
        self.policy_id = policy_id

        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,
        )

    # ------------------------------------------------------------------
    # Guardrail hook entry points
    # ------------------------------------------------------------------

    @log_guardrail_information
    async def async_pre_call_hook(
        self,
        user_api_key_dict: UserAPIKeyAuth,
        cache,
        data: dict,
        call_type: Literal[
            "completion",
            "text_completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "pass_through_endpoint",
            "rerank",
            "mcp_call",
            "anthropic_messages",
        ],
    ) -> Optional[Union[Exception, str, dict]]:
        if (
            self.should_run_guardrail(
                data=data, event_type=GuardrailEventHooks.pre_call
            )
            is not True
        ):
            return data

        verbose_proxy_logger.debug("Gray Swan Guardrail: pre-call hook triggered")

        messages = data.get("messages")
        if not messages:
            verbose_proxy_logger.debug("Gray Swan Guardrail: No messages in data")
            return data

        dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}

        payload = self._prepare_payload(messages, dynamic_body)
        if payload is None:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: no content to scan; skipping request"
            )
            return data

        await self.run_grayswan_guardrail(payload, data, GuardrailEventHooks.pre_call)
        add_guardrail_to_applied_guardrails_header(
            request_data=data, guardrail_name=self.guardrail_name
        )
        return data

    @log_guardrail_information
    async def async_moderation_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        call_type: Literal[
            "completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
            "responses",
            "mcp_call",
            "anthropic_messages",
        ],
    ) -> Optional[Union[Exception, str, dict]]:
        if (
            self.should_run_guardrail(
                data=data, event_type=GuardrailEventHooks.during_call
            )
            is not True
        ):
            return data

        verbose_proxy_logger.debug("GraySwan Guardrail: during-call hook triggered")

        messages = data.get("messages")
        if not messages:
            verbose_proxy_logger.debug("Gray Swan Guardrail: No messages in data")
            return data

        dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}

        payload = self._prepare_payload(messages, dynamic_body)
        if payload is None:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: no content to scan; skipping request"
            )
            return data

        await self.run_grayswan_guardrail(
            payload, data, GuardrailEventHooks.during_call
        )
        add_guardrail_to_applied_guardrails_header(
            request_data=data, guardrail_name=self.guardrail_name
        )
        return data

    @log_guardrail_information
    async def async_post_call_success_hook(
        self,
        data: dict,
        user_api_key_dict: UserAPIKeyAuth,
        response: LLMResponseTypes,
    ) -> LLMResponseTypes:
        if (
            self.should_run_guardrail(
                data=data, event_type=GuardrailEventHooks.post_call
            )
            is not True
        ):
            return response

        verbose_proxy_logger.debug("GraySwan Guardrail: post-call hook triggered")

        response_dict = response.model_dump() if hasattr(response, "model_dump") else {}  # type: ignore[union-attr]
        response_messages = [
            msg if isinstance(msg, dict) else msg.model_dump()
            for choice in response_dict.get("choices", [])
            if isinstance(choice, dict)
            for msg in [choice.get("message")]
            if msg is not None
        ]

        if not response_messages:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: no response messages detected; skipping post-call scan"
            )
            return response

        dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}

        payload = self._prepare_payload(response_messages, dynamic_body)
        if payload is None:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: no content to scan; skipping request"
            )
            return response

        await self.run_grayswan_guardrail(payload, data, GuardrailEventHooks.post_call)

        # If passthrough mode and detection info exists, replace response content with violation message
        if self.on_flagged_action == "passthrough" and "metadata" in data:
            guardrail_detections = data.get("metadata", {}).get(
                "guardrail_detections", []
            )
            if guardrail_detections:
                # Replace the model response content with guardrail violation message
                violation_message = self._format_violation_message(
                    guardrail_detections, is_output=True
                )

                # Handle ModelResponse (OpenAI-style chat/text completions)
                # Use isinstance to narrow the type for mypy
                if isinstance(response, ModelResponse) and response.choices:
                    verbose_proxy_logger.debug(
                        "Gray Swan Guardrail: Replacing response content in ModelResponse format"
                    )
                    for choice in response.choices:
                        # Handle chat completion format (message.content)
                        # Choices has message attribute, StreamingChoices has delta
                        if isinstance(choice, Choices) and hasattr(choice, "message") and hasattr(
                            choice.message, "content"
                        ):
                            choice.message.content = violation_message
                        # Handle text completion format (text)
                        # Text attribute might be set dynamically, use setattr
                        elif hasattr(choice, "text"):
                            setattr(choice, "text", violation_message)

                        # Update finish_reason to indicate content filtering
                        if hasattr(choice, "finish_reason"):
                            choice.finish_reason = "content_filter"

                # Handle AnthropicMessagesResponse format
                elif hasattr(response, "content") and isinstance(response.content, list):  # type: ignore
                    verbose_proxy_logger.debug(
                        "Gray Swan Guardrail: Replacing response content in Anthropic Messages format"
                    )
                    # Replace content blocks with text block containing violation message
                    response.content = [  # type: ignore
                        {"type": "text", "text": violation_message}
                    ]
                    # Update stop_reason if present
                    if hasattr(response, "stop_reason"):
                        response.stop_reason = "end_turn"  # type: ignore

                else:
                    verbose_proxy_logger.warning(
                        "Gray Swan Guardrail: Passthrough mode enabled but response format not recognized. "
                        "Cannot replace content. Response type: %s",
                        type(response).__name__,
                    )

        add_guardrail_to_applied_guardrails_header(
            request_data=data, guardrail_name=self.guardrail_name
        )
        return response

    # ------------------------------------------------------------------
    # Core GraySwan interaction
    # ------------------------------------------------------------------

    async def run_grayswan_guardrail(
        self,
        payload: dict,
        data: Optional[dict] = None,
        hook_type: Optional[GuardrailEventHooks] = None,
    ):
        headers = self._prepare_headers()

        try:
            response = await self.async_handler.post(
                url=self.monitor_url,
                headers=headers,
                json=payload,
                timeout=30.0,
            )
            response.raise_for_status()
            result = response.json()
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: monitor response %s", safe_dumps(result)
            )
        except HTTPException:
            raise
        except Exception as exc:  # pragma: no cover - depends on HTTP client behaviour
            verbose_proxy_logger.exception(
                "Gray Swan Guardrail: API request failed: %s", exc
            )
            raise GraySwanGuardrailAPIError(str(exc)) from exc

        self._process_grayswan_response(result, data, hook_type)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _prepare_headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "grayswan-api-key": self.api_key,
        }

    def _prepare_payload(
        self, messages: list[dict], dynamic_body: dict
    ) -> Optional[Dict[str, Any]]:
        payload: Dict[str, Any] = {}
        payload["messages"] = messages

        categories = dynamic_body.get("categories") or self.categories
        if categories:
            payload["categories"] = categories

        policy_id = dynamic_body.get("policy_id") or self.policy_id
        if policy_id:
            payload["policy_id"] = policy_id

        reasoning_mode = dynamic_body.get("reasoning_mode") or self.reasoning_mode
        if reasoning_mode:
            payload["reasoning_mode"] = reasoning_mode

        return payload

    def _process_grayswan_response(
        self,
        response_json: Dict[str, Any],
        data: Optional[dict] = None,
        hook_type: Optional[GuardrailEventHooks] = None,
    ) -> None:
        violation_score = float(response_json.get("violation", 0.0) or 0.0)
        violated_rules = response_json.get("violated_rules", [])
        mutation_detected = response_json.get("mutation")
        ipi_detected = response_json.get("ipi")

        flagged = violation_score >= self.violation_threshold
        if not flagged:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: request passed (score=%s, rules=%s)",
                violation_score,
                violated_rules,
            )
            return

        verbose_proxy_logger.warning(
            "Gray Swan Guardrail: violation score %.3f exceeds threshold %.3f",
            violation_score,
            self.violation_threshold,
        )

        if self.on_flagged_action == "block":
            # Determine if violation was in input or output
            violation_location = (
                "output"
                if hook_type == GuardrailEventHooks.post_call
                else "input"
            )
            raise HTTPException(
                status_code=400,
                detail={
                    "error": "Blocked by Gray Swan Guardrail",
                    "violation_location": violation_location,
                    "violation": violation_score,
                    "violated_rules": violated_rules,
                    "mutation": mutation_detected,
                    "ipi": ipi_detected,
                },
            )
        elif self.on_flagged_action == "monitor":
            verbose_proxy_logger.info(
                "Gray Swan Guardrail: Monitoring mode - allowing flagged content to proceed"
            )
        elif self.on_flagged_action == "passthrough":
            # Store detection info
            detection_info = {
                "guardrail": "grayswan",
                "flagged": True,
                "violation_score": violation_score,
                "violated_rules": violated_rules,
                "mutation": mutation_detected,
                "ipi": ipi_detected,
            }

            # For pre_call and during_call, raise exception to short-circuit LLM call
            if hook_type in (
                GuardrailEventHooks.pre_call,
                GuardrailEventHooks.during_call,
            ):
                verbose_proxy_logger.info(
                    "Gray Swan Guardrail: Passthrough mode - raising exception to short-circuit LLM call"
                )
                violation_message = self._format_violation_message(
                    [detection_info], is_output=False
                )
                self.raise_passthrough_exception(
                    violation_message=violation_message,
                    request_data=data or {},
                    detection_info=detection_info,
                )

            # For post_call, store in metadata to replace response later
            verbose_proxy_logger.info(
                "Gray Swan Guardrail: Passthrough mode - storing detection info in metadata"
            )
            if data is not None:
                if "metadata" not in data:
                    data["metadata"] = {}
                if "guardrail_detections" not in data["metadata"]:
                    data["metadata"]["guardrail_detections"] = []
                data["metadata"]["guardrail_detections"].append(detection_info)

    def _format_violation_message(
        self, guardrail_detections: list, is_output: bool = False
    ) -> str:
        """
        Format guardrail detections into a user-friendly violation message.

        Args:
            guardrail_detections: List of detection info dictionaries
            is_output: True if violation is in model output (post_call), False if in input (pre_call/during_call)

        Returns:
            Formatted violation message string
        """
        if not guardrail_detections:
            return "Content was flagged by guardrail"

        # Get the most recent detection (should be from this guardrail)
        detection = guardrail_detections[-1]

        violation_score = detection.get("violation_score", 0.0)
        violated_rules = detection.get("violated_rules", [])
        mutation = detection.get("mutation", False)
        ipi = detection.get("ipi", False)

        # Indicate whether violation was in input or output
        violation_location = "the model response" if is_output else "input query"

        message_parts = [
            f"Sorry I can't help with that. According to the Gray Swan Cygnal Guardrail, the {violation_location} has a violation score of {violation_score:.2f}.",
        ]

        if violated_rules:
            message_parts.append(
                f"It was violating the rule(s): {', '.join(map(str, violated_rules))}."
            )

        if mutation:
            message_parts.append(
                "Mutation effort to make the harmful intention disguised was DETECTED."
            )

        if ipi:
            message_parts.append("Indirect Prompt Injection was DETECTED.")

        return "\n".join(message_parts)

    def _resolve_threshold(self, threshold: Optional[float]) -> float:
        if threshold is not None:
            return min(max(threshold, 0.0), 1.0)
        return 0.5

    def _resolve_reasoning_mode(self, candidate: Optional[str]) -> Optional[str]:
        if candidate is None:
            return None
        normalised = candidate.strip().lower()
        if normalised in self.SUPPORTED_REASONING_MODES:
            return normalised
        verbose_proxy_logger.warning(
            "Gray Swan Guardrail: ignoring unsupported reasoning_mode '%s'",
            candidate,
        )
        return None

    @staticmethod
    def get_config_model():
        from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
            GraySwanGuardrailConfigModel,
        )

        return GraySwanGuardrailConfigModel
