"""
Content Moderation NemoGuard client
Detects inappropriate, harmful, or policy-violating content
"""
import json
import logging
import dataiku

from typing import Dict, Any, List, Optional

CONTENT_SAFETY_SYSTEM_INSTRUCTION = (
    "Check if there is unsafe content in the user or agent messages in the below conversation in "
    "accordance with our safety policy with the below categories. Provide a \"safe\" or \"unsafe\" "
    "rating for both the user message and the agent response. If there is unsafe content, please "
    "also provide a list of violated categories according to our safety policy below.\n\n"
    "<BEGIN UNSAFE CONTENT CATEGORIES>\n"
    "S1: Violence.\n"
    "S2: Sexual.\n"
    "S3: Criminal Planning/Confessions.\n"
    "S4: Guns and Illegal Weapons.\n"
    "S5: Controlled/Regulated Substances.\n"
    "S6: Suicide and Self Harm.\n"
    "S7: Sexual (minor).\n"
    "S8: Hate/Identity Hate.\n"
    "S9: PII/Privacy.\n"
    "S10: Harassment.\n"
    "S11: Threat.\n"
    "S12: Profanity.\n"
    "S13: Needs Caution.\n"
    "S14: Other.\n"
    "S15: Manipulation.\n"
    "S16: Fraud/Deception.\n"
    "S17: Malware.\n"
    "S18: High Risk Gov Decision Making.\n"
    "S19: Political/Misinformation/Conspiracy.\n"
    "S20: Copyright/Trademark/Plagiarism.\n"
    "S21: Unauthorized Advice.\n"
    "S22: Illegal Activity.\n"
    "S23: Immoral/Unethical.\n"
    "<END UNSAFE CONTENT CATEGORIES>\n\n"
    "Provide your safety assessment for the above conversation in the below output JSON format:\n"
    "{\"User Safety\": \"safe or unsafe rating\", \"Response Safety\": \"safe or unsafe rating\", "
    "\"Safety Categories\": \"comma-separated list\"}"
)


class ContentModerationClient:
    """
    Client for NVIDIA Content Moderation NemoGuard NIM
    Checks content for policy violations, toxicity, and inappropriate material
    Uses Dataiku LLM Mesh Python SDK
    """

    def __init__(self, llm_id: str, api_key: str = None):
        """
        Initialize Content Moderation client using Dataiku LLM Mesh

        Args:
            llm_id: LLM identifier from Dataiku LLM Mesh
            api_key: API key for authentication (optional, handled by LLM Mesh)
        """
        self.llm_id = llm_id
        self.logger = logging.getLogger(__name__)

        try:
            # Get the LLM handle from Dataiku LLM Mesh
            client = dataiku.api_client()
            project = client.get_default_project()
            self.llm = project.get_llm(self.llm_id)
            self.logger.info(f"Content Moderation client initialized with LLM ID: {llm_id}")
        except Exception as e:
            self.logger.error(f"Failed to initialize LLM from Mesh: {str(e)}")
            raise

    def check(self, messages: List[Dict[str, str]], **kwargs) -> tuple[bool, str]:
        """
        Check messages for content moderation violations

        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 violations detected, False if content is safe
            - violation_message: Empty string if safe, description of violation if detected
        """

        trace_payload = kwargs.get("trace_payload")
        trace_messages = [] if isinstance(trace_payload, dict) else None

        # Extract only the last user message and the last assistant message with content.
        last_user_idx = -1
        last_assistant_idx = -1
        for i in range(len(messages) - 1, -1, -1):
            role = messages[i].get("role")
            if last_assistant_idx == -1 and role == "assistant":
                if messages[i].get("content"):
                    last_assistant_idx = i
                else:
                    self.logger.info("Skipping assistant message without content after last user message")
            if role == "user":
                last_user_idx = i
                break

        if last_user_idx == -1:
            self.logger.warning("No user messages found in the messages array")
            return False, ""

        filtered_messages = [messages[last_user_idx]]
        if last_assistant_idx > last_user_idx:
            filtered_messages.append(messages[last_assistant_idx])

        #self.logger.info(f"Content Moderation check with messages: {filtered_messages}")

        try:
            # Create completion query using Dataiku LLM Mesh SDK
            completion = self.llm.new_completion()

            # Add system instruction (do not use any system message from Dataiku)
            completion.with_message(CONTENT_SAFETY_SYSTEM_INSTRUCTION, role="system")
            if trace_messages is not None:
                trace_messages.append({
                    "role": "system",
                    "content": CONTENT_SAFETY_SYSTEM_INSTRUCTION
                })

            # Add filtered messages to the completion
            for message in filtered_messages:
                role = message.get("role", "user")
                content = message.get("content", "")
                completion.with_message(content, role=role)
                if trace_messages is not None:
                    trace_messages.append({"role": role, "content": content})

            # Execute the completion
            response = completion.execute()

            self.logger.info(f"Content Moderation API RESPONSE: {response.success}")

            # Check if the completion was successful
            if not response.success:
                # Try to get error details from response text
                error_detail = response.text if response.text else "Unknown error"
                raise ValueError(f"LLM completion failed: {error_detail}")

            # Get the response text
            if not response.text:
                raise ValueError("LLM completion returned empty response")

            content_str = response.text.strip()
            self.logger.info(f"Content Moderation API response content: '{content_str}'")

            if isinstance(trace_payload, dict):
                trace_payload["inputs"] = {"messages": trace_messages or filtered_messages}
                trace_payload["outputs"] = {"text": response.text}

            try:
                moderation_result = json.loads(content_str)
            except json.JSONDecodeError:
                raise ValueError(f"Failed to parse JSON from Content Moderation API response: {content_str}")

            user_safety = moderation_result.get("User Safety", "safe").lower()

            # Check if either user or response is unsafe
            user_unsafe = user_safety == "unsafe"
            response_unsafe = False

            # Response Safety is optional, only check if it exists
            if "Response Safety" in moderation_result:
                response_safety = moderation_result.get("Response Safety", "safe").lower()
                response_unsafe = response_safety == "unsafe"

            if user_unsafe or response_unsafe:
                # Build violation message based on what's unsafe
                unsafe_parts = []
                if user_unsafe:
                    unsafe_parts.append("user query")
                if response_unsafe:
                    unsafe_parts.append("response")

                unsafe_description = " and ".join(unsafe_parts)
                categories = moderation_result.get("Safety Categories", "N/A")
                return True, f"Content flagged as unsafe: {unsafe_description}. Categories: {categories}"
            else:
                return False, ""  # Content is safe

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