from abc import ABC, abstractmethod
from synthex import Synthex
from synthex.models import JobOutputSchemaDefinition, JobStatus, JobStatusResponseModel
from transformers.modeling_utils import PreTrainedModel
from transformers.tokenization_utils_base import BatchEncoding, PreTrainedTokenizerBase
import time
from datasets import DatasetDict, disable_caching
from typing import Callable, Sequence, Any, Optional, Union
import os
from transformers.trainer_utils import TrainOutput
from rich.progress import Progress
from rich.console import Console
import torch

from artifex.config import config
from artifex.core import auto_validate_methods, BadRequestError, ServerError, ValidationError, \
    NERInstructions, ClassificationInstructions, ParsedModelInstructions
from artifex.utils import get_dataset_output_path, get_model_output_path

# TODO: While this appears to be the only way to suppress the tedious warning about the 
# BaseModel._tokenize_dataset.tokenize function not being hashable, the solution is not ideal as it 
# disables caching entirely, loading to potentially slower data processing.
disable_caching()

console = Console()

@auto_validate_methods
class BaseModel(ABC):
    """
    Base class for all models in the artifex package.
    """
    
    def __init__(self, synthex: Synthex):
        self._synthex_val: Synthex = synthex
        self._tokenizer_val: PreTrainedTokenizerBase = PreTrainedTokenizerBase()
        self._model_val: Optional[PreTrainedModel] = None

    ##### Abstract properties #####

    @property
    @abstractmethod
    def _synthetic_data_schema(self) -> JobOutputSchemaDefinition:
        """
        Schema of the synthetically generated data, used to train the model.
        """
        pass
    
    @property
    @abstractmethod
    def _system_data_gen_instr(self) -> str:
        """
        The system instructions for synthetic data generation.
        """
        pass
    
    @property
    @abstractmethod
    def _token_keys(self) -> list[str]:
        """
        The keys in the dataset that contain the text to be tokenized. These are used to tokenize
        the dataset during training and inference.
        """
        pass
    
    @property
    @abstractmethod
    def _base_model_name(self) -> str:
        """
        The name of the base model to use for the model.
        """
        pass
    
    ##### Abstract methods #####
    
    @abstractmethod
    def _parse_user_instructions(
        self, user_instructions: Union[list[str], NERInstructions, ClassificationInstructions],
        language: str
    ) -> ParsedModelInstructions:
        """
        Turn the data generation job instructions provided by the user into a list of strings that can be used 
        to generate synthetic data through Synthex.
        Args:
            user_instructions (Any): data generation instructions provided by the user.
            language (str): The language to use for generating the training dataset.
        Returns:
            ParsedModelInstructions: the parsed user instructions.
        """
        pass
    
    @abstractmethod
    def _get_data_gen_instr(self, user_instr: ParsedModelInstructions) -> list[str]:
        """
        Generate data generation instructions by combining system instructions with user-provided
        instructions.
        Args:
            user_instr (ParsedModelInstructions): A list of user instructions where the last element is the
                domain string, and preceding elements are class names and their descriptions.
        Returns:
            list[str]: A list containing the formatted system instructions followed by the
                class-related instructions (all elements except the domain).
        """
        
        pass
    
    @abstractmethod
    def _post_process_synthetic_dataset(self, synthetic_dataset_path: str) -> None:
        f"""
        A function to perform custom post processing logic on the synthetic dataset generated by 
        Synthex, prior to training the model.
        Args:
            synthetic_dataset_path (str): The path to the synthetic dataset file.
        """
        pass

    @abstractmethod
    def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> DatasetDict:
        """
        Load the generated synthetic dataset from the specified path into a `datasets.Dataset` and prepare it 
        for training.
        Args:
            synthetic_dataset_path (str): The path to the synthetic dataset file.
        Returns:
            DatasetDict: A `datasets.DatasetDict` object containing the synthetic data, split into training 
                and validation sets.
        """
        pass
        
    @abstractmethod
    def _perform_train_pipeline(
        self, user_instructions: ParsedModelInstructions, output_path: str, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3,
        train_datapoint_examples: Optional[list[dict[str, Any]]] = None,
        device: Optional[int] = None
    ) -> TrainOutput:
        f"""
        Perform the actual model training using the provided user instructions and training configuration.
        Args:
            user_instructions (ParsedModelInstructions): A ParsedModelInstructions object containing user 
                instruction strings to be used for generating the training dataset.
            output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved.
            num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to 
                {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}.
            num_epochs (Optional[int]): The number of training epochs. Defaults to 3.
            train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to guide 
                the synthetic data generation.
            device (Optional[int]): The device to perform training on. If None, it will use the GPU
                if available, otherwise it will use the CPU.
        Returns:
            TrainOutput: The output object containing training results and metrics.
        """
        pass
    
    # TODO: we need to centralize device management inside BaseModel.train.
    @abstractmethod
    def train(
        self, 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,
        *args: Any, **kwargs: Any, 
    ) -> TrainOutput:
        f"""
        Public entrypoint to train the model.
        NOTE: The only logic that should be implemented by any concrete methods of this abstract method is the 
        transformation of use-provided instructions into Synthex-specific instructions. Once this is done, a call 
        must be made to a concrete `_perform_train_pipeline` method, which is where the actual training logic 
        must be implemented.
        Args:
            language (str): The language to use for generating the training dataset. Defaults to "english".
            output_path (str, optional): Path to save the trained model or outputs.
            num_samples (int, optional): Number of synthetic data points to generate for training. Defaults to 
                {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}.
            num_epochs (int, optional): Number of training epochs. 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 result of the training process, including metrics and model artifacts.
        """
        pass
    
    # TODO: we need to centralize device management inside BaseModel.__call__.
    @abstractmethod
    def __call__(
        self, device: Optional[int] = None, disable_logging: Optional[bool] = False, 
        *args: Any, **kwargs: Any
    ) -> Any:
        """
        Perform inference.
        Args:
            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:
            Any: The inference results.
        """
        pass
    
    @abstractmethod
    def _load_model(self, model_path: str) -> None:
        """
        Execute all logic necessary to load a pre-trained model from the specified path.
        Args:
            model_path (str): The path to the pre-trained model.
        """
        pass
    
    ##### Methods #####
    
    @property
    def _synthex(self) -> Synthex:
        """
        The Synthex instance used to generate synthetic data for training the model.
        """
        return self._synthex_val
    
    @property
    def _tokenizer(self) -> PreTrainedTokenizerBase:
        """
        The tokenizer used to preprocess text data for the model.
        """
        return self._tokenizer_val
    
    @property
    def _model(self) -> Optional[PreTrainedModel]:
        """
        A trainable model to perform inference with. It may be None if the task characteristics 
        make it impossible to select a model at instantiation time (e.g., in a multiclass 
        classification task, where the number of classes is not known upfront).
        """
        return self._model_val
    
    @_model.setter
    def _model(self, model: PreTrainedModel) -> None:
        """
        Set the trainable model to perform inference with.
        Args:
            model (PreTrainedModel): The model to set.
        """
        self._model_val = model
        
    @staticmethod
    def _determine_default_device() -> int:
        """
        Determine the default device to use for inference and training
        Returns:
            int: The device to use for inference. -1 for CPU or MPS, 0 for GPU.
        """
        
        if torch.cuda.is_available():
            return 0  # Use the first GPU
        elif torch.backends.mps.is_available():
            return -1  # Use MPS (Apple Silicon)
        else:
            return -1  # Use CPU
        
    @staticmethod
    def _should_disable_cuda(device: Optional[int]) -> bool:
        """
        Determine whether CUDA should be turned off based on the provided device.
        Args:
            device (Optional[int]): The device to use for inference/training.
        Returns:
            bool: True if CUDA should be turned off, False otherwise.
        """
        
        return device == -1
    
    @staticmethod
    def _sanitize_output_path(output_path: Optional[str] = None) -> str:
        """
        Ensure that the output path is valid; if it is not, sanitize it.
        Args:
            output_path (str): The output path to sanitize.
        Returns:
            str: The sanitized output path.
        """
        
        if output_path is not None:
            # If a filename is provided in the output_path, raise an error.
            if "." in os.path.basename(output_path):
                directory_part = os.path.dirname(output_path)
                raise ValidationError(
                    message=f"The output_path parameter must be a directory path, not a file path. Try with: '{directory_part}'."
                )
        else:
            # If output_path is None, set it to an empty string
            output_path = ""

        if output_path.strip() != "":
            output_path = output_path.strip()
            if not output_path.endswith("/"):
                output_path += "/"
        else:
            output_path = config.DEFAULT_OUTPUT_PATH
            
        return output_path
    
    def _generate_synthetic_data(
        self, schema_definition: JobOutputSchemaDefinition, requirements: list[str], 
        output_path: str, num_samples: int, examples: Optional[list[dict[str, Any]]] = None
    ) -> str:
        """
        Use Synthex to generate synthetic data based on the provided requirements.
        Args:
            requirements (list[str]): A list of requirements for the synthetic data generation.
            output_path (str): The path where the generated synthetic data will be saved.
            num_samples (int): The number of synthetic data samples to generate.
        Returns:
            str: The ID of the newly created data generation job.
        """
        
        # Trigger the data generation job in Synthex.
        job_creation_response = self._synthex.jobs.generate_data(
            schema_definition=schema_definition,
            examples=examples or [],
            requirements=requirements,
            output_path=output_path,
            number_of_samples=num_samples,
            output_type="csv"
        )
            
        return job_creation_response.job_id

    def _await_data_generation(
        self,
        # A function which takes a job ID and returns a JobStatusResponseModel.
        get_status_fn: Callable[[str], JobStatusResponseModel],
        job_id: str,
        check_interval: float = 10.0,
    ) -> JobStatusResponseModel:
        """
        Polls the synthetic data generation job status until progress reaches 1.0, or a timeout is reached.
        Use a progress bar to visualize the progress of the data generation.
        Args:
            get_status_fn (Callable[[str], JobStatusResponseModel]): A function that returns a 
                synthex.JobStatusResponseModel object.
            check_interval (float): Time to wait between checks (in seconds).
            timeout (float): Maximum time to wait for progress updates (in seconds); if no progress is made
                within this time, the job is assumed to have errored out and a TimeoutError is raised.
        Returns:
            JobStatusResponseModel: The final job status object with progress == 1.0.
        """
                
        # Initialize a progress bar using Rich.
        with Progress(transient=True) as progress:
            task = progress.add_task("Generating training data...", total=1)
            
            status = get_status_fn(job_id)

            while status.status not in [JobStatus.COMPLETED, JobStatus.FAILED]:
                time.sleep(check_interval)
                progress.update(task, completed=status.progress)
                status = get_status_fn(job_id)
                
        console.print("[green]✔ Generating training data [/green]")
        
        # If the data generation job resulted in an error, raise an exception.      
        if status.status == JobStatus.FAILED:
            raise ServerError(
                message=config.DATA_GENERATION_ERROR
            )

        return status

    def _tokenize_dataset(self, dataset: DatasetDict, token_keys: list[str]) -> DatasetDict:
        """
        Tokenize the dataset using a pre-trained tokenizer.
        Args:
            dataset (DatasetDict): The dataset to be tokenized.
            token_keys (list[str]): The keys in the dataset to tokenize.
        Returns:
            DatasetDict: The tokenized dataset.
        """

        def tokenize(example: dict[str, Sequence[str]]) -> BatchEncoding:
            inputs = [example[token_key] for token_key in token_keys]
            if len(inputs) > 2:
                raise ValidationError(
                    message="Tokenization for more than two input keys is not supported."
                )
            return self._tokenizer(
                text=list(inputs[0]),
                text_pair=list(inputs[1]) if len(inputs) == 2 else None,
                truncation=True,
                padding="max_length", 
                max_length=config.DEFAULT_TOKENIZER_MAX_LENGTH
            )

        return dataset.map(tokenize, batched=True)

    def _build_tokenized_train_ds(
        self, user_instructions: ParsedModelInstructions, output_path: str,
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, 
        train_datapoint_examples: Optional[list[dict[str, Any]]] = None
    ) -> DatasetDict:
        """
        Build a training dataset by generating synthetic data based on user-provided instructions and 
        system instructions, then tokenize it.
        Args:
            user_instructions (ParsedModelInstructions): A list of instructions, provided by the user, for 
                generating synthetic data.
            output_path (Optional[str]): The path where the generated synthetic data will be saved.
            num_samples (int): The number of training data samples to generate.
            train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints 
                to guide the synthetic data generation.
        Returns:
            DatasetDict: The tokenized dataset ready for training.
        """
        
        output_dataset_path = get_dataset_output_path(output_path)

        # Build the data generation instructions by combining user instructions and system instructions
        # NOTE: the system instructions MUST be prepended to the user instructions, as they provide 
        # context for the data generation.
        full_instructions = self._get_data_gen_instr(user_instructions)

        # Generate synthetic data.
        job_id = self._generate_synthetic_data(
            schema_definition=self._synthetic_data_schema,
            requirements=full_instructions,
            output_path=output_dataset_path,
            num_samples=num_samples,
            examples=train_datapoint_examples
        )

        # Await the completion of the synthetic data generation job.
        self._await_data_generation(
            get_status_fn=self._synthex.jobs.status, job_id=job_id
        )

        with console.status("Creating training dataset..."):
            # Once the job is complete, clean up the synthetic dataset (which may contain errors or
            # inaccurate data).
            self._post_process_synthetic_dataset(output_dataset_path)
            
            # Turn synthetic data into a training dataset with train/test split.
            dataset = self._synthetic_to_training_dataset(output_dataset_path)
            
            # Tokenize the dataset.
            tokenized_dataset = self._tokenize_dataset(dataset, self._token_keys)
        console.print(f"[green]✔ Creating training dataset[/green]")

        return tokenized_dataset

    def _train_pipeline(
        self, user_instructions: ParsedModelInstructions, output_path: Optional[str] = None, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3,
        train_datapoint_examples: Optional[list[dict[str, Any]]] = None,
        device: Optional[int] = None
    ) -> TrainOutput:
        f"""
        NOTE: This method contains training-related logic that is common across all models. As such, it must 
        be called by any concrete train function, after user instruction parsing, if any, has been performed.
        
        This method does the following:
        1. Validate the output_path parameter and silently sanitizes it if necessary.
        2. Validate the train_datapoint_examples parameter, if provided.
        3. Call the concrete `_perform_train_pipeline` method to perform the actual model training.
        4. Print a success message with the model output path.
        Args:
            user_instructions (list[str]): A list of user instruction strings to be used for generating the training 
                dataset.
            output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved.
            num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to 
                {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}.
            num_epochs (Optional[int]): The number of training epochs. Defaults to 3.
            train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to guide 
                the synthetic data generation.
            device (Optional[int]): The device to perform training on. If None, it will use the GPU
                if available, otherwise it will use the CPU.
        Returns:
            TrainOutput: The output object containing training results and metrics.
        """
        
        # Sanitize the output path provided by the user.
        sanitized_output_path = self._sanitize_output_path(output_path)
        
        # If the output path already exists, raise an error.
        if os.path.exists(sanitized_output_path):
            raise BadRequestError(
                message=f"The specified output_path already exists. Please provide a different path."
            )
        
        # Validate train_datapoint_examples, if provided.
        if train_datapoint_examples is not None:
            # Each dictionary in the train_datapoint_examples must have exactly the same keys as the 
            # synthetic data schema.
            if not all(
                [ set(self._synthetic_data_schema.keys()) == set(example.keys()) 
                    for example in train_datapoint_examples ]
            ):
                raise BadRequestError(
                    message=f"Each dictionary in the train_datapoint_examples must have exactly the following keys: {list(self._synthetic_data_schema.keys())}."
                )

        out = self._perform_train_pipeline(
            user_instructions=user_instructions,
            output_path=sanitized_output_path,
            num_samples=num_samples,
            num_epochs=num_epochs,
            train_datapoint_examples=train_datapoint_examples,
            device=device
        )

        # Get model output path based on the sanitized output path and print a success message
        model_output_path = get_model_output_path(sanitized_output_path)
        console.print(f"\n🚀 Model generation complete!\n➡️  Find your new model at {model_output_path}")

        return out
    
    def load(self, model_path: str) -> None:
        """
        Load a pre-trained model from the specified path.
        Args:
            model_path (str): The path to the pre-trained model.
        """
        
        # If the specified path does not exist, raise an error.
        if not os.path.exists(model_path):
            raise OSError(
                f"The specified model path '{model_path}' does not exist."
            )
            
        # If the specified path does not contain the necessary files, raise an error.
        expected_files = ["config.json", "model.safetensors"]
        for file in expected_files:
            if not os.path.exists(os.path.join(model_path, file)):
                raise OSError(
                    f"The specified model path '{model_path}' is missing the required file '{file}'."
                )
        
        self._load_model(model_path)