from tqdm import tqdm
from time import sleep
import logging
from pynonymizer.database.provider import DatabaseProvider, SEED_TABLE_NAME
from pynonymizer.database.exceptions import UnsupportedTableStrategyError
from pynonymizer.database.mysql import execution, query_factory
from pynonymizer.database.basic.input import resolve_input
from pynonymizer.database.basic.output import resolve_output
from pynonymizer.strategy.table import TableStrategyTypes


class MySqlProvider(DatabaseProvider):
    """
    A command-line based mysql provider. Uses `mysql` and `mysqldump`,
    Because of the efficiency of piping mass amounts of sql into the command-line client.
    Unfortunately, this implementation provides limited feedback when things go wrong.
    """

    __CHUNK_SIZE = 8192
    __DUMPSIZE_ESTIMATE_INFLATION = 1.15
    logger = logging.getLogger(__name__)

    def __init__(
        self,
        db_host,
        db_user,
        db_pass,
        db_name,
        seed_rows,
        db_port=None,
        cmd_opts=None,
        dump_opts=None,
    ):
        if db_host is None:
            db_host = "127.0.0.1"
        if db_port is None:
            db_port = "3306"
        if cmd_opts is None:
            cmd_opts = ""
        if dump_opts is None:
            dump_opts = ""

        self.db_host = db_host
        self.db_user = db_user
        self.db_pass = db_pass
        self.db_name = db_name
        self.db_port = db_port

        self.seed_rows = int(seed_rows)

        self.__runner = execution.MySqlCmdRunner(
            db_host, db_user, db_pass, db_name, db_port, additional_opts=cmd_opts
        )
        self.__dumper = execution.MySqlDumpRunner(
            db_host, db_user, db_pass, db_name, db_port, additional_opts=dump_opts
        )

    def __seed(self, qualifier_map):
        """
        'Seed' the database with a bunch of pre-generated random records so updates can be performed in batch updates
        """
        for i in tqdm(
            range(0, self.seed_rows), desc="Inserting seed data", unit="rows"
        ):
            self.logger.debug(f"Inserting seed row {i}")
            self.__runner.db_execute(
                query_factory.get_insert_seed_row(SEED_TABLE_NAME, qualifier_map)
            )

    def __estimate_dumpsize(self):
        """
        Makes a guess on the dump size using internal database metrics
        :return: A value in bytes, or None (unknown)
        """
        statement = query_factory.get_dumpsize_estimate(self.db_name)
        process_output = self.__runner.get_single_result(statement)

        try:
            return int(process_output) * self.__DUMPSIZE_ESTIMATE_INFLATION
        except ValueError:
            # Value unparsable, likely NULL
            return None

    def __read_until_empty_byte(self, data):
        return iter(lambda: data.read(self.__CHUNK_SIZE), b"")

    def __run_scripts(self, script_list, title=""):
        for i, script in enumerate(script_list):
            self.logger.info(f'Running {title} script #{i} "{script[:50]}"')
            self.logger.info(self.__runner.db_execute(script))

    def create_database(self):
        """Create the working database"""
        self.__runner.execute(query_factory.get_create_database(self.db_name))

    def drop_database(self):
        """Drop the working database"""
        self.__runner.execute(query_factory.get_drop_database(self.db_name))

    def anonymize_database(self, database_strategy):
        """
        Anonymize a restored database using the passed database strategy
        :param database_strategy: a strategy.DatabaseStrategy configuration
        :return:
        """
        qualifier_map = database_strategy.fake_update_qualifier_map

        if len(qualifier_map) > 0:
            self.logger.info("creating seed table with %d columns", len(qualifier_map))
            create_seed_table_sql = query_factory.get_create_seed_table(
                SEED_TABLE_NAME, qualifier_map
            )
            self.__runner.db_execute(create_seed_table_sql)

            self.logger.info("Inserting seed data")
            self.__seed(qualifier_map)

        self.__run_scripts(database_strategy.before_scripts, "before")

        table_strategies = database_strategy.table_strategies
        self.logger.info("Anonymizing %d tables", len(table_strategies))

        anonymization_errors = []

        with tqdm(
            desc="Anonymizing database", total=len(table_strategies)
        ) as progressbar:
            for table_strategy in table_strategies:
                try:
                    if table_strategy.schema is not None:
                        self.logger.warning(
                            "%s: MySQL provider does not support table schema. This option will be ignored.",
                            table_strategy.table_name,
                        )

                    if table_strategy.strategy_type == TableStrategyTypes.TRUNCATE:
                        progressbar.set_description(
                            "Truncating {}".format(table_strategy.table_name)
                        )
                        self.__runner.db_execute(
                            query_factory.get_truncate_table(table_strategy.table_name)
                        )

                    elif table_strategy.strategy_type == TableStrategyTypes.DELETE:
                        progressbar.set_description(
                            "Deleting {}".format(table_strategy.table_name)
                        )
                        self.__runner.db_execute(
                            query_factory.get_delete_table(table_strategy.table_name)
                        )

                    elif (
                        table_strategy.strategy_type
                        == TableStrategyTypes.UPDATE_COLUMNS
                    ):
                        progressbar.set_description(
                            "Anonymizing {}".format(table_strategy.table_name)
                        )
                        statements = query_factory.get_update_table(
                            SEED_TABLE_NAME, table_strategy
                        )
                        self.__runner.db_execute(statements)

                    else:
                        raise UnsupportedTableStrategyError(table_strategy)
                except Exception as e:
                    anonymization_errors.append(e)
                    self.logger.exception(
                        f"Error while anonymizing table {table_strategy.qualified_name}"
                    )

                progressbar.update()

        if len(anonymization_errors) > 0:
            raise Exception("Error during anonymization")

        self.__run_scripts(database_strategy.after_scripts, "after")

        self.logger.info("dropping seed table")
        self.__runner.db_execute(query_factory.get_drop_seed_table(SEED_TABLE_NAME))

        # Wait an arbitrary amount of time here to prevent this step from interacting with
        # transactional dump operations
        self.logger.debug("Waiting for trailing operations to complete...")
        sleep(0.2)

    def restore_database(self, input_path):
        """
        Feed a mysqldump dumpfile to the mysql binary on stdin.
        :param input_path:
        :return:
        """
        input_obj = resolve_input(input_path)
        dumpsize = input_obj.get_size()

        try:
            batch_processor = self.__runner.open_batch_processor()
            with input_obj.open() as dumpfile_data:
                with tqdm(
                    desc="Restoring",
                    total=dumpsize,
                    unit="B",
                    unit_scale=True,
                    unit_divisor=1000,
                ) as bar:
                    for chunk in self.__read_until_empty_byte(dumpfile_data):
                        batch_processor.write(chunk)
                        batch_processor.flush()
                        bar.update(len(chunk))
        finally:
            self.__runner.close_batch_processor()

    def dump_database(self, output_path):
        """
        Feed an output with stdout from the mysqldump binary
        :param output_path:
        :return:
        """
        output_obj = resolve_output(output_path)
        dumpsize_estimate = self.__estimate_dumpsize()

        dump_process = self.__dumper.open_dumper()
        with output_obj.open() as output_file:
            with tqdm(
                desc="Dumping",
                total=dumpsize_estimate,
                unit="B",
                unit_scale=True,
                unit_divisor=1000,
            ) as bar:
                for chunk in self.__read_until_empty_byte(dump_process):
                    output_file.write(chunk)
                    bar.update(len(chunk))
