"""
Topic Control NemoGuard client
Filters content based on allowed/denied topics
"""

import logging
import dataiku
from typing import Dict, Any, List

TOPIC_SAFETY_OUTPUT_RESTRICTION = (
    'If any of the above conditions are violated, please respond with "off-topic". '
    'Otherwise, respond with "on-topic". You must respond with "on-topic" or "off-topic".'
)


class TopicControlClient:
    """
    Client for NVIDIA Topic Control NemoGuard NIM
    Checks if content relates to allowed or denied topics
    Uses Dataiku LLM Mesh Python SDK
    """

    def __init__(self, llm_id: str, api_key: str = None, additional_system_instruction: str = ""):
        """
        Initialize Topic Control 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)
            additional_system_instruction: Optional system instruction to prepend
        """
        self.llm_id = llm_id
        self.logger = logging.getLogger(__name__)
        additional_instruction = (additional_system_instruction or "").strip()
        if additional_instruction:
            self.system_instruction = f"{additional_instruction}\n\n{TOPIC_SAFETY_OUTPUT_RESTRICTION}"
        else:
            self.system_instruction = TOPIC_SAFETY_OUTPUT_RESTRICTION

        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"Topic Control 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 if conversation messages relate to allowed/denied topics

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

        Returns:
            Tuple of (violation_detected: bool, violation_message: str)
            - violation_detected: True if violation detected, False if content is safe
            - violation_message: Empty string if safe, description of violation if detected
        """
        self.logger.info(f"Topic Control check with messages: {messages}")

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

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

            appended_instruction = False
            for message in messages:
                role = message.get("role")
                if role not in {"system", "user", "assistant"}:
                    continue
                content = message.get("content")
                if role == "assistant" and not content:
                    continue
                if content is None:
                    content = ""
                if self.system_instruction and not appended_instruction and role == "system":
                    content = f"{content}\n\n{self.system_instruction}"
                    appended_instruction = True
                completion.with_message(content, role=role)
                if trace_messages is not None:
                    trace_messages.append({"role": role, "content": content})

            if self.system_instruction and not appended_instruction:
                completion.with_message(self.system_instruction, role="system")
                if trace_messages is not None:
                    trace_messages.append({"role": "system", "content": self.system_instruction})

            # Execute the completion
            response = completion.execute()

            self.logger.info(f"Topic Control 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 = response.text.strip().lower()
            self.logger.info(f"Topic Control API response content: '{content}'")

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

            if "off-topic" in content:
                violation_message = "Content was found to be off-topic"
                return True, violation_message
            else:
                return False, ""  # Content is safe

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