from abc import ABC, abstractmethod
from synthex import Synthex
from synthex.exceptions import BadRequestError as SynthexBadRequestError, RateLimitError as SynthexRateLimitError
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 # type: ignore
from typing import Callable, Sequence, Any
import os
from transformers.trainer_utils import TrainOutput
from typing import Optional
from rich.progress import Progress
from rich.console import Console

from artifex.config import config
from artifex.core import auto_validate_methods, BadRequestError, RateLimitError, ServerError
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.
    """

    ##### Abstract properties #####
    
    @property
    @abstractmethod
    def _synthex(self) -> Synthex:
        """
        An instance of the Synthex class.
        """
        pass

    @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) -> list[str]:
        """
        Instructions for the synthetic data generation job, which are not provided by the user but passed by default.
        """
        pass
    
    @property
    @abstractmethod
    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, selecting a model at instantiation is impossible).
        """
        pass
    
    @property
    @abstractmethod
    def _tokenizer(self) -> PreTrainedTokenizerBase:
        """
        The tokenizer used during training and inference.
        """
        pass
    
    @property
    @abstractmethod
    def _token_key(self) -> str:
        """
        The key in the dataset that contains the text to be tokenized. This is used to tokenize the dataset 
        during training and inference.
        """
        pass
    
    ##### Abstract methods #####
    
    @abstractmethod
    def _parse_user_instructions(self, user_instructions: Any) -> list[str]:
        """
        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.
        Returns:
            list[str]: the user instructions parsed into a list of strings.
        """
        pass
    
    @abstractmethod
    def _cleanup_synthetic_dataset(self, synthetic_dataset_path: str) -> None:
        f"""
        A function to perform custom cleaning logic on the synthetic dataset generated by Synthex. Especially on a
        {config.SYNTHEX_TIER_1_PLAN_NAME} plan, the synthetic training dataset may contain errors or inaccurate data, 
        so this function should be used to clean it up before 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: list[str], output_path: str, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3
    ) -> TrainOutput:
        f"""
        Perform the actual training of the model using the provided user instructions and training configuration.
        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.
        Returns:
            TrainOutput: The output object containing training results and metrics.
        """
        pass
    
    @abstractmethod
    def train(
        self, output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, 
        num_epochs: int = 3, *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:
            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.
        Returns:
            TrainOutput: The result of the training process, including metrics and model artifacts.
        """
        pass
    
    @abstractmethod
    def load(self, model_path: str) -> None:
        """
        Load the model from the specified file path.
        Args:
            model_path (str): The path to the model file to be loaded.
        """
        pass
    
    @abstractmethod
    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """
        Perform inference.
        Returns:
            Any: The inference results.
        """
        pass
    
    ##### Methods #####
    
    @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 None, set it to an empty string
        output_path = output_path or ""
        # If output_path is not an empty string, use a default directory
        default_dir = config.DEFAULT_OUTPUT_PATH

        if output_path != "":
            output_path = output_path.strip()
            # Extract the directory and file name from the output path, use only the directory
            directory, filename = os.path.split(output_path)
            # If the filename does not have an extension, use it as a directory
            if "." not in filename:
                directory = os.path.join(directory, filename).rstrip("/")
            # Add run identifier to the directory
            date_string = default_dir.split("/")[-2]  # Extract the date string from the default directory
            directory = f"{os.path.join(directory, date_string)}/"
        else:
            directory = default_dir
            
        return directory
    
    def _generate_synthetic_data(
        self, schema_definition: JobOutputSchemaDefinition, requirements: list[str], 
        output_path: str, num_samples: int
    ) -> 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.
        try:
            job_creation_response = self._synthex.jobs.generate_data(
                schema_definition=schema_definition,
                examples=[],
                requirements=requirements,
                output_path=output_path,
                number_of_samples=num_samples,
                output_type="csv"
            )
        except SynthexBadRequestError:
            raise BadRequestError(
                message=f"On a {config.SYNTHEX_TIER_1_PLAN_NAME} plan, `number_of_samples` must not be higher than {config.MAX_SYNTHEX_DATASET_SIZE_TIER_1_PLAN}."
            )
        except SynthexRateLimitError as e:
            current = e.details["details"]["current_monthly_datapoints"] # type: ignore
            requested = e.details["details"]["requested_datapoints"] # type: ignore
            raise RateLimitError(
                message=f"On a {config.SYNTHEX_TIER_1_PLAN_NAME} plan, you can generate up to {config.MAX_SYNTHEX_DATAPOINTS_PER_MONTH_TIER_1_PLAN} training samples per month.",
                details=f"This month you have generated {current} training samples, and you requested {requested} more."
            )
            
        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=100)
            
            status = get_status_fn(job_id)

            while status.status not in [JobStatus.COMPLETED, JobStatus.FAILED]:
                time.sleep(check_interval)
                new_status = get_status_fn(job_id)

                # If the new status has a higher progress value, update the progress bar.
                if new_status.progress > status.progress:
                    progress.update(task, advance=int(new_status.progress*100 - status.progress*100))

                status = new_status
                
        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_key: str) -> DatasetDict:
        """
        Tokenize the dataset using a pre-trained tokenizer.
        Args:
            dataset (DatasetDict): The dataset to be tokenized.
            token_key (str): The key in the dataset to tokenize.
        Returns:
            DatasetDict: The tokenized dataset.
        """

        def tokenize(example: dict[str, Sequence[str]]) -> BatchEncoding:
            return self._tokenizer(
                list(example[token_key]), truncation=True, padding="max_length", max_length=128
            )

        return dataset.map(tokenize, batched=True) # type: ignore

    def _build_tokenized_train_ds(
        self, user_instructions: list[str], output_path: str,
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM
    ) -> DatasetDict:
        """
        Build a training dataset by generating synthetic data based on user-provided instructions and 
        system instructions, then tokenize it.
        Args:
            user_instructions (list[str]): 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.
        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._system_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
        )

        # 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._cleanup_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_key)
        console.print(f"[green]✔ Creating training dataset[/green]")

        return tokenized_dataset

    def _train_pipeline(
        self, user_instructions: list[str], output_path: Optional[str] = None, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3
    ) -> TrainOutput:
        f"""
        NOTE: This method must be called by each concrete train function, after user instruction parsing, if any, has been 
        performed.
        Validate the output_path parameter and silently sanitizes it if necessary, then calls the concrete
        `_perform_train_pipeline` method to perform the actual training of the model using the provided user instructions and 
        training configuration.
        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.
        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)
        
        # Get model output path based on the sanitized output path
        model_output_path = get_model_output_path(sanitized_output_path)

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

        console.print(f"\n🚀 Model generation complete!\n➡️  Find your new model at {model_output_path}")

        return out