from synthex import Synthex
from typing import Optional, Union, Literal
from transformers.trainer_utils import TrainOutput

from .multi_label_classification_model import MultiLabelClassificationModel

from artifex.core import auto_validate_methods, ParsedModelInstructions, \
    ClassificationInstructions, track_training_calls, track_inference_calls, ValidationError, \
    GuardrailResponseModel, GuardrailResponseScoresModel
from artifex.config import config


@auto_validate_methods
class Guardrail(MultiLabelClassificationModel):
    """
    A Guardrail Model for user inputs or LLM-generated outputs using multi-label classification.
    
    This model can identify multiple safety concerns simultaneously in LLM-generated text,
    outputting independent probabilities for each unsafe content category.
    """

    def __init__(self, synthex: Synthex, language: Literal["english", "spanish", "german"] = "english"):
        """
        Initializes the Guardrail with multi-label classification capabilities.
        Args:
            synthex (Synthex): An instance of the Synthex class to generate the synthetic data used to train 
                the model.
            language (Literal["english", "spanish", "german"]): The language of the text data.
        """
        
        language_to_model = {
            "english": config.GUARDRAIL_ENGLISH_HF_BASE_MODEL,
            "spanish": config.GUARDRAIL_SPANISH_HF_BASE_MODEL,
            "german": config.GUARDRAIL_GERMAN_HF_BASE_MODEL,
        }
        super().__init__(
            synthex, tokenizer_max_length=config.GUARDRAIL_TOKENIZER_MAX_LENGTH,
            base_model_name=language_to_model[language]
        )
        self._system_data_gen_instr_val: list[str] = [
            "The 'text' field should contain user inputs to LLMs or LLM-generated responses, answers, or outputs.",
            "It is imperative that the 'text' field contains both instances of user inputs to LLMs and instances of LLM-generated outputs.",
            "The 'text' field must be in the following language, and only this language: {language}.",
            "The 'labels' field should contain an array of all unsafe content categories that apply to the 'text'.",
            "The 'labels' field can contain zero, one, or multiple categories from the list of unsafe content types.",
            "The 'labels' field must be formatted as a valid JSON array. it must start with a square bracket and end with a square bracket. Each element must be enclosed in quotes. For example: ['hate_speech', 'violence'].",
            "The 'labels' field must always be a JSON array, even if it is empty or contains only one category.",
            "If the text is completely safe and doesn't match any unsafe category, 'labels' should be an empty array [].",
            "The following are the unsafe content categories and their descriptions: {unsafe_categories}.",
            "The dataset should contain a mix of safe text (empty labels) and unsafe text (one or more category labels).",
            "You must include samples of safe text that does not fall under any unsafe category, and that is therefore labeled with an empty array [].",
            "The dataset should contain text with multiple unsafe categories applied simultaneously when appropriate.",
        ]
    
    def _get_data_gen_instr(self, user_instr: ParsedModelInstructions) -> list[str]:
        """
        Overrides `MultiLabelClassificationModel._get_data_gen_instr` to account for the different structure of
        `Guardrail.train`.
        Args:
            user_instr (ParsedModelInstructions): Parsed user instructions containing unsafe categories.
        Returns:
            list[str]: A list containing the formatted system instructions.
        """
        
        # Format the unsafe categories into a single string
        unsafe_categories_str = ", ".join(user_instr.user_instructions)
        
        out = [
            instr.format(
                language=user_instr.language, 
                unsafe_categories=unsafe_categories_str
            ) for instr in self._system_data_gen_instr_val
        ]
        return out
    
    def _parse_user_instructions(
        self, user_instructions: "ClassificationInstructions"
    ) -> ParsedModelInstructions:
        """
        Convert the unsafe categories passed by the user into ParsedModelInstructions.
        Args:
            user_instructions (ClassificationInstructions): Instructions with classes (categories), 
                domain, and language.
        Returns:
            ParsedModelInstructions: Parsed instructions for data generation.
        """
        
        # Format each category with its description
        formatted_instructions = [
            f"{category}: {description}" 
            for category, description in user_instructions.classes.items()
        ]

        return ParsedModelInstructions(
            user_instructions=formatted_instructions,
            language=user_instructions.language,
            domain=user_instructions.domain
        )
        
    @track_training_calls
    def train(
        self, unsafe_categories: Optional[dict[str, str]] = None, language: str = "english", 
        output_path: Optional[str] = None, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3,
        device: Optional[int] = None, disable_logging: Optional[bool] = False,
        train_dataset_path: Optional[str] = None
    ) -> TrainOutput:
        f"""
        Train the Guardrail model to detect multiple unsafe content categories in LLM outputs.
        
        Args:
            unsafe_categories (dict[str, str]): A dictionary mapping category names to their descriptions.
                For example: {{
                    "hate_speech": "Content containing hateful or discriminatory language",
                    "violence": "Content describing or encouraging violent acts",
                    "explicit": "Sexually explicit or inappropriate content"
                }}
                Each text can be tagged with zero, one, or multiple categories.
            language (str): The language to use for generating the training dataset. Defaults to "english".
            output_path (Optional[str]): The path where the synthetic training data and the
                output model will be saved.
            num_samples (int): The number of training data samples to generate. Defaults to 
                {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}.
            num_epochs (int): The number of epochs for training the model. Defaults to 3.
            device (Optional[int]): The device to perform training on. If None, it will use the GPU
                if available, otherwise it will use the CPU.
            disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False.
        Returns:
            TrainOutput: The output of the training process.
        """
        
        if train_dataset_path is None and unsafe_categories is None:
            raise ValidationError(
                message="The `unsafe_categories` parameter is required when `train_dataset_path` is not provided."
            )

        # Call the parent train method with the domain parameter
        return super().train(
            domain="LLM output safety and content moderation",
            labels=unsafe_categories,
            language=language,
            output_path=output_path,
            num_samples=num_samples,
            num_epochs=num_epochs,
            device=device,
            disable_logging=disable_logging,
            train_dataset_path=train_dataset_path
        )
    
    @track_inference_calls
    def __call__(
        self, text: Union[str, list[str]], unsafe_threshold: float = 0.55,
        device: Optional[int] = None, disable_logging: Optional[bool] = False
    ) -> list[GuardrailResponseModel]:
        """
        Classify LLM-generated outputs for multiple unsafe content categories simultaneously.
        
        Args:
            text (str | list[str]): The LLM-generated output(s) to be classified.
            unsafe_threshold (float): The probability threshold above which a label is considered unsafe.
            device (Optional[int]): The device to perform inference on. If None, it will use the GPU
                if available, otherwise it will use the CPU.
            disable_logging (Optional[bool]): Whether to disable logging during inference. Defaults to False.
        Returns:
            list[GuardrailResponseModel]: Classification results containing:
                - labels: dict mapping each category to its probability (0-1)
        """
        
        if unsafe_threshold < 0.0 or unsafe_threshold > 1.0:
            raise ValidationError(
                message="`unsafe_threshold` must be between 0.0 and 1.0."
            )
            
        if not self._model or not self._model.config.id2label:
            raise ValueError("Model not trained or loaded. Please call train() or load() first.")
        
        probs = super().__call__(
            text=text,
            device=device,
            disable_logging=disable_logging
        )
        
        # Build responses
        out = []

        for prob_vector in probs:
            partial_response = {}
            for i, prob in enumerate(prob_vector):
                partial_response[self._model.config.id2label[i]] = round(prob.item(), 4)
            out.append(GuardrailResponseModel(
                is_safe = all(prob_vector < unsafe_threshold),
                scores = GuardrailResponseScoresModel(**partial_response)
            ))

        return out
