import ast
import pandas as pd

from typing import Union, Dict, Tuple, Optional, List
from pathlib import Path
from sklearn.utils import resample

from emissary_ml_toolkit.config.dataset import DatasetConfig, ProcessedOutputConfig
from emissary_ml_toolkit.storage.manager import StorageManager
from emissary_ml_toolkit.webhooks.client import WebhookClient
from emissary_ml_toolkit.utils.exceptions import DatasetError


class DataPreProcessor:
    def __init__(self, storage_manager: StorageManager, webhook_client: WebhookClient = None):
        self.storage_manager = storage_manager
        self.webhook_client = webhook_client

    @staticmethod
    def load_raw_data(file_path: Union[Path, str]) -> pd.DataFrame:
        from .file_loaders import load_file
        return load_file(file_path)

    @staticmethod
    def split_data(data: pd.DataFrame,
                   split_ratio: float,
                   seed: int = 42) -> Tuple[pd.DataFrame, pd.DataFrame]:
        train_size = int(len(data) * (1 - split_ratio))
        shuffled = data.sample(frac=1, random_state=seed).reset_index(drop=True)
        return shuffled[:train_size], shuffled[train_size:]

    @staticmethod
    def downsample_data(data: pd.DataFrame, sample_size: float, seed: int = 42) -> pd.DataFrame:
        return data.sample(frac=sample_size, random_state=seed).reset_index(drop=True)

    @staticmethod
    def rebalance_data(data: pd.DataFrame, rebalance_config: Dict[str, Union[int, float]]) -> pd.DataFrame:
        def _parse_label_dict(label_data: Union[str, Dict]) -> Dict[str, int]:
            if isinstance(label_data, str):
                return ast.literal_eval(label_data)
            return label_data

        def _rebalance(pre_data: pd.DataFrame, label_column: str, config: Dict[str, Union[int, float]]) -> pd.DataFrame:
            rebalance_values = list(config.values())
            if all(isinstance(v, int) for v in rebalance_values):
                mode = "absolute"
            elif all(isinstance(v, (int, float)) for v in rebalance_values):
                mode = "fraction"
            else:
                raise DatasetError(f"All rebalance values must be either int or float")

            dataset_labels = set(pre_data[label_column].unique())
            config_labels = set(config.keys())

            # Validation
            if dataset_labels != config_labels:
                missing_in_config: set[str] = dataset_labels - config_labels
                missing_in_data: set[str] = config_labels - dataset_labels
                error_msg: list[str] = []
                if missing_in_config:
                    error_msg.append(f"Labels in data but not in config: {missing_in_config}")
                if missing_in_data:
                    error_msg.append(f"Labels in config but not in data: {missing_in_data}")
                raise DatasetError(". ".join(error_msg))

            grouped_data = pre_data.groupby(label_column)
            balanced_dfs: List[pd.DataFrame] = []

            if mode == "absolute":
                for label, group in grouped_data:
                    target_size = config[str(label)]
                    current_size = len(group)

                    if current_size == target_size:
                        balanced_dfs.append(group)
                    elif current_size > target_size:
                        downsampled = resample(group,
                                               n_samples=target_size,
                                               random_state=42,
                                               replace=False)
                        balanced_dfs.append(downsampled)
                    else:
                        upsampled: pd.DataFrame = resample(group,
                                                           n_samples=target_size,
                                                           random_state=42,
                                                           replace=True)
                        balanced_dfs.append(upsampled)
            else:
                total_size = len(pre_data)
                for label, group in grouped_data:
                    fraction = config[str(label)]
                    target_size = int(fraction * total_size)
                    current_size = len(group)

                    if current_size == target_size:
                        balanced_dfs.append(group)
                    elif current_size > target_size:
                        downsampled = resample(group,
                                               n_samples=target_size,
                                               random_state=42,
                                               replace=False)
                        balanced_dfs.append(downsampled)
                    else:
                        upsampled = resample(group,
                                             n_samples=target_size,
                                             random_state=42,
                                             replace=True)
                        balanced_dfs.append(upsampled)

            balanced_df = pd.concat(balanced_dfs, ignore_index=True)
            balanced_df = balanced_df.sample(frac=1, random_state=42).reset_index(drop=True)

            return balanced_df


        df = data.copy()
        df['label_for_rebalance'] = df['completion'].apply(
            lambda x: next(k for k, v in _parse_label_dict(x).items() if v == 1)
            if isinstance(x, (str, dict)) else x
        )

        df = _rebalance(df, "label_for_rebalance", rebalance_config)

        rebalanced_data = df[['prompt', 'completion']]

        return rebalanced_data

    # TODO: Improve
    def upload_and_notify(self,
                          output_config: ProcessedOutputConfig,
                          train_dataset: pd.DataFrame,
                          test_dataset: Optional[pd.DataFrame] = None):
        # Upload location provided in dataset's processed_output
        storage = self.storage_manager.get_storage(output_config.storage_ref)
        train_dataset_path = f"{output_config.base_path}/{output_config.train_dataset}"
        storage.upload_dataframe(train_dataset, train_dataset_path)
        if test_dataset is not None:
            test_dataset_path = f"{output_config.base_path}/{output_config.test_dataset}"
            storage.upload_dataframe(test_dataset, test_dataset_path)

        # Send webhook to the platform with event type datasets.uploaded
        train_dataset_id = output_config.train_dataset.split('/')[0]
        test_dataset_id = output_config.test_dataset.split('/')[0] if test_dataset is not None else None
        self.webhook_client.send_dataset_upload(train_dataset_id)
        if test_dataset_id:
            self.webhook_client.send_dataset_upload(test_dataset_id)

        return

    def preprocess_dataset(self,
                           train_dataset: DatasetConfig,
                           test_dataset: Optional[DatasetConfig] = None) -> Tuple[pd.DataFrame, Optional[pd.DataFrame]]:
        storage = self.storage_manager.get_storage(train_dataset.storage_ref)
        train_dataset_path = storage.download_file(train_dataset.storage_path)
        train_data = self.load_raw_data(train_dataset_path)
        if test_dataset:
            test_dataset_path = storage.download_file(test_dataset.storage_path)
            test_data = self.load_raw_data(test_dataset_path)
        else:
            test_data = None

        if not train_dataset.preprocessing:
            # skip preprocessing
            return train_data, test_data

        if train_dataset.preprocessing.split_ratio:
            train_data, test_data = self.split_data(train_data, train_dataset.preprocessing.split_ratio / 100)
        if train_dataset.preprocessing.downsample:
            train_data = self.downsample_data(train_data, train_dataset.preprocessing.downsample / 100)
        if train_dataset.preprocessing.rebalance:
            train_data = self.rebalance_data(train_data, train_dataset.preprocessing.rebalance)

        # Uploading processed train_data and test_data to S3 and send webhook to the platform
        output_config = train_dataset.processed_output
        self.upload_and_notify(output_config, train_data, test_data)

        return train_data, test_data

