from __future__ import annotations

import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict
from urllib.parse import urlparse
from uuid import uuid4

import httpx
import requests
from fastapi import HTTPException
from httpx import HTTPStatusError
from requests.auth import HTTPBasicAuth
from typing_extensions import ReadOnly

from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
    CustomGuardrail,
    log_guardrail_information,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
    HiddenlayerAction,
    HiddenlayerMessages,
)
from litellm.types.utils import GenericGuardrailAPIInputs

if TYPE_CHECKING:
    from pydantic import BaseModel

    from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel


class _HiddenlayerEvaluation(TypedDict, total=False):
    action: str
    threat_level: str


class _HiddenlayerAnalysisEntry(TypedDict, total=False):
    name: str
    detected: bool


class _HiddenlayerModifiedSide(TypedDict):
    messages: Any


class _HiddenlayerResponse(TypedDict, total=False):
    evaluation: _HiddenlayerEvaluation
    analysis: Sequence[_HiddenlayerAnalysisEntry]
    modified_data: Mapping[str, _HiddenlayerModifiedSide]


class _LoggedCallMetadata(TypedDict, total=False):
    headers: ReadOnly[Mapping[str, str]]


class _LoggedCallLitellmParams(TypedDict, total=False):
    metadata: ReadOnly[_LoggedCallMetadata]


class _HiddenlayerOutputMessage(TypedDict, total=False):
    content: ReadOnly[str | Sequence[Mapping[str, str]]]


class _HiddenlayerChoiceMessage(TypedDict, total=False):
    content: ReadOnly[str]


class _HiddenlayerChoice(TypedDict, total=False):
    message: ReadOnly[_HiddenlayerChoiceMessage]


class _HiddenlayerV2Output(TypedDict, total=False):
    messages: ReadOnly[Sequence[_HiddenlayerOutputMessage]]
    choices: ReadOnly[Sequence[_HiddenlayerChoice]]


class _LoggedCallDetails(Protocol):
    """Logging object view that exposes its untyped call details with the shape this guardrail reads."""

    @property
    def model_call_details(self) -> Mapping[str, _LoggedCallLitellmParams]: ...


class _TokenPayloadSource(Protocol):
    """Response view that decodes the HiddenLayer OAuth token body as a string mapping."""

    def json(self) -> Mapping[str, str]: ...


def _logged_request_headers(logging_obj: _LoggedCallDetails) -> Mapping[str, str]:
    return logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})


def _token_payload(response: _TokenPayloadSource) -> Mapping[str, str]:
    return response.json()


def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:
    return headers.get(key, default)


def is_saas(host: str) -> bool:
    """Checks whether the connection is to the SaaS platform"""

    o: Final = urlparse(host)

    if o.hostname and o.hostname.endswith("hiddenlayer.ai"):
        return True

    return False


def _get_jwt(auth_url, api_id, api_key) -> str:
    token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials"

    resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key))

    if not resp.ok:
        raise RuntimeError(
            f"Unable to get authentication credentials for the HiddenLayer API: {resp.status_code}: {resp.text}"
        )

    if "access_token" not in resp.json():
        raise RuntimeError(
            f"Unable to get authentication credentials for the HiddenLayer API - invalid response: {resp.json()}"
        )

    return _token_payload(resp)["access_token"]


class HiddenlayerGuardrail(CustomGuardrail):
    """Custom guardrail wrapper for HiddenLayer's safety checks."""

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

    def __init__(
        self,
        api_id: str | None = None,
        api_key: str | None = None,
        api_base: str | None = None,
        auth_url: str | None = None,
        **kwargs: Any,
    ) -> None:
        kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
        self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
        self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")
        self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai"
        self.jwt_token = None

        auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai"

        if is_saas(self.api_base):
            if not self.hiddenlayer_client_id:
                raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.")

            if not self.hiddenlayer_client_secret:
                raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.")

            self.jwt_token = _get_jwt(
                auth_url=auth_url,
                api_id=self.hiddenlayer_client_id,
                api_key=self.hiddenlayer_client_secret,
            )
            self.refresh_jwt_func = lambda: _get_jwt(
                auth_url=auth_url,
                api_id=self.hiddenlayer_client_id,
                api_key=self.hiddenlayer_client_secret,
            )

        self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
        super().__init__(**kwargs)

    @log_guardrail_information
    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: LiteLLMLoggingObj | None = None,
    ) -> GenericGuardrailAPIInputs:
        """Validate (and optionally redact) text via HiddenLayer before/after LLM calls."""

        # The model in the request and the response can be inconsistent
        # I.e request can specify gpt-4o-mini but the response from the server will be
        # gpt-4o-mini-2025-11-01. We need the model to be consistent so that inferences
        # will be grouped correctly on the Hiddenlayer side
        model_name: Final = logging_obj.model if logging_obj and logging_obj.model else "unknown"
        hl_request_metadata: Final = {"model": model_name}

        # We need the hiddenlayer project id and requester id on both the input and output
        # Since headers aren't available on the response back from the model, we get them
        # from the logging object. It ends up working out that on the request, we parse the
        # hiddenlayer params from the raw request and then retrieve those same headers
        # from the logger object on the response from the model.
        headers = request_data.get("proxy_server_request", {}).get("headers", {})
        if not headers and logging_obj and logging_obj.model_call_details:
            headers = _logged_request_headers(logging_obj)

        hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
        project_id: Final = headers.get("hl-project-id")

        if scan_params := inputs.get("structured_messages"):
            last_msg: Final = scan_params[-1]
            result: _HiddenlayerResponse = await self._call_hiddenlayer(
                project_id,
                hl_request_metadata,
                {
                    "messages": [
                        {
                            "role": last_msg.get("role", "user"),
                            "content": str(last_msg.get("content", "")),
                        }
                    ]
                },
                input_type,
            )
        elif text := inputs.get("texts"):
            result = await self._call_hiddenlayer(
                project_id,
                hl_request_metadata,
                {"messages": [{"role": "user", "content": text[-1]}]},
                input_type,
            )
        else:
            result = {}

        if result.get("evaluation", {}).get("action") == HiddenlayerAction.BLOCK:
            detected_reasons: Final = [
                entry.get("name", "unknown") for entry in result.get("analysis", []) if entry.get("detected")
            ]
            threat_level: Final = result.get("evaluation", {}).get("threat_level")
            raise HTTPException(
                status_code=400,
                detail={
                    "error": "Violated guardrail policy",
                    "hiddenlayer_guardrail_response": HiddenlayerMessages.BLOCK_MESSAGE.value,
                    "block_reasons": detected_reasons,
                    "threat_level": threat_level,
                },
            )

        if result.get("evaluation", {}).get("action") == HiddenlayerAction.REDACT:
            modified_data: Final = result.get("modified_data", {})
            if modified_data.get("input") and input_type == "request":
                last_content = modified_data["input"]["messages"][-1]["content"]
                if isinstance(last_content, list):
                    texts = [
                        item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text"
                    ]
                    inputs["texts"] = texts if texts else [""]
                else:
                    inputs["texts"] = [last_content]
                inputs["structured_messages"] = modified_data["input"]["messages"]

            if modified_data.get("output") and input_type == "response":
                last_content = modified_data["output"]["messages"][-1]["content"]
                if isinstance(last_content, list):
                    texts = [
                        item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text"
                    ]
                    inputs["texts"] = texts if texts else [""]
                else:
                    inputs["texts"] = [last_content]

        return inputs

    async def _call_hiddenlayer(
        self,
        project_id: str | None,
        metadata: Mapping[str, str],
        payload: Mapping[str, Sequence[Mapping[str, str]]],
        input_type: Literal["request", "response"],
    ) -> _HiddenlayerResponse:
        data: Final[dict[str, object]] = {"metadata": metadata}

        if input_type == "request":
            data["input"] = payload
        else:
            data["output"] = payload

        headers: Final = {
            "Content-Type": "application/json",
            "hl-runtime-edge-provider": "litellm",
            "hl-runtime-edge-provider-version": "1",
        }

        if project_id:
            headers["HL-Project-Id"] = project_id

        if self.jwt_token:
            headers["Authorization"] = f"Bearer {self.jwt_token}"

        try:
            response = await self._http_client.post(
                f"{self.api_base}/detection/v1/interactions",
                json=data,
                headers=headers,
            )
            response.raise_for_status()
            result: _HiddenlayerResponse = response.json()

            verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result)

            return result
        except HTTPStatusError as e:
            # Try the request again by refreshing the jwt if we get 401
            # since the Hiddenlayer jwt timeout is an hour and this is
            # a long lived session application
            if e.response.status_code == 401 and self.jwt_token is not None:
                verbose_proxy_logger.debug(
                    "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
                )
                self.jwt_token = self.refresh_jwt_func()
                headers["Authorization"] = f"Bearer {self.jwt_token}"
                response = await self._http_client.post(
                    f"{self.api_base}/detection/v1/interactions",
                    json=data,
                    headers=headers,
                )
            else:
                raise e

            response.raise_for_status()
            result = response.json()

            verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result)
            return result

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

        return HiddenlayerGuardrailConfigModel


class HiddenlayerGuardrailV2(CustomGuardrail):
    """Custom guardrail wrapper for HiddenLayer's safety checks."""

    def __init__(
        self,
        api_id: str | None = None,
        api_key: str | None = None,
        api_base: str | None = None,
        auth_url: str | None = None,
        **kwargs: Any,
    ) -> None:
        self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
        self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")
        self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai"
        self.jwt_token = None

        auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai"

        if is_saas(self.api_base):
            if not self.hiddenlayer_client_id:
                raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.")

            if not self.hiddenlayer_client_secret:
                raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.")

            self.jwt_token = _get_jwt(
                auth_url=auth_url,
                api_id=self.hiddenlayer_client_id,
                api_key=self.hiddenlayer_client_secret,
            )
            self.refresh_jwt_func = lambda: _get_jwt(
                auth_url=auth_url,
                api_id=self.hiddenlayer_client_id,
                api_key=self.hiddenlayer_client_secret,
            )

        self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
        super().__init__(**kwargs)

    @log_guardrail_information
    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: LiteLLMLoggingObj | None = None,
    ) -> GenericGuardrailAPIInputs:
        """Validate (and optionally redact) text via HiddenLayer before/after LLM calls."""

        # We need the hiddenlayer project id and requester id on both the input and output
        # Since headers aren't available on the response back from the model, we get them
        # from the logging object. It ends up working out that on the request, we parse the
        # hiddenlayer params from the raw request and then retrieve those same headers
        # from the logger object on the response from the model.
        headers = request_data.get("proxy_server_request", {}).get("headers", {})
        if not headers and logging_obj and logging_obj.model_call_details:
            headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})

        # put our roundtrip id in the header to the model so we get it on the way back from the model
        if "hl-roundtrip-id" not in headers:
            proxy_req: Final = request_data.get("proxy_server_request")
            if proxy_req is not None and "headers" in proxy_req:
                proxy_req["headers"]["hl-roundtrip-id"] = str(uuid4())
                headers["hl-roundtrip-id"] = proxy_req["headers"]["hl-roundtrip-id"]

        hl_headers: Final = {h.lower(): v for h, v in headers.items() if h.lower().startswith("hl-")}

        if "hl-requester-id" not in hl_headers:
            hl_headers["hl-requester-id"] = "LiteLLM"

        payload: object
        if input_type == "request":
            payload = {
                "messages": inputs.get("structured_messages"),
                "model": inputs.get("model"),
                "tools": inputs.get("tools"),
            }
        else:
            if inputs.get("texts"):
                payload = {
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": (inputs["texts"][0] if inputs.get("texts") else ""),
                            },
                            "finish_reason": "stop",
                        }
                    ]
                }
            elif tool_calls := inputs.get("tool_calls"):
                payload = tool_calls
            else:
                payload = {}

        response: Final = await self._call_hiddenlayer(payload, input_type, hl_headers)
        output: Final = response.json()
        evaluated_output: Final[_HiddenlayerV2Output] = output

        if _header_value(response.headers, "hl-runtime-action", "").lower() == "block":
            raise HTTPException(
                status_code=400,
                detail={
                    "error": "Violated guardrail policy",
                    "hiddenlayer_guardrail_response": HiddenlayerMessages.BLOCK_MESSAGE.value,
                },
            )

        new_texts: Final = []
        if input_type == "request":
            inputs["structured_messages"] = output

            modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = evaluated_output.get("messages", [])
            for message in modified_messages:
                content = message.get("content", "")
                if isinstance(content, list):
                    text_parts = [
                        item["text"] for item in content if isinstance(item, dict) and item.get("type") == "text"
                    ]
                    if text_parts:
                        new_texts.append(" ".join(text_parts))
                elif content:
                    new_texts.append(content)

            inputs["texts"] = new_texts

        elif input_type == "response" and inputs.get("texts"):
            redacted_choices: Final[Sequence[_HiddenlayerChoice]] = evaluated_output.get("choices", [{}])
            inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")]
        elif input_type == "response" and inputs.get("tool_calls"):
            inputs["tool_calls"] = output

        return inputs

    async def _call_hiddenlayer(
        self,
        payload: Any,
        input_type: Literal["request", "response"],
        hl_headers: dict[str, str],
    ) -> httpx.Response:
        if input_type == "request":
            path = "detection/v2/request-evaluations"
        else:
            path = "detection/v2/response-evaluations"

        headers: Final = {
            "Content-Type": "application/json",
            "hl-runtime-edge-provider": "litellm",
            "hl-runtime-edge-provider-version": "2",
        }
        if self.jwt_token:
            headers["Authorization"] = f"Bearer {self.jwt_token}"

        headers.update(hl_headers)

        try:
            response = await self._http_client.post(
                f"{self.api_base}/{path}",
                json=payload,
                headers=headers,
            )
            response.raise_for_status()

            verbose_proxy_logger.debug("Hiddenlayer reponse: %s", response)

            return response
        except HTTPStatusError as e:
            # Try the request again by refreshing the jwt if we get 401
            # since the Hiddenlayer jwt timeout is an hour and this is
            # a long lived session application
            if e.response.status_code == 401 and self.jwt_token is not None:
                verbose_proxy_logger.debug(
                    "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
                )
                self.jwt_token = self.refresh_jwt_func()
                headers["Authorization"] = f"Bearer {self.jwt_token}"
                response = await self._http_client.post(
                    f"{self.api_base}/{path}",
                    json=payload,
                    headers=headers,
                )
            else:
                raise e

            response.raise_for_status()

            verbose_proxy_logger.debug("Hiddenlayer reponse: %s", response)
            return response

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

        return HiddenlayerGuardrailConfigModel
