# This file contains the implementation of the custom guardrail nemoguard-nim
import logging
import dataiku
from dataiku.llm.guardrails import BaseGuardrail
from dku_nvidia.nemoguard import (
    TopicControlClient,
    ContentModerationClient,
    JailbreakDetectionClient
)


class NemoGuardrailNims(BaseGuardrail):
    def _get_api_key_from_connection(self, connection_name: str) -> str:
        """
        Extract API key from Dataiku connection

        Args:
            connection_name: Name of the Dataiku connection

        Returns:
            API key string
        """
        try:
            if not connection_name:
                self.logger.warning("No connection name provided")
                return ""

            client = dataiku.api_client()
            connection = client.get_connection(connection_name)
            conn_info = connection.get_info()
            api_key = conn_info.get_params().get('apiKey', '')

            if not api_key:
                self.logger.warning(f"API key not found in connection '{connection_name}'")

            return api_key
        except Exception as e:
            self.logger.error(f"Failed to get API key from connection '{connection_name}': {str(e)}")
            return ""

    def set_config(self, config, plugin_config):
        """
        Capture and store configuration from guardrail.json and parameter set
        Initialize the appropriate NemoGuard client based on guardrail type
        """
        self.config = config
        self.logger = logging.getLogger(__name__)

        self.logger.info(f"Full Config Received: '{config}'")

        # Capture guardrail type selection
        self.guardrail_type = config.get("nim_guardrail_type", "content_moderation")

        # Capture action to take when guardrail triggers
        self.action = config.get("action", "REJECT")

        # Capture decline message if action is DECLINE
        self.decline_msg = config.get("decline_msg", "I do not feel comfortable talking about this!")
        self.topic_control_system_instruction = config.get("topic_control_system_instruction", "")

        # Capture configuration based on guardrail type
        if self.guardrail_type == "jailbreak_protection":
            # Jailbreak protection uses manual endpoint and requires API key
            endpoint = config.get("model_endpoint", "")
            connection_name = config.get("nvidia_nim_connection", "")
            api_key = self._get_api_key_from_connection(connection_name)
        else:
            # Topic Control and Content Moderation use LLM ID from Mesh
            # API key is handled automatically by LLM Mesh
            llm_id = config.get("model_id", "")

        # Initialize the appropriate client based on guardrail type
        self.client = None

        try:
            if self.guardrail_type == "topic_control":
                if not llm_id:
                    raise ValueError(f"Topic Control llm_id not configured")
                self.client = TopicControlClient(
                    llm_id=llm_id,
                    additional_system_instruction=self.topic_control_system_instruction
                )
                self.logger.info(f"Topic Control client initialized with LLM ID: {llm_id}")

            elif self.guardrail_type == "content_moderation":
                if not llm_id:
                    raise ValueError(f"Content Moderation llm_id not configured")
                self.client = ContentModerationClient(llm_id=llm_id)
                self.logger.info(f"Content Moderation client initialized with LLM ID: {llm_id}")

            elif self.guardrail_type == "jailbreak_protection":
                if not endpoint:
                    raise ValueError(f"Jailbreak Detection endpoint not configured")
                self.client = JailbreakDetectionClient(
                    endpoint=endpoint,
                    api_key=api_key
                )
                self.logger.info("Jailbreak Detection client initialized")

            else:
                raise ValueError(f"Unknown guardrail type: {self.guardrail_type}")

            self.logger.info(f"NemoGuard initialized - Type: {self.guardrail_type}, Action: {self.action}")

        except Exception as e:
            self.logger.error(f"Failed to initialize NemoGuard client: {str(e)}")
            self.client = None
            self.initialization_error = str(e)

    def process(self, input, trace):
        """
        Main processing method - uses NemoGuard client to check content
        Handles both completionResponse and completionQuery
        """
        # Check if client was initialized successfully
        self.logger.info(f"Processing input: {input}")

        if not self.client:
            error_msg = getattr(self, 'initialization_error', f"No client initialized for {self.guardrail_type}")
            self.logger.error(f"Client initialization failed: {error_msg}")

            # Determine if this is a response or query check and fail appropriately
            if "completionResponse" in input:
                return self._handle_response_failure(input, trace, f"Guardrail initialization failed: {error_msg}")
            elif "completionQuery" in input:
                return self._handle_query_failure(input, trace, f"Guardrail initialization failed: {error_msg}")
            else:
                return input

        # Handle response checking
        if "completionResponse" in input:
            self.logger.info(f"Processing response with {self.guardrail_type} guardrail")

            # Jailbreak protection typically only operates on queries, not responses
            if self.guardrail_type == "jailbreak_protection":
                self.logger.info("Jailbreak Detection only operates on queries, skipping response")
                return input

            # Check if response contains text (skip if it's a tool call or doesn't have text)
            if "text" not in input["completionResponse"]:
                self.logger.info("Response doesn't contain text (e.g., tool calling), skipping guardrail check")
                return input

            # Build messages array from completionQuery (includes conversation history + response)
            messages = []
            if "completionQuery" in input and "messages" in input["completionQuery"]:
                messages = input["completionQuery"]["messages"].copy()

            # Append the assistant's response as the last message
            response_text = input["completionResponse"]["text"]
            messages.append({"role": "assistant", "content": response_text})

            try:
                # Call the NemoGuard client to check the messages
                # Returns (violation_detected: bool, violation_message: str)
                violation_detected, violation_message = self._check_with_trace(messages, trace)

                if not violation_detected:
                    self.logger.info("Response passed guardrail check")
                    return input
                else:
                    self.logger.info(f"Response failed guardrail check: {violation_message}")
                    return self._handle_response_failure(input, trace, violation_message)

            except Exception as e:
                error_message = f"Guardrail check failed with error: {str(e)}"
                self.logger.error(f"Error during response guardrail check: {str(e)}")
                return self._handle_response_failure(input, trace, error_message)

        # Handle query checking
        elif "completionQuery" in input:
            messages = input["completionQuery"]["messages"]
            self.logger.info(f"Processing query with {self.guardrail_type} guardrail")

            if not messages:
                self.logger.warning("No messages found in query, skipping guardrail check")
                return input
            if messages[-1].get("role") == "tool":
                self.logger.info("Last message is a tool message, skipping guardrail check")
                return input

            try:
                # Call the NemoGuard client to check the messages
                # Returns (violation_detected: bool, violation_message: str)
                violation_detected, violation_message = self._check_with_trace(messages, trace)

                if not violation_detected:
                    self.logger.info("Query passed guardrail check")
                    return input
                else:
                    self.logger.info(f"Query failed guardrail check: {violation_message}")
                    return self._handle_query_failure(input, trace, violation_message)

            except Exception as e:
                error_message = f"Guardrail check failed with error: {str(e)}"
                self.logger.error(f"Error during query guardrail check: {str(e)}")
                return self._handle_query_failure(input, trace, error_message)

        else:
            self.logger.warning("Input contains neither completionResponse nor completionQuery")
            return input

    def _check_with_trace(self, messages, trace):
        if trace is None:
            return self.client.check(messages)

        trace_payload = {}
        span_label = f"Calling NemoGuard {self.guardrail_type}"
        attributes = {
            "NemoGuardType": self.guardrail_type
        }
        if self.guardrail_type == "jailbreak_protection":
            span_label = f"Calling NemoGuard Jailbreak Detection {self.client.endpoint}"
            attributes["NemoGuardEndpoint"] = self.client.endpoint
        else:
            span_label = f"Calling NemoGuard {self.guardrail_type} {self.client.llm_id}"
            attributes["NemoGuardLlmId"] = self.client.llm_id

        with trace.subspan(span_label) as subspan:
            try:
                violation_detected, violation_message = self.client.check(
                    messages,
                    trace_payload=trace_payload
                )
                subspan.span["outputs"] = trace_payload.get("outputs", {})
                return violation_detected, violation_message
            except Exception as e:
                subspan.span["outputs"] = {"error": str(e)}
                raise
            finally:
                subspan.span["inputs"] = trace_payload.get("inputs", {"messages": messages})
                subspan.span["attributes"] = attributes

    def _handle_response_failure(self, input, trace, violation_message):
        """
        Handle response that failed the guardrail check

        Args:
            input: The input dictionary
            trace: Trace object for logging
            violation_message: Description of the violation from NemoGuard client

        Returns:
            Modified input with Dataiku guardrail response format
        """
        self.logger.info(f"Response guardrail check failed (action:{self.action}): {violation_message}")

        if self.action == "REJECT":
            input["responseGuardrailResponse"] = {
                "action": "FAIL",
                "error": {
                    "message": violation_message
                }
            }
        elif self.action == "AUDIT":
            input["responseGuardrailResponse"] = {
                "action": "PASS_WITH_AUDIT",
                "auditData": [{
                    "origin": f"NemoGuard {self.guardrail_type}",
                    "violation": violation_message
                }]
            }
        elif self.action == "DECLINE":
            input["responseGuardrailResponse"] = {
                "action": "RESPOND"
            }
            input["completionResponse"]["text"] = self.decline_msg

        return input

    def _handle_query_failure(self, input, trace, violation_message):
        """
        Handle query that failed the guardrail check

        Args:
            input: The input dictionary
            trace: Trace object for logging
            violation_message: Description of the violation from NemoGuard client

        Returns:
            Modified input with Dataiku guardrail response format
        """
        self.logger.info(f"Query guardrail check failed (action:{self.action}): {violation_message}")

        if self.action == "REJECT":
            input["queryGuardrailResponse"] = {
                "action": "FAIL",
                "error": {
                    "message": violation_message
                }
            }
        elif self.action == "AUDIT":
            input["queryGuardrailResponse"] = {
                "action": "PASS_WITH_AUDIT",
                "auditData": [{
                    "origin": f"NemoGuard {self.guardrail_type}",
                    "violation": violation_message
                }]
            }
        elif self.action == "DECLINE":
            input["queryGuardrailResponse"] = {
                "action": "RESPOND",
                "overriddenResponseText": self.decline_msg,
                "auditData": [{
                    "origin": f"NemoGuard {self.guardrail_type}",
                    "violation": violation_message
                }]
            }

        return input
