"""Akto guardrail integration for LiteLLM proxy.

Uses a two-config-entry pattern:
  - akto-validate (pre_call): Checks request against Akto guardrails, blocks if flagged.
  - akto-ingest (post_call): Sends request+response to Akto for data ingestion.

For monitor-only mode, enable only akto-ingest without akto-validate.
"""

import asyncio
import json
import os
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Type

from fastapi import HTTPException

import httpx

from litellm._logging import verbose_proxy_logger
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

HTTP_PROXY_PATH = "/api/http-proxy"
AKTO_CONNECTOR_NAME = "litellm"
DEFAULT_GUARDRAIL_TIMEOUT = 5


class AktoGuardrail(CustomGuardrail):
    """LiteLLM guardrail hook that validates and ingests LLM traffic via the Akto API."""

    # Maps event_hook to the input_type it should handle; mismatches are no-ops
    HOOK_TO_INPUT = {"pre_call": "request", "post_call": "response"}

    @staticmethod
    def get_config_model() -> Type["GuardrailConfigModel"]:
        """Return the Pydantic config model for YAML-based initialization."""
        from litellm.types.proxy.guardrails.guardrail_hooks.akto import (
            AktoConfigModel,
        )

        return AktoConfigModel

    def __init__(
        self,
        akto_base_url: Optional[str] = None,
        akto_api_key: Optional[str] = None,
        akto_account_id: Optional[str] = None,
        akto_vxlan_id: Optional[str] = None,
        unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
        guardrail_timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize the Akto guardrail.

        Args:
            akto_base_url: Akto API base URL. Falls back to AKTO_GUARDRAIL_API_BASE env var.
            akto_api_key: Akto API key. Falls back to AKTO_API_KEY env var.
            akto_account_id: Akto account ID. Falls back to AKTO_ACCOUNT_ID env var, then "1000000".
            akto_vxlan_id: Akto VXLAN ID. Falls back to AKTO_VXLAN_ID env var, then "0".
            unreachable_fallback: Behavior when Akto is unreachable — block or allow.
            guardrail_timeout: HTTP timeout in seconds for Akto API calls.
        """
        self.async_handler = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.GuardrailCallback,
        )
        self.background_tasks: set = set()

        self.akto_base_url = (
            akto_base_url or os.environ.get("AKTO_GUARDRAIL_API_BASE", "")
        ).rstrip("/")
        if not self.akto_base_url:
            raise ValueError(
                "akto_base_url is required. Set AKTO_GUARDRAIL_API_BASE or pass it in litellm_params."
            )

        self.akto_api_key = akto_api_key or os.environ.get("AKTO_API_KEY", "")
        if not self.akto_api_key:
            raise ValueError(
                "akto_api_key is required. Set AKTO_API_KEY or pass it in litellm_params."
            )

        self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
            unreachable_fallback
        )
        self.guardrail_timeout = guardrail_timeout or DEFAULT_GUARDRAIL_TIMEOUT
        self.akto_account_id = akto_account_id or os.environ.get(
            "AKTO_ACCOUNT_ID", "1000000"
        )
        self.akto_vxlan_id = akto_vxlan_id or os.environ.get("AKTO_VXLAN_ID", "0")

        kwargs["supported_event_hooks"] = [
            GuardrailEventHooks.pre_call,
            GuardrailEventHooks.post_call,
        ]
        super().__init__(**kwargs)

        verbose_proxy_logger.debug(
            "Akto guardrail initialized: base_url=%s fallback=%s",
            self.akto_base_url,
            self.unreachable_fallback,
        )

    @staticmethod
    def resolve_metadata_value(request_data: Optional[dict], key: str) -> Optional[str]:
        """Look up a metadata value from litellm_metadata or metadata dicts."""
        if request_data is None:
            return None
        for dict_key in ("litellm_metadata", "metadata"):
            container = request_data.get(dict_key) or {}
            if isinstance(container, dict) and container:
                value = container.get(key)
                if value is not None:
                    return str(value).strip()
        return None

    @staticmethod
    def extract_request_path(request_data: dict) -> str:
        """Extract the API route from request metadata, defaulting to /v1/chat/completions."""
        metadata = request_data.get("metadata") or {}
        if not isinstance(metadata, dict):
            metadata = {}
        route = metadata.get("user_api_key_request_route")
        return route if route else "/v1/chat/completions"

    def prepare_headers(self) -> Dict[str, str]:
        """Build HTTP headers for the Akto API call."""
        return {
            "content-type": "application/json",
            "Authorization": self.akto_api_key,
        }

    @staticmethod
    def build_query_params(*, guardrails: bool, ingest_data: bool) -> Dict[str, str]:
        """Build query params that control Akto backend behavior (guardrail check and/or data ingestion)."""
        params: Dict[str, str] = {"akto_connector": AKTO_CONNECTOR_NAME}
        if guardrails:
            params["guardrails"] = "true"
        if ingest_data:
            params["ingest_data"] = "true"
        return params

    @staticmethod
    def build_request_headers(request_data: dict) -> Dict[str, str]:
        """Build the requestHeaders field from proxy request headers."""
        headers: Dict[str, str] = {"content-type": "application/json"}
        proxy_req = request_data.get("proxy_server_request", {})
        if not isinstance(proxy_req, dict):
            return headers
        proxy_req_headers = proxy_req.get("headers")
        if isinstance(proxy_req_headers, dict):
            for key, val in proxy_req_headers.items():
                if key and val:
                    headers[str(key).lower()] = str(val)
        return headers

    @staticmethod
    def build_request_body(
        inputs: GenericGuardrailAPIInputs,
        request_data: Optional[dict] = None,
    ) -> Dict[str, Any]:
        """Build the LLM request body from guardrail inputs (messages, model, tools)."""
        model = inputs.get("model", "") or ""
        body: Dict[str, Any] = {"model": model}

        structured = inputs.get("structured_messages")
        if structured:
            body["messages"] = structured
        elif request_data is not None and request_data.get("messages"):
            body["messages"] = request_data["messages"]
            if request_data.get("model"):
                body["model"] = request_data["model"]
        else:
            texts = inputs.get("texts", [])
            body["messages"] = (
                [{"role": "user", "content": t} for t in texts] if texts else []
            )

        tools = inputs.get("tools")
        if tools:
            body["tools"] = tools
        elif request_data is not None and request_data.get("tools"):
            body["tools"] = request_data["tools"]

        tool_calls = inputs.get("tool_calls")
        if tool_calls:
            body["tool_calls"] = tool_calls

        return body

    @staticmethod
    def build_response_body(
        inputs: GenericGuardrailAPIInputs,
        request_data: Optional[dict] = None,
    ) -> Dict[str, Any]:
        """Build the LLM response body, preferring the actual model response if available."""
        model_response = request_data.get("response") if request_data else None
        if model_response is not None and hasattr(model_response, "model_dump"):
            return model_response.model_dump()

        texts = inputs.get("texts", [])
        if texts:
            return {
                "choices": [
                    {"message": {"content": t, "role": "assistant"}} for t in texts
                ]
            }
        return {}

    @staticmethod
    def build_tag_metadata(request_data: dict) -> Dict[str, str]:
        """Build tag/metadata dict with user_id and team_id for Akto tracking."""
        tag: Dict[str, str] = {"gen-ai": "Gen AI"}
        user_id = AktoGuardrail.resolve_metadata_value(
            request_data, "user_api_key_user_id"
        )
        team_id = AktoGuardrail.resolve_metadata_value(
            request_data, "user_api_key_team_id"
        )
        if user_id:
            tag["user_id"] = user_id
        if team_id:
            tag["team_id"] = team_id
        return tag

    def build_akto_payload(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        *,
        status_code: int = 200,
        include_response: bool = False,
    ) -> Dict[str, Any]:
        """Build the flat MIRRORING payload sent to Akto's HTTP proxy endpoint.

        All body fields use double-encoding: json.dumps({"body": json.dumps(actual_body)})
        to match the canonical CLI hook format.
        """
        request_path = self.extract_request_path(request_data)
        request_headers = self.build_request_headers(request_data)
        request_body = self.build_request_body(inputs, request_data)
        tag = self.build_tag_metadata(request_data)

        response_payload = json.dumps({})  # Empty body wrapper when no response yet
        response_headers: Dict[str, str] = {}
        if include_response:
            response_body = self.build_response_body(inputs, request_data)
            response_payload = json.dumps(
                {"body": json.dumps(response_body)}
            )  # Double-encoded
            response_headers = {"content-type": "application/json"}

        # Extract client IP from proxy headers
        ip = ""
        proxy_req = request_data.get("proxy_server_request", {})
        proxy_headers = (
            proxy_req.get("headers", {}) if isinstance(proxy_req, dict) else {}
        )
        if isinstance(proxy_headers, dict):
            ip = (
                proxy_headers.get("x-forwarded-for")
                or proxy_headers.get("x-real-ip")
                or ""
            )
            if "," in ip:
                ip = ip.split(",")[0].strip()

        return {
            "path": request_path,
            "requestHeaders": json.dumps(request_headers),
            "responseHeaders": json.dumps(response_headers),
            "method": "POST",
            "requestPayload": json.dumps(
                {"body": json.dumps(request_body)}
            ),  # Double-encoded
            "responsePayload": response_payload,
            "ip": ip,
            "destIp": "127.0.0.1",
            "time": str(int(datetime.now().timestamp() * 1000)),
            "statusCode": str(status_code),
            "type": "HTTP/1.1",
            "status": str(status_code),
            "akto_account_id": self.akto_account_id,
            "akto_vxlan_id": self.akto_vxlan_id,
            "is_pending": "false",
            "source": "MIRRORING",
            "direction": None,
            "process_id": None,
            "socket_id": None,
            "daemonset_id": None,
            "enabled_graph": None,
            "tag": json.dumps(tag),
            "metadata": json.dumps(tag),
            "contextSource": "AGENTIC",
        }

    async def send_request(
        self,
        *,
        guardrails: bool,
        ingest_data: bool,
        payload: dict,
    ) -> httpx.Response:
        """Send an HTTP POST to the Akto API endpoint."""
        endpoint = f"{self.akto_base_url}{HTTP_PROXY_PATH}"
        params = self.build_query_params(guardrails=guardrails, ingest_data=ingest_data)
        headers = self.prepare_headers()
        return await self.async_handler.post(
            url=endpoint,
            data=json.dumps(payload),
            params=params,
            headers=headers,
            timeout=self.guardrail_timeout,
        )

    @staticmethod
    def handle_guardrail_response(response: httpx.Response) -> Tuple[bool, str]:
        """Parse the Akto guardrail response. Returns (allowed, reason)."""
        if response.status_code != 200:
            verbose_proxy_logger.error("Akto returned HTTP %d", response.status_code)
            raise httpx.HTTPStatusError(
                f"Akto returned unexpected status {response.status_code}",
                request=response.request,
                response=response,
            )
        try:
            result = response.json()
        except (json.JSONDecodeError, ValueError) as e:
            response_text = getattr(response, "text", "")
            verbose_proxy_logger.error(
                "Akto returned non-JSON body for status 200: %r",
                response_text[:200],
            )
            raise httpx.RequestError(
                "Akto returned non-JSON body",
                request=response.request,
            ) from e
        if not isinstance(result, dict):
            return True, ""
        data = result.get("data") or {}
        if not isinstance(data, dict):
            return True, ""
        guardrails_result = data.get("guardrailsResult") or {}
        if not isinstance(guardrails_result, dict):
            return True, ""
        return (
            bool(guardrails_result.get("Allowed", True)),
            str(guardrails_result.get("Reason", "")),
        )

    def handle_unreachable(
        self,
        inputs: GenericGuardrailAPIInputs,
        error: Exception,
    ) -> GenericGuardrailAPIInputs:
        """Handle Akto being unreachable based on fail_open/fail_closed config."""
        if self.unreachable_fallback == "fail_open":
            verbose_proxy_logger.critical(
                "Akto unreachable (fail-open): %s",
                str(error),
                exc_info=error,
            )
            return inputs

        verbose_proxy_logger.error("Akto unreachable (fail-closed): %s", str(error))
        raise HTTPException(
            status_code=503,
            detail="Akto guardrail service unreachable",
        )

    async def fire_and_forget_request(
        self,
        *,
        guardrails: bool,
        ingest_data: bool,
        payload: dict,
    ) -> None:
        """Send a request without awaiting it in the caller. Errors are logged, not raised."""
        try:
            response = await self.send_request(
                guardrails=guardrails,
                ingest_data=ingest_data,
                payload=payload,
            )
            if response.status_code != 200:
                verbose_proxy_logger.error(
                    "Akto fire-and-forget returned HTTP %d",
                    response.status_code,
                )
        except Exception as e:
            verbose_proxy_logger.error("Akto fire-and-forget error: %s", str(e))

    @log_guardrail_information
    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj=None,
    ) -> GenericGuardrailAPIInputs:
        """Main entry point called by LiteLLM's guardrail framework.

        Pre_call (input_type="request"):
          - Awaits guardrail check. If blocked, fires off ingest with 403 marker and raises.
        Post_call (input_type="response"):
          - Fire-and-forget combined guardrail + ingest call.
        """
        # Skip if this hook doesn't handle the current input_type
        expected = self.HOOK_TO_INPUT.get(str(self.event_hook))
        if expected and expected != input_type:
            return inputs

        if input_type == "request":
            # Pre_call: awaited guardrail check (no ingestion)
            payload = self.build_akto_payload(
                inputs, request_data, include_response=False
            )
            try:
                response = await self.send_request(
                    guardrails=True,
                    ingest_data=False,
                    payload=payload,
                )
                allowed, reason = self.handle_guardrail_response(response)
            except HTTPException:
                raise
            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                return self.handle_unreachable(
                    inputs=inputs,
                    error=e,
                )

            if not allowed:
                # Build a blocked marker payload with 403 status and reason
                blocked_payload = self.build_akto_payload(
                    inputs,
                    request_data,
                    include_response=False,
                    status_code=403,
                )
                blocked_payload["responsePayload"] = json.dumps(
                    {
                        "body": json.dumps(
                            {"x-blocked-by": "Akto Proxy", "reason": reason}
                        ),
                    }
                )
                blocked_payload["responseHeaders"] = json.dumps(
                    {"content-type": "application/json"},
                )
                # Fire-and-forget ingest of the blocked request, then raise 403
                task = asyncio.create_task(
                    self.fire_and_forget_request(
                        guardrails=False,
                        ingest_data=True,
                        payload=blocked_payload,
                    )
                )
                self.background_tasks.add(task)
                task.add_done_callback(self.background_tasks.discard)
                raise HTTPException(
                    status_code=403,
                    detail=reason or "Blocked by Akto Guardrails",
                )

        elif input_type == "response":
            # Post_call: fire-and-forget combined guardrail + ingest
            payload = self.build_akto_payload(
                inputs, request_data, include_response=True
            )
            task = asyncio.create_task(
                self.fire_and_forget_request(
                    guardrails=True,
                    ingest_data=True,
                    payload=payload,
                )
            )
            self.background_tasks.add(task)
            task.add_done_callback(self.background_tasks.discard)

        return inputs
