import atexit
from datetime import date, datetime
from threading import Thread
from time import sleep
from typing import Any, Tuple

import pymysql
from pymysql.connections import Connection
from pymysql.cursors import Cursor

from src import config, logger
from src.clients.base_client import BaseClient
from src.constants import DEFAULT_DATE_FORMAT
from src.constants.queries.mysql import CREATE_INDEX, CREATE_TABLE, DELETE_DATE_RANGE, DROP_TABLE, GET_VALUE, \
    RENAME_TABLE, SET_VALUE, SHOW_CREATE_TABLE
from src.exceptions import MySQLError, MySQLNotFinishedError


class MySQLClient(BaseClient):
    def __init__(self):
        self._connection, self._cursor = self._get_connection()
        atexit.register(self.cleanup)

    @property
    def is_closed(self):
        return not self._connection or not self._connection.open

    def cleanup(self):
        self._cursor.close()
        self._cursor = None
        self._connection.close()
        self._connection = None

    @staticmethod
    def _get_connection() -> Tuple[Connection, Cursor]:
        connection = pymysql.connect(
            host=config.MySQL.HOST,
            port=config.MySQL.PORT,
            user=config.MySQL.USER,
            passwd=config.MySQL.PASSWORD,
            db=config.MySQL.DATABASE,
            autocommit=False,
        )
        cursor = connection.cursor()
        return connection, cursor

    def create_temp_table(self, new_table: str, old_table: str, id_column: str):
        self._cursor.execute(CREATE_TABLE.format(new_table=new_table, old_table=old_table, id_column=id_column))
        self._connection.commit()

    def create_temp_indexes(self, new_table: str, old_table: str):
        self._cursor.execute(SHOW_CREATE_TABLE.format(table=old_table))
        create_line_list = self._cursor.fetchone()[-1].split("\n")
        for line in create_line_list:
            line = line.strip(" ,")
            if not line.startswith("KEY"):
                continue
            self._cursor.execute(CREATE_INDEX.format(table=new_table, index=line))
        self._connection.commit()

    def drop_new_table(self, table_name):
        self._cursor.execute(DROP_TABLE.format(table=table_name))
        self._connection.commit()

    def replace_table(self, new_table: str, old_table: str):
        self._cursor.execute(DROP_TABLE.format(table=old_table))
        self._cursor.execute(RENAME_TABLE.format(old_name=new_table, new_name=old_table))
        self._connection.commit()

    def get_single_value(self, query: str) -> Any:
        self._cursor.execute(query)
        self._connection.commit()
        result = self._cursor.fetchone()
        return result[0] if result else None

    def execute_query_sync(self, query: str):
        self._cursor.execute(query)
        self._connection.commit()

    @staticmethod
    def _execute_query_thread(cursor: Cursor, query: str):
        try:
            cursor.execute(query)
        except Exception as ex:
            logger.log.warning(str(ex))

    def execute_query_async(self, query: str) -> int:
        connection, cursor = self._get_connection()
        connection.autocommit(True)
        result = connection.thread_id()
        thread = Thread(target=self._execute_query_thread, args=(cursor, query))
        thread.daemon = True
        thread.start()
        sleep(config.MySQL.SLEEP)
        cursor.close()
        connection.close()
        return result

    def cancel_query(self, thread_id: int):
        try:
            self._connection.kill(thread_id)
        except Exception as ex:
            logger.log.warning(f"Cannot cancel query: {ex}")

    def _check_process_exists(self, thread_id: int) -> bool:
        self._cursor.execute(
            "SELECT ID FROM INFORMATION_SCHEMA.PROCESSLIST WHERE COMMAND = 'Query' AND ID = %s", (thread_id,)
        )
        self._connection.commit()
        result = self._cursor.fetchone()
        return bool(result)

    def _calc_rows(self, load_prefix: str) -> int:
        self._cursor.execute("SELECT COUNT(1) FROM mysql.aurora_s3_load_history WHERE load_prefix = %s", (load_prefix,))
        self._connection.commit()
        result = self._cursor.fetchone()
        return result[0]

    def check_query_finished(self, thread_id: int, load_prefix: str, file_count: int):
        if self._check_process_exists(thread_id):
            raise MySQLNotFinishedError()
        db_count = self._calc_rows(load_prefix)
        if db_count != file_count:
            raise MySQLError(f"S3 {file_count} != DB {db_count} ({load_prefix})")

    def get_date_value(self, key: str, default_value: date or None = None) -> date or None:
        """Get value from DB Apollo key-value storage.

        Args:
            key: Key.
            default_value: Default value if None.

        Returns:
            Value.

        """
        value = self.get_single_value(GET_VALUE.format(key=key))
        if not value:
            return default_value
        return datetime.strptime(value, DEFAULT_DATE_FORMAT).date()

    def set_value(self, key: str, value: str):
        """Set value for a key to key-value storage.

        Args:
            key: Key.
            value: Value.

        """
        self.execute_query_sync(SET_VALUE.format(key=key, value=value))

    def set_date_value(self, key: str, value: date):
        """Set value to DB Apollo key-value storage.

        Args:
            key: Key.
            value: Value.

        """
        self.set_value(key, value.strftime(DEFAULT_DATE_FORMAT))

    def delete_date_range(self, table: str, start_date: date, end_date: date):
        """Delete date range.

        Args:
            table: Table name.
            start_date: Delete data from date.
            end_date: Delete date to date.

        """
        self.execute_query_sync(
            DELETE_DATE_RANGE.format(table=table, start_date=start_date.isoformat(), end_date=end_date.isoformat())
        )