"""Ovalix guardrail integration: pre- and post-call checks via the Tracker service.

Use Ovalix Guardrails for your LLM calls. Supports pre_call (user input) and
post_call (model output) checkpoints with optional correction/blocking.
"""

import datetime
import hashlib
import os
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type

import httpx

from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
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.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs

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


BLOCKED_BY_OVALIX_FALLBACK_MESSAGE = "This message was blocked by Ovalix"
BLOCKED_ACTION_TYPE = "block"


class OvalixGuardrailMissingSecrets(Exception):
    """Raised when required Ovalix config (API base, key, application/checkpoint IDs) is missing."""

    pass


class OvalixGuardrailBlockedException(GuardrailRaisedException):
    """
    Raised when Ovalix blocks a message. Sets status_code=400 so the proxy
    returns 400 and HTTP clients do not retry (they retry on 5xx).
    """

    status_code = 400

    def __init__(
        self,
        guardrail_name: Optional[str] = None,
        message: str = "",
        should_wrap_with_default_message: bool = True,
    ):
        super().__init__(
            guardrail_name=guardrail_name,
            message=message,
            should_wrap_with_default_message=should_wrap_with_default_message,
        )


class OvalixGuardrail(CustomGuardrail):
    """
    Ovalix guardrail: pre-prompt (pre_call) and post-prompt (post_call) checks
    via the Tracker service, with application and checkpoint resolution from the
    Monolith backend.
    """

    def __init__(
        self,
        tracker_api_base: Optional[str] = None,
        tracker_api_key: Optional[str] = None,
        application_id: Optional[str] = None,
        pre_checkpoint_id: Optional[str] = None,
        post_checkpoint_id: Optional[str] = None,
        **kwargs: Any,
    ):
        self._tracker_api_base = tracker_api_base or os.environ.get(
            "OVALIX_TRACKER_API_BASE"
        )
        self._tracker_api_key = tracker_api_key or os.environ.get(
            "OVALIX_TRACKER_API_KEY"
        )
        self._application_id = application_id or os.environ.get("OVALIX_APPLICATION_ID")
        self._pre_checkpoint_id = pre_checkpoint_id or os.environ.get(
            "OVALIX_PRE_CHECKPOINT_ID"
        )
        self._post_checkpoint_id = post_checkpoint_id or os.environ.get(
            "OVALIX_POST_CHECKPOINT_ID"
        )

        if "supported_event_hooks" not in kwargs:
            kwargs["supported_event_hooks"] = []

        self._validate_config(kwargs["supported_event_hooks"])

        self._tracker_headers = httpx.Headers(
            {
                "Authorization": f"Bearer {self._tracker_api_key}",
                "Content-Type": "application/json",
            },
            encoding="utf-8",
        )

        self._async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback
        )

        super().__init__(**kwargs)
        verbose_proxy_logger.debug(
            "Ovalix Guardrail initialized: tracker=%s, application_id=%s, pre_checkpoint_id=%s, post_checkpoint_id=%s",
            self._tracker_api_base,
            self._application_id,
            self._pre_checkpoint_id,
            self._post_checkpoint_id,
        )

    def _validate_config(
        self, supported_event_hooks: List[GuardrailEventHooks]
    ) -> None:
        """Ensure required secrets and checkpoint IDs are set; auto-add hooks when IDs are present."""
        errors: List[str] = []

        if not self._tracker_api_base:
            errors.append(
                "Tracker API base, set OVALIX_TRACKER_API_BASE or pass tracker_api_base"
            )
        if not self._tracker_api_key:
            errors.append(
                "Tracker API key, set OVALIX_TRACKER_API_KEY or pass tracker_api_key"
            )
        if not self._application_id:
            errors.append(
                "Application ID, set OVALIX_APPLICATION_ID or pass application_id"
            )
        if (
            not self._pre_checkpoint_id
            and GuardrailEventHooks.pre_call in supported_event_hooks
        ):
            errors.append(
                "Pre-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or pass pre_checkpoint_id"
            )
        if (
            not self._post_checkpoint_id
            and GuardrailEventHooks.post_call in supported_event_hooks
        ):
            errors.append(
                "Post-checkpoint ID, set OVALIX_POST_CHECKPOINT_ID or pass post_checkpoint_id"
            )
        if not self._pre_checkpoint_id and not self._post_checkpoint_id:
            errors.append(
                "Pre-checkpoint ID or Post-checkpoint ID, set OVALIX_PRE_CHECKPOINT_ID or OVALIX_POST_CHECKPOINT_ID or pass pre_checkpoint_id or post_checkpoint_id"
            )

        if errors:
            raise OvalixGuardrailMissingSecrets(
                "Missing Ovalix guardrail configuration errors: " + ". ".join(errors)
            )

        # auto-add hooks when checkpoint IDs are present
        if (
            self._pre_checkpoint_id
            and GuardrailEventHooks.pre_call not in supported_event_hooks
        ):
            supported_event_hooks.append(GuardrailEventHooks.pre_call)
        if (
            self._post_checkpoint_id
            and GuardrailEventHooks.post_call not in supported_event_hooks
        ):
            supported_event_hooks.append(GuardrailEventHooks.post_call)

    def _get_actor(self, data: dict) -> str:
        """Return a stable actor identifier from request metadata (e.g. user email or id)."""
        metadata = data.get("metadata") or data.get("litellm_metadata") or {}
        if metadata.get("user_api_key_user_email"):
            return metadata["user_api_key_user_email"]
        if metadata.get("user_api_key_user_id"):
            return metadata["user_api_key_user_id"]
        return "unknown"

    def _get_tracker_actor_id(self, data: dict) -> str:
        """Normalize the actor string into a short, stable id for Tracker API payloads."""
        # NOTE: this hash is purely for normalization — it collapses an arbitrary actor
        # string (email, user id, or "unknown") into a compact, fixed-length, consistent
        # key. It is not a privacy/security measure and the actor value is not sensitive,
        # so a plain SHA-256 (truncated) is sufficient; no salting/KDF is needed here.
        actor_id = self._get_actor(data).encode()
        normalized_actor_id = hashlib.sha256(actor_id).hexdigest()[:8]
        return normalized_actor_id

    def _get_session_id(self, data: dict) -> str:
        """Return a unique identifier for the chat/session (actor + date + application_id)."""
        actor_hash = self._get_tracker_actor_id(data)
        today = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
        return f"{actor_hash}_{today}_{self._application_id}"

    async def _call_checkpoint(
        self,
        content: str,
        checkpoint_id: str,
        actor: str,
        session_id: str,
    ) -> Dict[str, Any]:
        """Call the Ovalix Tracker checkpoint API and return the JSON response."""
        application_id = self._application_id
        if not application_id or not checkpoint_id:
            raise ValueError("Ovalix: application_id or checkpoint_id not resolved")

        url = f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
        headers = dict(self._tracker_headers)
        payload = {
            "application_id": application_id,
            "checkpoint_id": checkpoint_id,
            "actor": actor,
            "session_id": session_id,
            "data_type": "TEXT",
            "data": {"content": content},
        }
        response = await self._async_handler.post(url, headers=headers, json=payload)
        response.raise_for_status()
        return response.json()

    @log_guardrail_information
    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional[Any] = None,
    ) -> GenericGuardrailAPIInputs:
        """
        Apply Ovalix guardrail to the given inputs (request or response text).

        Used by the unified guardrail flow and the /apply_guardrail API.
        For "request", uses the pre-checkpoint; for "response", uses the post-checkpoint.

        Args:
            inputs: Guardrail API inputs (e.g. texts to check).
            request_data: Full request payload (messages, metadata, response).
            input_type: "request" (pre_call) or "response" (post_call).
            logging_obj: Optional logging context.

        Returns:
            Updated inputs (e.g. with replaced/corrected texts, or unchanged).
        """
        if not self._pre_checkpoint_id and not self._post_checkpoint_id:
            return inputs

        tracker_actor_id = self._get_tracker_actor_id(request_data)
        session_id = self._get_session_id(request_data)
        texts = inputs.get("texts") or []
        if not texts or not isinstance(texts, list):
            return inputs

        if input_type == "response":
            if not self._post_checkpoint_id:
                return inputs
            corrected_llm_responses = await self._generate_post_guardrail_llm_texts(
                texts, tracker_actor_id, session_id, self._post_checkpoint_id
            )
            return {**inputs, "texts": corrected_llm_responses}

        if self._pre_checkpoint_id:
            post_guardrail_texts = await self._generate_post_guardrail_llm_texts(
                texts, tracker_actor_id, session_id, self._pre_checkpoint_id
            )
            return {**inputs, "texts": post_guardrail_texts}
        return inputs

    async def _generate_post_guardrail_llm_texts(
        self, texts: List[str], actor: str, session_id: str, checkpoint_id: str
    ) -> List[str]:
        """Generate post-guardrail LLM responses for the given LLM responses."""
        post_guardrail_texts: List[str] = []

        is_first_response = True
        for llm_response in reversed(texts):
            try:
                resp = await self._call_checkpoint(
                    llm_response, checkpoint_id, actor, session_id
                )
            except Exception as e:
                verbose_proxy_logger.exception(
                    "Ovalix apply_guardrail checkpoint call failed: %s", e
                )
                raise GuardrailRaisedException(
                    guardrail_name=self.guardrail_name,
                    message=f"Ovalix guardrail error: {e!s}",
                    should_wrap_with_default_message=False,
                ) from e

            action_type = (resp.get("action_type") or "").lower()
            blocking_message = (
                self._get_trackers_corrected_message(resp)
                or BLOCKED_BY_OVALIX_FALLBACK_MESSAGE
            )
            if action_type == BLOCKED_ACTION_TYPE and is_first_response:
                self._block_current_message(blocking_message)
            elif action_type == BLOCKED_ACTION_TYPE:
                post_guardrail_texts.insert(0, blocking_message)
            else:
                corrected_text = (
                    self._get_trackers_corrected_message(resp) or llm_response
                )
                post_guardrail_texts.insert(0, corrected_text)
            is_first_response = False
        return post_guardrail_texts

    def _block_current_message(self, blocking_message: str) -> None:
        """Raise OvalixGuardrailBlockedException with the given message (no default wrapper)."""
        raise OvalixGuardrailBlockedException(
            guardrail_name=self.guardrail_name,
            message=blocking_message,
            should_wrap_with_default_message=False,
        )

    def _get_trackers_corrected_message(self, resp: dict) -> Optional[str]:
        """Extract corrected/blocking message content from Tracker checkpoint response."""
        modified = resp.get("modified_data")
        if isinstance(modified, dict) and "content" in modified:
            return modified["content"]
        return None

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

        return OvalixGuardrailConfigModel
