"""
Jailbreak Detection NemoGuard client
Detects prompt injection and jailbreak attempts
"""

import logging
import requests
from typing import Dict, Any


class JailbreakDetectionClient:
    """
    Client for NVIDIA Jailbreak Detection NemoGuard NIM
    Detects attempts to bypass AI safety measures through prompt injection
    Uses direct REST API calls
    """

    def __init__(self, endpoint: str, api_key: str):
        """
        Initialize Jailbreak Detection client

        Note: Jailbreak detection doesn't use a model_name or threshold parameter.
        Pass/fail is determined solely by the API's jailbreak field.

        Args:
            endpoint: Full URL to the Jailbreak Detection NIM endpoint
            api_key: API key for authentication
        """
        self.endpoint = endpoint
        self.api_key = api_key
        self.logger = logging.getLogger(__name__)

    def _get_headers(self) -> Dict[str, str]:
        """
        Build HTTP headers with authentication

        Returns:
            Dictionary of HTTP headers
        """
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json"
        }

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

        return headers

    def _make_request(self, payload: Dict[str, Any], timeout: int = 30) -> Dict[str, Any]:
        """
        Make HTTP POST request to Jailbreak Detection NIM endpoint

        Args:
            payload: Request payload
            timeout: Request timeout in seconds

        Returns:
            Response JSON as dictionary

        Raises:
            requests.exceptions.RequestException: On request failure
        """
        try:
            self.logger.debug(f"Making request to {self.endpoint}")
            self.logger.debug(f"Payload: {payload}")

            response = requests.post(
                self.endpoint,
                json=payload,
                headers=self._get_headers(),
                timeout=timeout
            )

            response.raise_for_status()

            result = response.json()
            self.logger.debug(f"Response: {result}")

            return result

        except requests.exceptions.Timeout:
            self.logger.error(f"Request to {self.endpoint} timed out after {timeout}s")
            raise
        except requests.exceptions.HTTPError as e:
            self.logger.error(f"HTTP error from {self.endpoint}: {e.response.status_code} - {e.response.text}")
            raise
        except requests.exceptions.RequestException as e:
            self.logger.error(f"Request failed: {str(e)}")
            raise
        except ValueError as e:
            self.logger.error(f"Failed to parse response JSON: {str(e)}")
            raise

    def check(self, messages: list, **kwargs) -> tuple[bool, str]:
        """
        Check messages for jailbreak attempts using NVIDIA NemoGuard Jailbreak Detection API

        Args:
            messages: List of message dictionaries with 'role' and 'content' keys
            **kwargs: Additional parameters

        Returns:
            Tuple of (violation_detected: bool, violation_message: str)
            - violation_detected: True if jailbreak detected, False if content is safe
            - violation_message: Empty string if safe, description of violation if detected
        """
        # Get the last user message for jailbreak detection
        if not messages:
            self.logger.warning("No messages provided")
            return False, ""

        last_message = messages[-1]
        if last_message.get("role") != "user":
            self.logger.info("Last message is not a user message, skipping jailbreak check")
            return False, ""

        text = last_message.get("content", "")

        trace_payload = kwargs.get("trace_payload")

        # Build the request payload according to NVIDIA API specification
        payload = {
            "input": text
        }

        # self.logger.info(f"Jailbreak Detection API PAYLOAD #### : '{payload}'")

        try:
            response = self._make_request(payload)

            self.logger.info(f"Jailbreak Detection API RESPONSE #### : '{response}'")

            # Parse API response
            # Response format: {"jailbreak": true|false (boolean), "score": 0.8785 (float)}
            jailbreak_value = response.get("jailbreak", False)
            score = float(response.get("score", 0.0))

            if isinstance(trace_payload, dict):
                trace_payload["inputs"] = {"payload": payload}
                trace_payload["outputs"] = {"response": response}

            # Handle both boolean and string responses for backward compatibility
            if isinstance(jailbreak_value, bool):
                is_jailbreak = jailbreak_value
            elif isinstance(jailbreak_value, str):
                is_jailbreak = jailbreak_value.lower() == "true"
            else:
                self.logger.warning(f"Unexpected jailbreak value type: {type(jailbreak_value)}, treating as False")
                is_jailbreak = False

            self.logger.info(f"Jailbreak Detection: jailbreak={is_jailbreak}, score={score:.4f}")

            # Return True if violation detected (jailbreak=true)
            if is_jailbreak:
                violation_msg = f"Jailbreak attempt identified with a score of {score:.4f}"
                return True, violation_msg
            else:
                return False, ""

        except Exception as e:
            self.logger.error(f"Jailbreak Detection check failed: {str(e)}")
            raise
