from abc import ABCMeta
from multiprocessing.pool import ThreadPool
from typing import Callable, Any, List

from msgbus import (
    MBEasyContext,
    TK_SYS_NAK,
    TK_SYS_ACK,
    TK_STORECFG_GET,
    FM_PARAM,
    TK_HV_GLOBALRESTART,
)

from persistence import (
    Connection,
    Driver,
    DbdException,
)


class BaseRepository(object):
    __metaclass__ = ABCMeta

    def __init__(self, mb_context):
        # type: (MBEasyContext) -> None
        self.mb_context = mb_context

    def execute_with_connection(self, method, db_name=None, service_name=None):
        # type: (Callable[[Connection], Any], int, str) -> Any
        conn = None
        try:
            conn = Driver().open(self.mb_context,
                                 dbname=str(db_name) if db_name is not None else None,
                                 service_name=service_name)
            return method(conn)
        finally:
            if conn is not None:
                conn.close()

    def execute_with_transaction(self, method, db_name=None, service_name=None):
        # type: (Callable[[Connection], Any], int, str) -> Any
        def transaction_method(conn):
            # type: (Connection) -> Any

            if conn.get_number_of_instances() != 1 and db_name is None:
                raise DbdException("Cannot open transaction without DB instance")

            conn.transaction_start()
            try:
                conn.query("BEGIN TRANSACTION")
                ret = method(conn)
                conn.query("COMMIT")
                return ret
            except Exception as _:
                conn.query("ROLLBACK")
                raise
            finally:
                conn.transaction_end()

        return self.execute_with_connection(transaction_method, db_name=db_name, service_name=service_name)

    def execute_in_all_databases(self, method, pos_list):
        # type: (Callable[[Connection], Any], List[int]) -> Any
        ret = []
        for pos_id in pos_list:
            ret.append(self._execute_in_database(method, pos_id))

        return ret

    def execute_in_all_databases_returning_flat_list(self, method, pos_list):
        # type: (Callable[[Connection], Any], List[int]) -> Any
        inner_ret = self.execute_in_all_databases(method, pos_list)

        ret = []
        for inner_list in inner_ret:
            ret.extend(inner_list)

        return ret

    def execute_parallel_in_all_databases_returning_flat_list(self, method, pos_list):
        # type: (Callable[[Connection], Any], List[int]) -> Any
        inner_ret = self.execute_parallel_in_all_databases(method, pos_list)

        ret = []
        for inner_list in inner_ret:
            ret.extend(inner_list)

        return ret

    def execute_parallel_in_all_databases(self, method, pos_list):
        # type: (Callable[[Connection], Any], List[int]) -> Any
        ret = []
        pool = ThreadPool(len(pos_list))

        for pos_id in pos_list:
            ret.append(pool.apply_async(self._execute_in_database, args=(method, pos_id)))

        pool.close()
        pool.join()

        result = [r.get() for r in ret]
        return result

    def _execute_in_database(self, method, pos_id):
        conn = None
        try:
            conn = Driver().open(self.mb_context, dbname=str(pos_id))
            return method(conn)
        finally:
            if conn is not None:
                conn.close()
