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

import os
import time
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol

from fastapi import HTTPException
from typing_extensions import NotRequired, ReadOnly, TypedDict

from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
    CustomGuardrail,
    ModifyResponseException,
    log_guardrail_information,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
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

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj

GRAYSWAN_BLOCK_ERROR_MSG: Final = "Blocked by Gray Swan Guardrail"


class _GraySwanMonitorResponse(TypedDict):
    """Body returned by Gray Swan's `/cygnal/monitor` endpoint."""

    violation: ReadOnly[NotRequired[float | None]]
    violated_rules: ReadOnly[NotRequired[list[object]]]
    violated_rule_descriptions: ReadOnly[NotRequired[list[object]]]
    mutation: ReadOnly[NotRequired[bool | None]]
    ipi: ReadOnly[NotRequired[bool | None]]


class _GraySwanMonitorHTTPResponse(Protocol):
    def raise_for_status(self) -> object: ...

    def json(self) -> _GraySwanMonitorResponse: ...


class _GraySwanMonitorHTTPClient(Protocol):
    async def post(
        self,
        *,
        url: str,
        headers: dict[str, str],
        json: dict[str, object],
        timeout: float,
    ) -> _GraySwanMonitorHTTPResponse: ...


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


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

    def __init__(self, message: str, status_code: int | None = None) -> None:
        super().__init__(message)
        self.status_code = status_code


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

    Uses the unified guardrail system via `apply_guardrail` method,
    which automatically works with all LiteLLM endpoints:
    - OpenAI Chat Completions
    - OpenAI Responses API
    - OpenAI Text Completions
    - Anthropic Messages
    - Image Generation
    - And more...

    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: str | None = "grayswan",
        api_key: str | None = None,
        api_base: str | None = None,
        on_flagged_action: str | None = None,
        violation_threshold: float | None = None,
        reasoning_mode: str | None = None,
        categories: dict[str, str] | None = None,
        policy_id: str | None = None,
        streaming_end_of_stream_only: bool = False,
        streaming_sampling_rate: int = 5,
        fail_open: bool | None = True,
        guardrail_timeout: float | None = 30.0,
        **kwargs: Any,
    ) -> None:
        self.async_handler: _GraySwanMonitorHTTPClient = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )

        api_key_value: Final = 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: Final = 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: Final = 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
        self.fail_open = True if fail_open is None else bool(fail_open)
        self.guardrail_timeout = 30.0 if guardrail_timeout is None else float(guardrail_timeout)

        # Streaming configuration
        self.streaming_end_of_stream_only = streaming_end_of_stream_only
        self.streaming_sampling_rate = streaming_sampling_rate

        verbose_proxy_logger.debug(
            "GraySwan __init__: streaming_end_of_stream_only=%s, streaming_sampling_rate=%s",
            streaming_end_of_stream_only,
            streaming_sampling_rate,
        )

        super().__init__(
            guardrail_name=guardrail_name,
            supported_event_hooks=list(self.get_supported_event_hooks()),
            **kwargs,
        )

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

    # ------------------------------------------------------------------
    # Debug override to trace post_call issues
    # ------------------------------------------------------------------

    def should_run_guardrail(self, data, event_type) -> bool:
        """Override to add debug logging."""
        result: Final = super().should_run_guardrail(data, event_type)
        # Check if apply_guardrail is in __dict__
        has_apply_guardrail: Final = "apply_guardrail" in type(self).__dict__
        verbose_proxy_logger.debug(
            "GraySwan DEBUG: should_run_guardrail event_type=%s, result=%s, event_hook=%s, has_apply_guardrail=%s, class=%s",
            event_type,
            result,
            self.event_hook,
            has_apply_guardrail,
            type(self).__name__,
        )
        return result

    # ------------------------------------------------------------------
    # Unified Guardrail Interface (works with ALL endpoints automatically)
    # ------------------------------------------------------------------

    @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 Gray Swan guardrail to extracted text content.

        This method is called by the unified guardrail system which handles
        extracting text from any request format (OpenAI, Anthropic, etc.).

        Args:
            inputs: Dictionary containing:
                - texts: List of texts to scan
                - images: Optional list of images (not currently used by GraySwan)
                - tool_calls: Optional list of tool calls (not currently used)
            request_data: The original request data
            input_type: "request" for pre-call, "response" for post-call
            logging_obj: Optional logging object

        Returns:
            GenericGuardrailAPIInputs - texts may be replaced with violation message in passthrough mode

        Raises:
            HTTPException: If content is blocked (block mode)
            Exception: If guardrail check fails
        """
        # DEBUG: Log when apply_guardrail is called
        verbose_proxy_logger.debug(
            "GraySwan DEBUG: apply_guardrail called with input_type=%s, texts=%s",
            input_type,
            inputs.get("texts", [])[:100] if inputs.get("texts") else "NONE",
        )

        texts: Final = inputs.get("texts", [])
        if not texts:
            verbose_proxy_logger.debug("Gray Swan Guardrail: No texts to scan")
            return inputs

        verbose_proxy_logger.debug(
            "Gray Swan Guardrail: Scanning %d text(s) for %s",
            len(texts),
            input_type,
        )

        # Convert texts to messages format for GraySwan API
        # Use "user" role for request content, "assistant" for response content
        role: Final = "assistant" if input_type == "response" else "user"
        messages: Final = [{"role": role, "content": text} for text in texts]

        # Get dynamic params from request metadata
        dynamic_body: Final = self.get_guardrail_dynamic_request_body_params(request_data) or {}
        if dynamic_body:
            verbose_proxy_logger.debug("Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body))

        # Prepare and send payload
        payload: Final = self._prepare_payload(messages, dynamic_body, request_data, logging_obj)
        if payload is None:
            return inputs

        start_time: Final = time.time()
        try:
            response_json: Final = await self._call_grayswan_api(payload)
            is_output: Final = input_type == "response"
            result: Final = self._process_response_internal(
                response_json=response_json,
                request_data=request_data,
                inputs=inputs,
                is_output=is_output,
            )
            return result
        except Exception as exc:
            if self._is_grayswan_exception(exc):
                raise
            end_time: Final = time.time()
            status_code: Final = getattr(exc, "status_code", None) or getattr(exc, "exception_status_code", None)
            self._log_guardrail_failure(
                exc=exc,
                request_data=request_data or {},
                start_time=start_time,
                end_time=end_time,
                status_code=status_code,
            )
            if self.fail_open:
                verbose_proxy_logger.warning(
                    "Gray Swan Guardrail: fail_open=True. Allowing request to proceed despite error: %s",
                    exc,
                )
                return inputs
            if isinstance(exc, GraySwanGuardrailAPIError):
                raise exc
            raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc

    def _is_grayswan_exception(self, exc: Exception) -> bool:
        # Guardrail decision (passthrough) should always propagate,
        # regardless of fail_open.
        if isinstance(exc, ModifyResponseException):
            return True
        detail: Final = getattr(exc, "detail", None)
        if isinstance(detail, dict):
            return detail.get("error") == GRAYSWAN_BLOCK_ERROR_MSG
        return False

    # ------------------------------------------------------------------
    # Legacy Test Interface (for backward compatibility)
    # ------------------------------------------------------------------

    async def run_grayswan_guardrail(self, payload: dict[str, object]) -> _GraySwanMonitorResponse:
        """
        Run the GraySwan guardrail on a payload.

        This is a legacy method for testing purposes.

        Args:
            payload: The payload to scan

        Returns:
            Dict containing the GraySwan API response
        """
        response_json: Final = await self._call_grayswan_api(payload)
        # Call the legacy response processor (for test compatibility)
        self._process_grayswan_response(response_json)
        return response_json

    def _process_grayswan_response(
        self,
        response_json: _GraySwanMonitorResponse,
        data: dict | None = None,
        hook_type: GuardrailEventHooks | None = None,
    ) -> None:
        """
        Legacy method for processing GraySwan API responses.

        This method is maintained for backward compatibility with existing tests.
        It handles the test scenarios where responses need to be processed with
        knowledge of the request context (pre/during/post call hooks).

        Args:
            response_json: Response from GraySwan API
            data: Optional request data (for passthrough exceptions)
            hook_type: Optional GuardrailEventHooks for determining behavior
        """
        violation_score: Final = float(response_json.get("violation", 0.0) or 0.0)
        violated_rules: Final = response_json.get("violated_rules", [])
        mutation_detected: Final = response_json.get("mutation")
        ipi_detected: Final = response_json.get("ipi")

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

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

        detection_info: Final = {
            "guardrail": "grayswan",
            "flagged": True,
            "violation_score": violation_score,
            "violated_rules": violated_rules,
            "mutation": mutation_detected,
            "ipi": ipi_detected,
        }

        # Determine if this is input (pre-call/during-call) or output (post-call)
        if hook_type is not None:
            is_input = hook_type in [
                GuardrailEventHooks.pre_call,
                GuardrailEventHooks.during_call,
            ]
        else:
            is_input = True

        if self.on_flagged_action == "block":
            violation_location: Final = "output" if (not is_input) else "input"
            raise HTTPException(
                status_code=400,
                detail={
                    "error": GRAYSWAN_BLOCK_ERROR_MSG,
                    "violation_location": violation_location,
                    "violation": violation_score,
                    "violated_rules": violated_rules,
                    "mutation": mutation_detected,
                    "ipi": ipi_detected,
                },
            )
        elif self.on_flagged_action == "passthrough":
            # For passthrough mode, we need to handle violations
            detections: Final = [detection_info]
            violation_message: Final = self._format_violation_message(detections, is_output=not is_input)
            verbose_proxy_logger.info("Gray Swan Guardrail: Passthrough mode - handling violation")

            # If hook_type is provided and in pre/during call, raise exception
            if hook_type in [
                GuardrailEventHooks.pre_call,
                GuardrailEventHooks.during_call,
            ]:
                # Raise ModifyResponseException to short-circuit LLM call
                if data is None:
                    data = {}
                self.raise_passthrough_exception(
                    violation_message=violation_message,
                    request_data=data,
                    detection_info=detection_info,
                )
            elif hook_type == GuardrailEventHooks.post_call:
                # For post-call, store detection info in metadata
                if data is None:
                    data = {}
                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)

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

    async def _call_grayswan_api(self, payload: dict[str, object]) -> _GraySwanMonitorResponse:
        """Call the GraySwan monitoring API."""
        headers: Final = self._prepare_headers()

        try:
            response: Final = await self.async_handler.post(
                url=self.monitor_url,
                headers=headers,
                json=payload,
                timeout=self.guardrail_timeout,
            )
            response.raise_for_status()
            result: Final = response.json()
            verbose_proxy_logger.debug("Gray Swan Guardrail: monitor response %s", safe_dumps(result))
            return result
        except Exception as exc:
            status_code: Final = getattr(exc, "status_code", None) or getattr(exc, "exception_status_code", None)
            raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc

    def _process_response_internal(
        self,
        response_json: _GraySwanMonitorResponse,
        request_data: dict,
        inputs: GenericGuardrailAPIInputs,
        is_output: bool,
    ) -> GenericGuardrailAPIInputs:
        """
        Process GraySwan API response and handle violations.

        Args:
            response_json: Response from GraySwan API
            request_data: Original request data
            inputs: The inputs being scanned
            is_output: True if scanning model output, False for input

        Returns:
            GenericGuardrailAPIInputs - possibly modified with violation message

        Raises:
            HTTPException: If content is blocked (block mode)
        """
        violation_score: Final = float(response_json.get("violation", 0.0) or 0.0)
        violated_rules: Final = response_json.get("violated_rule_descriptions", [])
        mutation_detected: Final = response_json.get("mutation")
        ipi_detected: Final = response_json.get("ipi")

        flagged: Final = violation_score >= self.violation_threshold
        if not flagged:
            verbose_proxy_logger.debug(
                "Gray Swan Guardrail: content passed (score=%s, threshold=%s)",
                violation_score,
                self.violation_threshold,
            )
            return inputs

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

        detection_info: Final = {
            "guardrail": "grayswan",
            "flagged": True,
            "violation_score": violation_score,
            "violated_rules": violated_rules,
            "mutation": mutation_detected,
            "ipi": ipi_detected,
        }

        if self.on_flagged_action == "block":
            violation_location: Final = "output" if is_output else "input"
            raise HTTPException(
                status_code=400,
                detail={
                    "error": GRAYSWAN_BLOCK_ERROR_MSG,
                    "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")
            return inputs
        elif self.on_flagged_action == "passthrough":
            # Replace content with violation message
            violation_message: Final = self._format_violation_message(detection_info, is_output=is_output)
            verbose_proxy_logger.info(
                "Gray Swan Guardrail: Passthrough mode - replacing content with violation message"
            )

            if not is_output:
                # For pre-call (request), raise exception to short-circuit LLM call
                # and return synthetic response with violation message
                self.raise_passthrough_exception(
                    violation_message=violation_message,
                    request_data=request_data,
                    detection_info=detection_info,
                )

            # For post-call (response), replace texts and let unified system apply them
            inputs["texts"] = [violation_message]
            return inputs

        return inputs

    # ------------------------------------------------------------------
    # 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 _extract_inbound_headers(
        self,
        request_data: dict,
        logging_obj: Optional["LiteLLMLoggingObj"] = None,
    ) -> dict[str, str] | None:
        headers = (request_data.get("proxy_server_request") or {}).get("headers")
        if not headers:
            headers = request_data.get("headers")
        if not headers:
            headers = (request_data.get("metadata") or {}).get("headers")
        if not headers and logging_obj and getattr(logging_obj, "model_call_details", None):
            headers = (
                (logging_obj.model_call_details or {}).get("litellm_params", {}).get("metadata", {}).get("headers")
            )
        if not isinstance(headers, dict):
            return None

        forwarded_header_names: Final = ("shade_scan_id",)
        forwarded_headers: Final = {}
        for key, value in headers.items():
            if str(key).lower() in forwarded_header_names:
                forwarded_headers[str(key)] = str(value)
        return forwarded_headers or None

    def _prepare_payload(
        self,
        messages: list[dict[str, str]],
        dynamic_body: dict,
        request_data: dict,
        logging_obj: Optional["LiteLLMLoggingObj"] = None,
    ) -> dict[str, object] | None:
        payload: Final[dict[str, object]] = {"messages": messages}

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

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

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

        # Pass through arbitrary metadata when provided via dynamic extra_body.
        if "metadata" in dynamic_body:
            payload["metadata"] = dynamic_body["metadata"]

        inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj)

        litellm_metadata: Final = request_data.get("litellm_metadata")
        cleaned_litellm_metadata: Final = dict(litellm_metadata) if isinstance(litellm_metadata, dict) else {}
        if inbound_headers:
            existing_headers: Final = cleaned_litellm_metadata.get("headers")
            cleaned_litellm_metadata["headers"] = (
                {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers
            )
        if cleaned_litellm_metadata:
            sanitized: Final[object] = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={})
            if isinstance(sanitized, dict) and sanitized:
                payload["litellm_metadata"] = sanitized

        return payload

    def _format_violation_message(self, detection_info: object, is_output: bool = False) -> str:
        """
        Format detection info into a user-friendly violation message.

        Args:
            detection_info: Can be either:
                - A single dict with violation_score, violated_rules, mutation, ipi keys
                - A list of such dicts (legacy format)
            is_output: True if violation is in model output, False if in input

        Returns:
            Formatted violation message string
        """
        # Handle legacy format where detection_info is a list
        if isinstance(detection_info, list) and len(detection_info) > 0:
            detection_info = detection_info[0]

        # Extract fields from detection_info dict
        detection_dict: Final[dict] = detection_info if isinstance(detection_info, dict) else {}
        violation_score: Final = detection_dict.get("violation_score", 0.0)
        violated_rules: Final = detection_dict.get("violated_rules", [])
        mutation: Final = detection_dict.get("mutation", False)
        ipi: Final = detection_dict.get("ipi", False)

        violation_location: Final = "the model response" if is_output else "input query"

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

        if violated_rules:
            formatted_rules: Final = self._format_violated_rules(violated_rules)
            if formatted_rules:
                message_parts.append(f"It was violating the rule(s): {formatted_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 _format_violated_rules(self, violated_rules: list) -> str:
        """Format violated rules list into a readable string."""
        formatted: Final[list[str]] = []
        for rule in violated_rules:
            if isinstance(rule, dict):
                # New format: {'rule': 6, 'name': 'Illegal Activities...', 'description': '...'}
                rule_num = rule.get("rule", "")
                rule_name = rule.get("name", "")
                rule_desc = rule.get("description", "")
                if rule_num and rule_name:
                    if rule_desc:
                        formatted.append(f"#{rule_num} {rule_name}: {rule_desc}")
                    else:
                        formatted.append(f"#{rule_num} {rule_name}")
                elif rule_name:
                    formatted.append(rule_name)
                else:
                    formatted.append(str(rule))
            else:
                # Legacy format: simple value
                formatted.append(str(rule))

        return ", ".join(formatted)

    def _resolve_threshold(self, value: float | None) -> float:
        if value is not None:
            return float(value)
        env_val: Final = os.getenv("GRAYSWAN_VIOLATION_THRESHOLD")
        if env_val:
            try:
                return float(env_val)
            except ValueError:
                pass
        return 0.5

    def _resolve_reasoning_mode(self, value: str | None) -> str | None:
        if value and value.lower() in self.SUPPORTED_REASONING_MODES:
            return value.lower()
        env_val: Final = os.getenv("GRAYSWAN_REASONING_MODE")
        if env_val and env_val.lower() in self.SUPPORTED_REASONING_MODES:
            return env_val.lower()
        return None

    def _log_guardrail_failure(
        self,
        exc: Exception,
        request_data: dict,
        start_time: float,
        end_time: float,
        status_code: int | None = None,
    ) -> None:
        """Log guardrail failure and attach standard logging metadata."""
        try:
            self.add_standard_logging_guardrail_information_to_request_data(
                guardrail_json_response=str(exc),
                request_data=request_data,
                guardrail_status="guardrail_failed_to_respond",
                start_time=start_time,
                end_time=end_time,
                duration=end_time - start_time,
                guardrail_provider="grayswan",
            )
        except Exception:
            verbose_proxy_logger.exception(
                "Gray Swan Guardrail: failed to log guardrail failure for error: %s",
                exc,
            )
        verbose_proxy_logger.error(
            "Gray Swan Guardrail: API request failed%s: %s",
            f" (status_code={status_code})" if status_code else "",
            exc,
        )
