import os
from collections import Counter
from typing import List, Tuple, Dict, Any, Optional, Callable, Union, Literal
import json
from json import JSONDecodeError
from slugify import slugify

from loguru import logger
import pandas as pd
import numpy as np
from syngen.ml.flatten_json import unflatten_list, safe_flatten

from syngen.ml.data_loaders import DataLoader, DataFrameFetcher
from syngen.ml.utils import fetch_unique_root, fetch_config, get_source_path_extension
from syngen.ml.context import get_context


class Processor:
    """
    Base class for data preprocessing and postprocessing

    Attributes:
        metadata (Dict): Processing metadata
        metadata_path (Optional[str]): Path to metadata file
        table_name (Optional[str]): Name of the table being processed
        loader (Optional[Callable]): Custom data loader function
    """
    def __init__(
        self,
        metadata: Dict,
        metadata_path: Optional[str],
        table_name: str,
        loader: Optional[Callable[[str], pd.DataFrame]] = None
    ):
        self.metadata_path = metadata_path
        self.metadata = metadata
        self.table_name = table_name
        self.loader = loader
        self.path_to_flatten_metadata = (
            "model_artifacts/system_store/flatten_configs/"
            f"flatten_metadata_{fetch_unique_root(self.table_name, self.metadata_path)}.json"
        )


class PreprocessHandler(Processor):
    """
    The class for the preprocessing of the data before the training process
    """
    def __init__(
        self,
        metadata: Dict,
        metadata_path: Optional[str],
        table_name: str,
        loader: Optional[Callable[[str], pd.DataFrame]] = None
    ):
        super().__init__(
            metadata=metadata,
            metadata_path=metadata_path,
            table_name=table_name,
            loader=loader
        )
        self.initial_data_shape: Tuple = ()
        self.row_subset: int = int()
        self.original_df: Optional[pd.DataFrame] = None
        self.initial_order_of_columns = list()
        self.schema = dict()
        self.original_schema = None

    def _save_original_schema(self):
        """
        Save the schema of the original data
        """
        path = (
            f"model_artifacts/tmp_store/{slugify(self.table_name)}"
            f"/original_schema_{slugify(self.table_name)}.pkl"
        )
        DataLoader(path).save_data(data=self.original_schema)

    def _save_initial_order_of_columns(self):
        """
        Save the initial order of columns of the original data
        """
        path = (
            f"model_artifacts/tmp_store/{slugify(self.table_name)}"
            f"/initial_order_of_columns_{slugify(self.table_name)}.pkl"
        )
        DataLoader(path).save_data(data=self.initial_order_of_columns)

    def _check_if_data_is_empty(self):
        """
        Check if the provided data is empty or not
        """
        if self.original_df.shape[0] < 1:
            raise ValueError(
                f"The empty table was provided. Unable to train the table - '{self.table_name}'"
            )

    def _remove_empty_columns(self) -> pd.DataFrame:
        """
        Remove completely empty columns from dataframe and mark them as 'removed' in the schema
        """
        data = self.original_df.dropna(how="all", axis=1)

        dropped_columns = set(self.initial_order_of_columns) - set(data.columns)
        list_of_dropped_columns = [f"'{column}'" for column in dropped_columns]
        if list_of_dropped_columns:
            logger.info(f"Empty columns - {', '.join(list_of_dropped_columns)} were removed")
            self.schema["fields"].update({column: "removed" for column in dropped_columns})
        return data

    def prepare_data(self) -> pd.DataFrame:
        """
        Prepare the subset of the data for the training process,
        and get the preprocessed data
        """
        self._load_source()
        self._check_if_data_is_empty()
        self._save_original_schema()
        self._save_initial_order_of_columns()
        preprocessed_data = self._preprocess_data()
        return preprocessed_data

    def _preprocess_data(self) -> pd.DataFrame:
        """
        Apply the parameters 'drop_null' and 'row_limit' to the data,
        and get the subset for the training process
        """
        drop_null = self.metadata[self.table_name]["train_settings"]["drop_null"]
        row_limit = self.metadata[self.table_name]["train_settings"]["row_limit"]
        data = self._remove_empty_columns()
        if self.loader:
            warning_message = (
                "parameter will be ignored because the retrieval of the data "
                "is handled by a callback function"
            )
            if drop_null:
                logger.warning(f"The 'drop_null' {warning_message}")
            if row_limit is not None:
                logger.warning(f"The 'row_limit' {warning_message}")
        else:
            if drop_null:
                if not data.dropna().empty:
                    data = data.dropna()
                    if count_of_dropped_rows := self.initial_data_shape[0] - data.shape[0]:
                        logger.info(
                            f"As the parameter 'drop_null' set to 'True', "
                            f"{count_of_dropped_rows} rows of the table - '{self.table_name}' "
                            f"that have empty values have been dropped. "
                            f"The count of remained rows is {data.shape[0]}."
                        )
                else:
                    logger.warning(
                        "The specified 'drop_null' argument results in the empty dataframe, "
                        "so it will be ignored"
                    )
                self.row_subset = len(data)

            if row_limit:
                self.row_subset = min(row_limit, len(data))
                data = data.sample(n=self.row_subset)

        if len(data) < 100:
            logger.warning(
                "The input table is too small to provide any meaningful results. "
                "Please consider: 1) disable drop_null argument, 2) provide bigger table"
            )
        elif len(data) < 500:
            logger.warning(
                f"The amount of data is {len(data)} rows. It seems that it isn't enough "
                f"to supply high-quality results. To improve the quality of generated data "
                f"please consider any of the steps: 1) provide a bigger table, "
                f"2) disable drop_null argument"
            )
        logger.info(f"The subset of rows was set to {self.row_subset}")
        return data

    @staticmethod
    def _run_script():
        """
        Run the script before the training process
        if it exists in the predefined path
        """
        path_to_script = f"{os.getcwd()}/model_artifacts/script.py"
        if os.path.exists(path_to_script):
            os.system(f"python3 {path_to_script}")

    @staticmethod
    def _get_json_columns(data: pd.DataFrame) -> List[str]:
        """
        Get the list of columns which contain JSON data convertable to a Python dictionary
        Returns:
            List[str]: List of column names containing valid JSON data convertable to dict
        """
        def is_json_column(series: pd.Series) -> bool:
            if pd.isnull(series).all():
                return False

            def is_json_value(x):
                if not isinstance(x, str):
                    return False
                try:
                    return isinstance(json.loads(x), dict)
                except JSONDecodeError:
                    return False

            return series.dropna().apply(is_json_value).any()
        return [col for col in data.columns if is_json_column(data[col])]

    @staticmethod
    def _get_artifacts(
        data: pd.DataFrame,
        json_columns: List[str]
    ) -> Tuple[pd.DataFrame, Dict[str, List[str]], List[str]]:
        """
        Flatten JSON columns in DataFrame and create mapping of original to flattened columns.

        Args:
            data (pd.DataFrame): Input DataFrame with JSON columns
            json_columns (List[str]): List of columns containing JSON data

        Returns:
            Tuple containing:
            - Flattened DataFrame
            - Mapping of original columns to flattened columns
            - List of duplicated column names
        """
        flattened_dfs = list()
        flattening_mapping = dict()
        for column in json_columns:
            result = [safe_flatten(i) for i in data[column]]
            flattened_data = pd.DataFrame(
                [i["flattened_data"] for i in result],
                index=data.index
            )
            flattened_data[f"{column}_"] = pd.DataFrame(
                [i["original_data"] for i in result],
                index=data.index
            )
            flattening_mapping[column] = [
                c for c in flattened_data.columns.to_list()
                if c != f"{column}_"
            ]
            flattened_dfs.append(flattened_data)
        flattened_data = pd.concat([data, *flattened_dfs], axis=1)
        duplicated_columns = [
            key
            for key, value in dict(Counter(flattened_data.columns.to_list())).items()
            if value > 1
        ]
        flattened_data.drop(columns=flattening_mapping.keys(), inplace=True)
        flattened_df = flattened_data.T.loc[~flattened_data.T.index.duplicated(), :].T
        flattened_df = flattened_df.applymap(lambda x: np.NaN if x in [list(), dict()] else x)
        return flattened_df, flattening_mapping, duplicated_columns

    def _save_flatten_metadata(self, metadata: Dict):
        """
        Save the metadata of the flattening process
        """
        if os.path.exists(self.path_to_flatten_metadata):
            with open(self.path_to_flatten_metadata, "r") as f:
                existing_metadata = json.load(f)
            metadata = {**existing_metadata, **metadata}
        with open(f"{self.path_to_flatten_metadata}", "w") as f:
            json.dump(metadata, f)

    def _load_source(self):
        """
        Load the data from the predefined source
        """
        if self.loader is not None:
            data_loader = DataFrameFetcher(
                loader=self.loader,
                table_name=self.table_name
            )
        else:
            path_to_source = self.metadata[self.table_name]["train_settings"]["source"]
            data_loader = DataLoader(path=path_to_source)
        self.original_df, self.schema = data_loader.load_data()
        self.original_schema = data_loader.original_schema
        self.initial_data_shape = self.original_df.shape
        self.initial_order_of_columns = self.original_df.columns.to_list()
        self.row_subset = len(self.original_df)

    def _handle_json_columns(self, data: pd.DataFrame):
        """
        Preprocess the data contained JSON columns before the training process
        """
        flatten_metadata = dict()
        order_of_columns = data.columns.to_list()
        if json_columns := self._get_json_columns(data):
            list_of_json_columns = [f"'{column}'" for column in json_columns]
            list_of_json_columns = ', '.join(list_of_json_columns)
            logger.info(
                f"The table '{self.table_name}' contains JSON columns: {list_of_json_columns}"
            )
            logger.info(f"Flattening the JSON columns in the table - '{self.table_name}'")
            (
                flattened_data,
                flattening_mapping,
                duplicated_columns
            ) = self._get_artifacts(data, json_columns)
            flatten_metadata[self.table_name] = {
                "flattening_mapping": flattening_mapping,
                "duplicated_columns": duplicated_columns,
                "order_of_columns": order_of_columns
            }
            logger.info(f"The table '{self.table_name}' has been successfully flattened")
            self._save_flatten_metadata(flatten_metadata)
            return flattened_data
        return data

    def run(self) -> Tuple[pd.DataFrame, Dict]:
        """
        Launch the preprocessing process:
        """
        self._run_script()
        preprocessed_data = self.prepare_data()
        data = self._handle_json_columns(preprocessed_data)
        return data, self.schema


class PostprocessHandler(Processor):
    """
    Class for the postprocessing of the generated data
    """

    def __init__(
        self,
        metadata: Dict,
        metadata_path: Optional[str],
        table_name: Optional[str],
        type_of_process: Literal["train", "infer"]
    ):
        super().__init__(metadata, metadata_path, table_name)
        self.type_of_process = type_of_process
        self.flatten_metadata = self._fetch_flatten_config()
        self.preview_rows = 10

    def _fetch_flatten_config(self) -> Dict:
        """
        Fetch the metadata of the flattening process
        """
        if os.path.exists(self.path_to_flatten_metadata):
            with open(self.path_to_flatten_metadata, "r") as f:
                return json.load(f)
        return dict()

    @staticmethod
    def _check_none_values(x: Any) -> bool:
        """
        Check if the value should be treated as null/none.

        Args:
            x: Value to check

        Returns:
            bool: True if value should be treated as null/none

        Note:
            Handles None, np.NaN, and empty containers as null values
        """
        if x is None:
            return True
        if isinstance(x, float) and np.isnan(x):
            return True
        if isinstance(x, (list, dict)) and not x:
            return True
        return False

    def _remove_none_from_struct(self, input_data: Any) -> Any:
        """
        Remove np.NaN or None values from nested data structures.

        Args:
            input_data: Input data structure (list, dict, or scalar)

        Returns:
            Cleaned data structure with None values removed
        """
        if isinstance(input_data, list):
            # Process list elements and filter empty results
            return [
                self._remove_none_from_struct(v)
                for v in input_data
                if not self._check_none_values(v)
            ]

        if isinstance(input_data, dict):
            # Process dictionary items and filter empty results
            return {
                k: self._remove_none_from_struct(v)
                for k, v in input_data.items()
                if not self._check_none_values(v)
            }
        return input_data

    @staticmethod
    def _restore_empty_values(df: pd.DataFrame) -> pd.DataFrame:
        """
        Restore empty dictionary values in the dataframe for nested fields.

        Args:
            df: Input DataFrame with potentially empty nested fields

        Returns:
            pd.DataFrame: DataFrame with restored empty dictionary values
        """
        df = df.copy().astype(object)
        columns = df.columns.tolist()

        # Find potential parent columns
        parent_columns = set()
        for col in columns:
            for potential_parent in columns:
                if f"{potential_parent}." in col:
                    parent_columns.add(potential_parent)

        # Restore empty dictionary values for parent columns
        for parent_col in parent_columns:
            df[parent_col] = df[parent_col].apply(lambda val: {} if pd.isna(val) else val)

        return df

    @staticmethod
    def _load_generated_data(path_to_generated_data: str, *args) -> pd.DataFrame:
        """
        Load generated data from the predefined path
        """
        data, schema = DataLoader(path=path_to_generated_data).load_data()
        return data

    @staticmethod
    def _remove_empty_elements(d: Union[Dict, List]) -> Optional[Union[Dict, List]]:
        """
        Recursively remove keys with empty dictionaries or lists from a nested dictionary or list
        """
        def clean(data):
            if isinstance(data, dict):
                cleaned = {
                    key: cleaned_value
                    for key, value in data.items()
                    if (cleaned_value := clean(value)) not in ({}, [])
                }
                return cleaned

            if isinstance(data, list):
                cleaned = [
                    cleaned_item
                    for item in data
                    if (cleaned_item := clean(item)) not in ({}, [])
                ]
                return cleaned

            return data
        result = clean(d)
        return result if pd.notnull(result) and result not in [{}, []] else None

    def _postprocess_generated_data(
        self,
        data: pd.DataFrame,
        flattening_mapping: Dict,
        duplicated_columns: List
    ) -> pd.DataFrame:
        """
        Postprocess the generated data
        """
        for old_column, new_columns in flattening_mapping.items():
            data[new_columns] = self._restore_empty_values(data[new_columns])

            data[old_column] = data.apply(
                lambda row: unflatten_list(row[new_columns].to_dict(), ".")
                if row[new_columns].notnull().any()
                else row[f"{old_column}_"],
                axis=1
            )
            data[old_column] = data[old_column].apply(
                lambda row: self._remove_none_from_struct(row)
            )
            data[old_column] = data[old_column].apply(lambda row: self._remove_empty_elements(row))
            data[old_column] = data[old_column].apply(
                lambda row: json.dumps(row, ensure_ascii=False)
                if isinstance(row, dict)
                else row
            )
            dropped_columns = set(i for i in new_columns if i not in duplicated_columns)
            dropped_columns.add(f"{old_column}_")
            data.drop(list(dropped_columns), axis=1, inplace=True)
        return data

    def run(self):
        """
        Launch the postprocessing of the generated data,
        and save the processed data to the predefined path
        """
        for table in self.metadata.keys():
            path_to_merged_infer = fetch_config(
                f"model_artifacts/resources/{slugify(table)}/vae/checkpoints/train_config.pkl"
            ).paths["path_to_merged_infer"]
            extension = get_source_path_extension(path=path_to_merged_infer)
            default_path_to_generated_data = (
                f"model_artifacts/tmp_store/{slugify(table)}/"
                f"merged_infer_{slugify(table)}{extension}"
            )
            path_to_generated_data = (
                self.metadata[table].get("infer_settings", {}).get(
                    "destination", default_path_to_generated_data
                ) if self.type_of_process == "infer" else default_path_to_generated_data
            )
            data = self._load_generated_data(path_to_generated_data, table)
            flatten_metadata_of_table = self.flatten_metadata.get(table)
            if flatten_metadata_of_table is not None:
                logger.info("Start postprocessing of the generated data")
                flattening_mapping = flatten_metadata_of_table.get("flattening_mapping")
                duplicated_columns = flatten_metadata_of_table.get("duplicated_columns")
                data = self._postprocess_generated_data(
                    data,
                    flattening_mapping,
                    duplicated_columns
                )
                logger.info("Finish postprocessing of the generated data")
            order_of_columns = fetch_config(
                config_pickle_path=f"model_artifacts/tmp_store/{slugify(table)}"
                                   f"/initial_order_of_columns_{slugify(table)}.pkl")
            self.__save_generated_data(
                data,
                path_to_generated_data,
                order_of_columns,
                table
            )

    @staticmethod
    def __save_preview_data(
        path_to_destination,
        preview_df
    ):
        """
        Save the preview data in '.csv' format for the fast review of the generated data
        """
        preview_path = f"{os.path.splitext(path_to_destination)[0]}_preview.csv"
        DataLoader(path=preview_path).save_data(
            preview_df,
            format=get_context().get_config(),
        )
        logger.info(
            f"Preview saved in '{preview_path}' (first {len(preview_df)} rows)"
        )

    def __save_generated_data(
        self,
        generated_data: pd.DataFrame,
        path_to_destination: str,
        order_of_columns: List[str],
        *args
    ):
        """
        Save generated data to the path
        """
        generated_data = generated_data[order_of_columns]
        DataLoader(path=path_to_destination).save_data(
            generated_data,
            format=get_context().get_config(),
        )
        if self.type_of_process == "infer":
            self.__save_preview_data(
                path_to_destination=path_to_destination,
                preview_df=generated_data.head(self.preview_rows)
            )
