"""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, Final, Literal

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.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel


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


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


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: str | None = 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,
            blocked_content=True,
        )


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.
    """

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

    def __init__(
        self,
        tracker_api_base: str | None = None,
        tracker_api_key: str | None = None,
        application_id: str | None = None,
        pre_checkpoint_id: str | None = None,
        post_checkpoint_id: str | None = 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: Final[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: Final = 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: Final = self._get_actor(data).encode()
        normalized_actor_id: Final = 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: Final = self._get_tracker_actor_id(data)
        today: Final = 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: Final = self._application_id
        if not application_id or not checkpoint_id:
            raise ValueError("Ovalix: application_id or checkpoint_id not resolved")

        url: Final = f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
        headers: Final = dict(self._tracker_headers)
        payload: Final = {
            "application_id": application_id,
            "checkpoint_id": checkpoint_id,
            "actor": actor,
            "session_id": session_id,
            "data_type": "TEXT",
            "data": {"content": content},
        }
        response: Final = 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: "LiteLLMLoggingObj | None" = 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: Final = self._get_tracker_actor_id(request_data)
        session_id: Final = self._get_session_id(request_data)
        texts: Final = 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: Final = 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: Final = 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: Final[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}",
                    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) -> str | None:
        """Extract corrected/blocking message content from Tracker checkpoint response."""
        modified: Final = resp.get("modified_data")
        if isinstance(modified, dict) and "content" in modified:
            return modified["content"]
        return None

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

        return OvalixGuardrailConfigModel
