from abc import ABC

from src import logger
from src.clients import S3Path, clients
from src.constants import Action
from src.constants.queries.mysql import DOWNLOAD_FROM_S3
from src.utils import generate_unique_id, get_temp_table_name


class BaseProcessor(ABC):
    _data_type: str = None
    _save_temp_table: bool = False
    _snowflake_query: str = None
    _destination_table: str = None
    _destination_columns: str = None
    _id_column: str = None

    @classmethod
    @property
    def data_type(cls):
        return cls._data_type

    @property
    def snowflake_query(self) -> str:
        return self._snowflake_query

    @property
    def destination_table(self) -> str:
        return self._destination_table

    @property
    def copy_table(self) -> str:
        if self._save_temp_table:
            return get_temp_table_name(self._run_id, self.destination_table)
        return self.destination_table

    @property
    def destination_columns(self) -> str:
        return self._destination_columns

    def __init__(self, action: str, payload: dict):
        """Init processor, parse args.

        Args:
            action: Action type.
            payload: Input payload.

        """
        self._action = action
        self._payload = payload

        self.parse_base_data()
        self.parse_type_specific_data()

        self._result = {}

    def parse_base_data(self):
        """Parse base input fields, not data type specific."""

        data = self._payload.get("data", {})

        if self._action not in Action.ORDER:
            raise ValueError(f"Invalid action {self._action}")

        if self._action == Action.SNOWFLAKE_EXECUTE:
            self._run_id = generate_unique_id()
        else:
            self._run_id = data["run_id"]
            if not self._run_id:
                raise ValueError(f"Invalid run ID {self._run_id}")

        self._query_id, self._thread_id, self._file_count = (
            data.get("snowflake", {}).get("query_id"),
            data.get("mysql", {}).get("thread_id"),
            data.get("s3", {}).get("file_count"),
        )

        if Action.is_after(self._action, Action.SNOWFLAKE_EXECUTE) and not self._query_id:
            raise ValueError(f"Invalid query ID {self._query_id}")

        if Action.is_after(self._action, Action.MYSQL_EXECUTE) and not self._thread_id:
            raise ValueError(f"Invalid thread ID {self._thread_id}")

    def parse_type_specific_data(self):
        """Parse current processed data type specific input."""

        return

    def pre_snowflake_query(self):
        """Pre SF query actions."""

        return

    def run_snowflake_query(self):
        """Run snowflake query to aggregate data and put it into S3."""

        self._query_id = clients.snowflake.execute_query_async(
            query=self.snowflake_query, s3_path=S3Path.get_full_path(self._data_type, self._run_id)
        )
        self._result.update({"run_id": self._run_id, "snowflake": {"query_id": self._query_id}})

    def check_snowflake_finished(self):
        """Check if snowflake query is finished."""

        clients.snowflake.check_query_finished(self._query_id)
        logger.log.info(f"{self._run_id}: counting S3 files")
        self._file_count = clients.s3.count_files(S3Path.get_relative_path(self._data_type, self._run_id))
        self._result["file_count"] = self._file_count

    def _delete_s3_files(self):
        logger.log.info(f"{self._run_id}: removing S3 files")
        clients.s3.delete_folder_files(S3Path.get_relative_path(self._data_type, self._run_id))

    def handle_snowflake_errors(self):
        """Handle SF errors."""

        logger.log.info(f"{self._run_id}: trying to cancel query {self._query_id}")
        clients.snowflake.cancel_query(self._query_id)
        self._delete_s3_files()

    def pre_mysql_query(self):
        """Pre mysql, post snowflake actions."""

        if self._save_temp_table:
            logger.log.info(f"{self._run_id}: creating a new table")
            clients.mysql.create_temp_table(self.copy_table, self.destination_table, self._id_column)

    def get_mysql_query(self) -> str:
        """Get MySQL query.

        Returns:
            MySQL query as str.

        """
        return DOWNLOAD_FROM_S3.format(
            s3_path=S3Path.get_full_path(self._data_type, self._run_id),
            table=self.copy_table,
            columns=self.destination_columns,
        )

    def run_mysql_query(self):
        """Execute MySQL download query."""

        logger.log.info(f"{self._run_id}: trying to run MySQL download query")
        try:
            self._thread_id = clients.mysql.execute_query_async(self.get_mysql_query())
            self._result["thread_id"] = self._thread_id
        except Exception as ex:
            logger.log.warning(f"Cannot download data {ex}")
            clients.mysql.drop_new_table(self.copy_table)

    def check_mysql_finished(self):
        """Check if MySQL query is finished."""

        clients.mysql.check_query_finished(
            thread_id=self._thread_id,
            load_prefix=S3Path.get_full_path(self._data_type, self._run_id),
            file_count=self._file_count,
        )

    def handle_mysql_errors(self):
        """Handle MySQL errors."""
        logger.log.info(f"{self._run_id}: trying to kill thread {self._thread_id}")
        clients.mysql.cancel_query(self._thread_id)
        if self._save_temp_table:
            logger.log.info(f"{self._run_id}: dropping the new table")
            clients.mysql.drop_new_table(self.copy_table)
        self._delete_s3_files()

    def post_mysql_query(self):
        """MySQL unload post-actions (kill thread, create indexes, replace the table, remove S3 files)."""

        logger.log.info(f"{self._run_id}: trying to kill thread {self._thread_id}")
        clients.mysql.cancel_query(self._thread_id)
        if self._save_temp_table:
            logger.log.info(f"{self._run_id}: creating indexes")
            clients.mysql.create_temp_indexes(self.copy_table, self.destination_table)
            logger.log.info(f"{self._run_id}: replacing the table")
            clients.mysql.replace_table(self.copy_table, self.destination_table)
        self._delete_s3_files()

    def run(self):
        """Execute action."""

        logger.log.info(f"{self._run_id}: {self._action}")

        if self._action == Action.SNOWFLAKE_EXECUTE:
            self.pre_snowflake_query()
            self.run_snowflake_query()
        elif self._action == Action.SNOWFLAKE_CHECK:
            self.check_snowflake_finished()
        elif self._action == Action.SNOWFLAKE_ERROR:
            self.handle_snowflake_errors()
        elif self._action == Action.MYSQL_EXECUTE:
            self.pre_mysql_query()
            self.run_mysql_query()
        elif self._action == Action.MYSQL_CHECK:
            self.check_mysql_finished()
        elif self._action == Action.MYSQL_ERROR:
            self.handle_mysql_errors()
        elif self._action == Action.MYSQL_FINALIZE:
            self.post_mysql_query()

        logger.log.info(f"{self._run_id}: {self._action} result {self._result}")

        return self._result