# Copyright Contributors to the Amundsen project.
# SPDX-License-Identifier: Apache-2.0

import logging
import time
from typing import (
    Any, Dict, Set, Type,
)

from amundsen_rds.models import RDSModel
from amundsen_rds.models.base import Base
from pyhocon import ConfigFactory, ConfigTree
from sqlalchemy import create_engine, func
from sqlalchemy.orm import sessionmaker

from databuilder import Scoped
from databuilder.task.base_task import Task

LOGGER = logging.getLogger(__name__)


class MySQLStalenessRemovalTask(Task):
    """
    A Specific task to remove stale records in MySQL. It will use either the "published_tag" or
    "publisher_last_updated_epoch_ms" attribute assigned from the MySqlCsvPublisher. If "published tag" is different
    from the one it is getting from the config or if the "publisher_last_updated_epoch_ms" is less than the cutoff,
    it will regard the record as stale. Not all resources are being published by MySQLCSVPublisher and you can only
    set specific model table names to perform this deletion.
    Note: This task performs a cascade delete and will delete all the orphan records in the child tables of the stale
    records.

    """
    # Connection string
    CONN_STRING = "conn_string"
    # If its value is true, SQLAlchemy engine will log all statements
    ENGINE_ECHO = 'engine_echo'
    # Additional arguments used for engine
    CONNECT_ARGS = 'connect_args'

    TARGET_TABLES = "target_tables"
    DRY_RUN = "dry_run"
    # Staleness max percentage. Safety net to prevent majority of data being deleted.
    STALENESS_MAX_PCT = "staleness_max_pct"
    # Staleness max percentage per table. Safety net to prevent majority of data being deleted.
    STALENESS_PCT_MAX_DICT = "staleness_max_pct_dict"
    # If the published_tag value is different from the one in config, the record is determined to be stale
    PUBLISHED_TAG = "job_publish_tag"
    # Using this milliseconds and published timestamp to determine staleness
    MS_TO_EXPIRE = "milliseconds_to_expire"
    MIN_MS_TO_EXPIRE = "minimum_milliseconds_to_expire"

    _DEFAULT_CONFIG = ConfigFactory.from_dict({STALENESS_MAX_PCT: 5,
                                              TARGET_TABLES: [],
                                              STALENESS_PCT_MAX_DICT: {},
                                              MIN_MS_TO_EXPIRE: 86400000,
                                              DRY_RUN: False,
                                              ENGINE_ECHO: False})

    def get_scope(self) -> str:
        return 'task.mysql_remove_stale_data'

    def init(self, conf: ConfigTree) -> None:
        conf = Scoped.get_scoped_conf(conf, self.get_scope()) \
            .with_fallback(conf) \
            .with_fallback(MySQLStalenessRemovalTask._DEFAULT_CONFIG)
        self.target_tables = set(conf.get_list(MySQLStalenessRemovalTask.TARGET_TABLES))
        self.target_table_model_dict = self._get_target_table_model_dict(self.target_tables)
        self.dry_run = conf.get_bool(MySQLStalenessRemovalTask.DRY_RUN)
        self.staleness_max_pct = conf.get_int(MySQLStalenessRemovalTask.STALENESS_MAX_PCT)
        self.staleness_max_pct_dict = conf.get(MySQLStalenessRemovalTask.STALENESS_PCT_MAX_DICT)

        if MySQLStalenessRemovalTask.PUBLISHED_TAG in conf and MySQLStalenessRemovalTask.MS_TO_EXPIRE in conf:
            raise Exception(f'Cannot have both {MySQLStalenessRemovalTask.PUBLISHED_TAG} and '
                            f'{MySQLStalenessRemovalTask.MS_TO_EXPIRE} in job config')

        self.ms_to_expire = None
        if MySQLStalenessRemovalTask.MS_TO_EXPIRE in conf:
            self.ms_to_expire = conf.get_int(MySQLStalenessRemovalTask.MS_TO_EXPIRE)
            if self.ms_to_expire < conf.get_int(MySQLStalenessRemovalTask.MIN_MS_TO_EXPIRE):
                raise Exception(f'{MySQLStalenessRemovalTask.MS_TO_EXPIRE} is too small')
            self.marker = self.ms_to_expire
        else:
            self.marker = conf.get_string(MySQLStalenessRemovalTask.PUBLISHED_TAG)

        connect_args = {k: v for k, v in conf.get_config(MySQLStalenessRemovalTask.CONNECT_ARGS,
                                                         default=ConfigTree()).items()}
        self._engine = create_engine(conf.get_string(MySQLStalenessRemovalTask.CONN_STRING),
                                     echo=conf.get_bool(MySQLStalenessRemovalTask.ENGINE_ECHO),
                                     connect_args=connect_args)
        self._session_factory = sessionmaker(bind=self._engine)
        self._session = self._session_factory()

    def _get_target_table_model_dict(self, target_tables: Set[str]) -> Dict[str, Type[RDSModel]]:
        """
        Returns a dictionary with a table name to the corresponding RDS model class mapping.
        :param target_tables:
        :return:
        """
        target_table_model_dict: Dict[str, Type[RDSModel]] = {}

        # Handle difference observed between Python 3.7 and other Python version runtimes:
        if hasattr(Base, '_decl_class_registry'):
            models_generator = Base._decl_class_registry.values()
        elif hasattr(Base, 'registry'):
            models_generator = Base.registry._class_registry.values()
        else:
            raise Exception(f'Failed to get models for target tables {target_tables}')

        for model in models_generator:
            if hasattr(model, '__tablename__') and model.__tablename__ in target_tables:
                target_table_model_dict[model.__tablename__] = model
        return target_table_model_dict

    def run(self) -> None:
        """
        For each table performs a safety check to make sure this operation does not delete more than a threshold
        where default threshold is 5%. Once it passes a safety check, it cascade deletes the stale records and all
        the referencing records. The target tables are traversed acc. to foreign key dependency (i.e in the order of
        child tables to parent tables) to try to keep the actual deleted record count consistent with the validation
        percentage (which reflects the staleness pct only from a single table and does not take into account the
        referenced tables data which will be deleted in a cascade delete)
        :return:
        """
        sorted_table_names = [table.name for table in Base.metadata.sorted_tables]
        sorted_target_tables = sorted(
            self.target_tables, key=lambda table: sorted_table_names.index(table), reverse=True)
        try:
            for table_name in sorted_target_tables:
                target_model_class = self.target_table_model_dict.get(table_name, None)
                if target_model_class:
                    staleness_pct = self._validate_record_staleness_pct(target_model_class=target_model_class)
                    if self.dry_run:
                        LOGGER.info('Skipping deleting records since it is a Dry Run.')
                        continue
                    if staleness_pct > 0:
                        self._delete_stale_records(target_model_class=target_model_class)
                else:
                    raise Exception(f'Failed to get corresponding model for {table_name}')
        except Exception as e:
            self._session.rollback()
            raise e
        finally:
            self._session.close()

    def _validate_record_staleness_pct(self, target_model_class: Type[RDSModel]) -> int:
        """
        Validation method.Focused on limit the risk of deleting records.
        - Check if deleted records for a table will be within 5% of the total records in the table.
        :param target_model_class:
        :return:
        """
        total_records_count = self._session.query(func.count('*')).select_from(target_model_class).scalar()
        stale_records_count = self._session.query(func.count('*')).select_from(target_model_class) \
            .filter(self._get_stale_records_filter_condition(target_model_class=target_model_class)).scalar()

        staleness_pct = 0
        target_table = target_model_class.__tablename__
        if stale_records_count:
            staleness_pct = stale_records_count * 100 / total_records_count
            threshold = self.staleness_max_pct_dict.get(target_table, self.staleness_max_pct)
            if staleness_pct >= threshold:
                raise Exception(f'Staleness percentage of {target_table} is {staleness_pct} %. '
                                f'Stopping due to over threshold {threshold} %')
            LOGGER.info(f'Will be deleting {stale_records_count} {target_table} '
                        f'record(s) or {staleness_pct}% of {target_table} data.')
        else:
            LOGGER.info(f'No stale records in {target_table}')
        return staleness_pct

    def _delete_stale_records(self, target_model_class: Type[RDSModel]) -> None:
        target_table = target_model_class.__tablename__
        try:
            deleted_records_count = self._session.query(target_model_class).filter(
                self._get_stale_records_filter_condition(target_model_class=target_model_class)).delete()
            self._session.commit()
            LOGGER.info(f'Deleted {deleted_records_count} record(s) of {target_table}')
        except Exception as e:
            LOGGER.exception(f'Failed to delete stale records for {target_table}')
            raise e

    def _get_stale_records_filter_condition(self, target_model_class: Type[RDSModel]) -> Any:
        """
        Return the appropriate stale records filter condition depending on which field is used to expire stale data.
        :param target_model_class:
        :return:
        """
        if self.ms_to_expire:
            current_time = int(time.time() * 1000)
            filter_condition = target_model_class.publisher_last_updated_epoch_ms < (
                current_time - self.marker)
        else:
            filter_condition = target_model_class.published_tag != self.marker
        return filter_condition
