from typing import Dict, List, Literal
import os
from dataclasses import dataclass, field
import json
from collections import defaultdict

from slugify import slugify
from loguru import logger
from syngen.ml.data_loaders import MetadataLoader, DataLoader
from syngen.ml.validation_schema import ValidationSchema, ReportTypes
from syngen.ml.utils import ValidationError


@dataclass
class Validator:
    """
    Class for validating the metadata
    """

    metadata: Dict
    metadata_path: str
    type_of_process: Literal["train", "infer"]
    validation_source: bool = True
    type_of_fk_keys = ["FK"]
    infer_report_types: List[str] = field(
        default_factory=lambda: ReportTypes().infer_report_types
    )
    merged_metadata: Dict = field(default_factory=dict)
    mapping: Dict = field(default_factory=dict)
    existed_columns_mapping: Dict = field(default_factory=dict)
    errors = defaultdict(defaultdict)

    @staticmethod
    def _check_encryption_process():
        """
        Check if the encryption process is enabled
        """
        fernet_key = os.getenv("FERNET_KEY")
        state = "is" if fernet_key else "isn't"
        process_status = "enabled" if fernet_key else "disabled"
        message = (
            f"As the environment variable 'FERNET_KEY' {state} set, "
            f"the encryption and decryption process is {process_status}."
        )
        logger.warning(message)

    def _launch_validation_of_schema(self):
        """
        Launch the validation of the schema of the metadata
        """
        return ValidationSchema(
            metadata=self.metadata,
            metadata_path=self.metadata_path,
            validation_source=self.validation_source,
            process=self.type_of_process
        ).validate_schema()

    def _define_mapping(self):
        """
        Define the mapping contained the information of the foreign keys
        defined in the metadata
        """
        for table_name, table_metadata in self.metadata.items():
            if table_name == "global":
                continue
            metadata_keys = (
                table_metadata.get("keys")
                if "keys" in table_metadata and table_metadata.get("keys") is not None
                else {}
            )
            for key_name, key_data in metadata_keys.items():
                if key_data["type"] not in self.type_of_fk_keys:
                    continue

                self.mapping[key_name] = {
                    "parent_table": key_data["references"]["table"],
                    "parent_columns": key_data["references"]["columns"],
                }

    def _check_conditions(self, metadata: Dict) -> bool:
        """
        Check conditions whether to launch validation or not
        """
        reports = metadata.get("train_settings", {}).get("reports", [])
        return (
            self.type_of_process == "infer"
            or (
                self.type_of_process == "train" and
                any([item in self.infer_report_types for item in reports])
            )
        )

    def _validate_metadata(self, table_name: str):
        """
        Validate the metadata
        """
        metadata_of_the_table = self.metadata[table_name]
        table_keys = metadata_of_the_table.get("keys", {})
        for key, config in table_keys.items():
            if config["type"] not in self.type_of_fk_keys:
                continue
            self._validate_referential_integrity(
                fk_name=key,
                fk_config=config,
                parent_config=self.merged_metadata[self.mapping[key]["parent_table"]],
            )
            parent_table = self.mapping[key]["parent_table"]
            if parent_table not in self.metadata:
                if self._check_conditions(metadata_of_the_table):
                    self._check_completion_of_training(parent_table)
                    self._check_existence_of_generated_data(parent_table)
                elif self.type_of_process == "train":
                    self._check_completion_of_training(parent_table)
            else:
                continue

    def _validate_referential_integrity(self, fk_name: str, fk_config: Dict, parent_config: Dict):
        """
        Validate whether the columns related to the primary key are the same as
        the referenced columns of the foreign key
        """
        keys = parent_config.get("keys", {})
        if not keys or all(config["type"] not in ["PK", "UQ"] for config in keys.values()):
            message = (
                "The information about columns of the primary or unique key "
                f"associated with the columns of the '{fk_config['type']}' key - '{fk_name}' "
                "wasn't found in the metadata of the parent table"
            )
            self.errors["validate referential integrity"][fk_name] = message
            return

        result = any(
            [
                config["columns"] == fk_config["references"]["columns"]
                for config in parent_config.get("keys", {}).values()
                if config["type"] in ["PK", "UQ"]
            ]
        )
        if result is False:
            message = (
                f"The columns of primary or unique key associated with the columns of "
                f"the '{fk_config['type']}' key - '{fk_name}' aren't the same"
            )
            self.errors["validate referential integrity"][fk_name] = message

    def _check_existence_of_generated_data(self, parent_table: str):
        """
        Check if the generated data of the certain parent table exists.
        The generated data is created after the successful execution of the inference process
        of the certain table.
        """
        destination = (
            self.merged_metadata[parent_table].get("infer_settings", {}).get("destination")
        )
        if destination is None:
            destination = (
                f"model_artifacts/tmp_store/{slugify(parent_table)}/"
                f"merged_infer_{slugify(parent_table)}.csv"
            )
        if not DataLoader(destination).has_existed_path:
            message = (
                f"The generated data of the table - '{parent_table}' hasn't been generated. "
                f"Please, generate the data related to the table '{parent_table}' first"
            )
            self.errors["check existence of the generated data"][parent_table] = message

    def _check_existence_of_source(self, table_name: str):
        """
        Check if the source of the certain table exists
        """
        if not DataLoader(
            self.merged_metadata[table_name]["train_settings"]["source"]
        ).has_existed_path:
            message = (
                f"It seems that the path to the source of the table - '{table_name}' "
                f"isn't correct. Please, check the path to the source of the table - "
                f"'{table_name}'"
            )
            self.errors["check existence of the source"][table_name] = message

    def _check_existence_of_destination(self, table_name: str):
        """
        Check if the destination of the certain table exists
        """
        destination = self.merged_metadata[table_name].get("infer_settings", {}).get("destination")
        if destination is None:
            logger.warning(
                f"As the destination path wasn't specified for the table - "
                f"'{table_name}', the synthetic data will be stored "
                f"at the default path - './model_artifacts/tmp_store/{slugify(table_name)}/"
                f"merged_infer_{slugify(table_name)}.csv'"
            )
        if destination is not None and not DataLoader(destination).has_existed_destination:
            message = (
                f"It seems that the directory path for storing the generated data of table "
                f"'{table_name}' isn't correct. Please, verify the destination path"
            )
            self.errors["check existence of the destination"][table_name] = message

    def _check_completion_of_training(self, table_name: str):
        """
        Check if the training process of a specific table has been completed.

        Args:
        table_name (str): The name of the table to check.

        Raises:
        FileNotFoundError: If the success file does not exist.
        ValueError: If the content of the success file does not indicate success.
        """
        path_to_success_file = f"model_artifacts/resources/{slugify(table_name)}/message.success"
        error_message = (
            f"The training of the table - '{table_name}' hasn't been completed. "
            "Please, retrain the table."
        )

        if os.path.exists(path_to_success_file):
            with open(path_to_success_file, "r") as file:
                content = file.read().strip()

        if not os.path.exists(path_to_success_file) or content != "SUCCESS":
            self.errors["check completion of the training process"][table_name] = error_message

    def _check_merged_metadata(self, parent_table: str):
        if parent_table not in self.merged_metadata:
            message = (
                f"The metadata of the parent table - '{parent_table}' hasn't been found. "
                f"Please, check whether the information of the parent table - '{parent_table}' "
                f"exists in the current metadata file or in the metadata files stored in "
                f"'model_artifacts/metadata' directory"
            )
            logger.error(message)
            raise ValidationError(message)

    def _merge_metadata(self):
        """
        Find the parent metadata contained the parent table
        in the metadata files stored in 'model_artifacts/metadata' directory,
        and merge it with the metadata of the child table
        """
        self.merged_metadata = self.metadata.copy()
        for key_name, config in self.mapping.items():
            parent_table = config.get("parent_table")
            if parent_table in self.metadata:
                continue
            path_to_metadata_storage = "model_artifacts/metadata"
            for file in os.listdir(path_to_metadata_storage):
                path_to_metadata_file = os.path.join(path_to_metadata_storage, file)
                metadata = MetadataLoader(path_to_metadata_file).load_data()
                if parent_table not in metadata:
                    continue
                self._launch_validation_of_schema()
                self.merged_metadata.update(metadata)
                logger.info(
                    f"The metadata located in the path - '{path_to_metadata_storage}' "
                    f"has been merged with the current metadata as it contains "
                    f"the information of the parent table - '{parent_table}'"
                )
            self._check_merged_metadata(parent_table)

    def _check_existence_of_key_columns(self, table_name):
        """
        Check whether the key columns exist in the table
        """
        metadata_of_table = self.merged_metadata[table_name]
        existed_columns = self.existed_columns_mapping[table_name]
        for key, config_of_key in metadata_of_table.get("keys", {}).items():
            if all([column in existed_columns for column in config_of_key["columns"]]):
                continue
            else:
                non_existed_columns = [
                    f"{col!r}"
                    for col in set(config_of_key["columns"]).difference(set(existed_columns))
                ]
                message = (
                    f"The columns of the {config_of_key['type']} key '{key}' - "
                    f"{', '.join(non_existed_columns)} "
                    f"don't exist in the table - '{table_name}'"
                )
                self.errors["check existence of the key columns in 'columns'"][key] = message

    def _check_existence_of_referenced_columns(self, table_name: str):
        """
        Check whether the key columns exist in the referenced table
        """
        metadata_of_table = self.merged_metadata[table_name]
        for key, config_of_key in metadata_of_table.get("keys", {}).items():
            if config_of_key["type"] in self.type_of_fk_keys:
                referenced_table = config_of_key["references"]["table"]
                referenced_columns = config_of_key["references"]["columns"]
                existed_columns = self.existed_columns_mapping.get(referenced_table, {})
                non_existed_columns = [
                    f"{col!r}"
                    for col in referenced_columns
                    if col not in existed_columns
                ]
                if non_existed_columns:
                    message = (
                        f"The 'references.columns' of the {config_of_key['type']} '{key}' - "
                        f"{', '.join(non_existed_columns)} "
                        f"don't exist in the referenced table - '{referenced_table}'"
                    )
                    self.errors[
                        "check existence of the key columns in 'references.columns'"
                    ][key] = message

    def _fetch_existed_columns(self, table_name: str) -> List[str]:
        """
        Fetch the list of the columns of the source table
        """
        metadata_of_table = self.merged_metadata[table_name]
        format_settings = metadata_of_table.get("format", {})
        path_to_source = self._fetch_path_to_source(table_name)
        return DataLoader(path_to_source).get_columns(**format_settings)

    def _gather_existed_columns(self, table_name: str):
        """
        Fetch the list of the existed columns of the table
        and put it into the mapping of existed columns
        """
        existed_columns = self._fetch_existed_columns(table_name)
        self.existed_columns_mapping[table_name] = existed_columns

    def _preprocess_metadata(self):
        """
        Preprocess the metadata, set the metadata and the merged metadata
        """
        self._launch_validation_of_schema()
        self._define_mapping()
        self._merge_metadata()
        self.merged_metadata.pop("global", None)
        self.metadata.pop("global", None)

    def _fetch_path_to_source(self, table_name):
        """
        Fetch the path to the source of the certain table
        """
        return self.merged_metadata[table_name]["train_settings"]["source"]

    def _launch_validation(self):
        """
        Launch the validation process
        """
        if self.type_of_process == "train" and self.validation_source:
            for table_name in self.merged_metadata.keys():
                self._gather_existed_columns(table_name)

        for table_name in self.metadata.keys():
            if self.type_of_process == "train" and self.validation_source:
                self._check_existence_of_source(table_name)
                self._check_existence_of_key_columns(table_name)
                self._check_existence_of_referenced_columns(table_name)
            elif self.type_of_process == "infer":
                self._check_completion_of_training(table_name)
                self._check_existence_of_destination(table_name)

        for table_name in self.metadata.keys():
            self._validate_metadata(table_name)

    def _collect_errors(self):
        """
        Collect the errors found during the validation process
        """
        error_logs = []
        for section, errors_details in self.errors.items():
            error_log = f'"{section}": {json.dumps(errors_details, indent=4)}'
            error_logs.append(error_log)
        if self.errors:
            message = (
                f"The validation of the metadata has been failed. "
                f"The error(s) found in - \n{', '.join(error_logs)}"
            )
            logger.error(message)
            raise ValidationError(message)
        logger.info("The validation of the metadata has been passed successfully")

    def run(self):
        """
        Run the validation process
        """
        self._check_encryption_process()
        self._preprocess_metadata()
        self._launch_validation()
        self._collect_errors()
