import abc
from typing import Dict, List, Tuple
import numpy
from dejavu.base_classes.base_database import BaseDatabase
import multiprocessing
import sys
import traceback


class CommonDatabase(BaseDatabase, metaclass=abc.ABCMeta):
    # Since several methods across different databases are actually just the same
    # I've built this class with the idea to reuse that logic instead of copy pasting
    # over and over the same code.

    def __init__(self):
        super().__init__()

    def get_cursor(self):
        pass

    def before_fork(self) -> None:
        """
        Called before the database instance is given to the new process
        """
        pass

    def after_fork(self) -> None:
        """
        Called after the database instance has been given to the new process

        This will be called in the new process.
        """
        pass

    def setup(self) -> None:
        """
        Called on creation or shortly afterwards.
        """
        with self.cursor() as cur:
            cur.execute(self.CREATE_SONGS_TABLE)
            cur.execute(self.CREATE_FINGERPRINTS_TABLE)

    def empty(self) -> None:
        """
        Called when the database should be cleared of all data.
        """
        with self.cursor() as cur:
            cur.execute(self.DROP_FINGERPRINTS)
            cur.execute(self.DROP_SONGS)

        self.setup()

    def delete_unfingerprinted_songs(self) -> None:
        """
        Called to remove any song entries that do not have any fingerprints
        associated with them.
        """
        with self.cursor() as cur:
            cur.execute(self.DELETE_UNFINGERPRINTED)

    def get_num_songs(self) -> int:
        """
        Returns the song's count stored.

        :return: the amount of songs in the database.
        """
        with self.cursor(buffered=True) as cur:
            cur.execute(self.SELECT_UNIQUE_SONG_IDS)
            count = cur.fetchone()[0] if cur.rowcount != 0 else 0

        return count

    def get_num_fingerprints(self) -> int:
        """
        Returns the fingerprints' count stored.

        :return: the number of fingerprints in the database.
        """
        with self.cursor(buffered=True) as cur:
            cur.execute(self.SELECT_NUM_FINGERPRINTS)
            count = cur.fetchone()[0] if cur.rowcount != 0 else 0

        return count

    def set_song_fingerprinted(self, song_id):
        """
        Sets a specific song as having all fingerprints in the database.

        :param song_id: song identifier.
        """
        with self.cursor() as cur:
            cur.execute(self.UPDATE_SONG_FINGERPRINTED, (song_id,))

    def get_songs(self) -> List[Dict[str, str]]:
        """
        Returns all fully fingerprinted songs in the database

        :return: a dictionary with the songs info.
        """
        with self.cursor(dictionary=True) as cur:
            cur.execute(self.SELECT_SONGS)
            return list(cur)

    def get_song_by_id(self, song_id: int) -> Dict[str, str]:
        """
        Brings the song info from the database.

        :param song_id: song identifier.
        :return: a song by its identifier. Result must be a Dictionary.
        """
        with self.cursor(dictionary=True) as cur:
            cur.execute(self.SELECT_SONG, (song_id,))
            return cur.fetchone()

    def get_songs_by_id(self, song_ids: List[int]) -> List[Dict[str, str]]:
        """
        Brings the song info from the database.

        :param song_id: song identifier.
        :return: a song by its identifier. Result must be a Dictionary.
        """
        with self.cursor(dictionary=True) as cur:
            query = self.SELECT_SONGS_BY_ID % ', '.join(['%s'] * len(song_ids))
            cur.execute(query, song_ids)
            return list(cur)

    def get_song_by_name(self, song_name: str) -> Dict[str, str]:
        """
        Brings the song info from the database.

        :param song_name: unique song name.
        :return: a song by its identifier. Result must be a Dictionary.
        """
        with self.cursor(dictionary=True) as cur:
            cur.execute(self.SELECT_SONG_ID_BY_NAME, (song_name,))
            return cur.fetchone()

    def insert(self, fingerprint: str, song_id: int, offset: int):
        """
        Inserts a single fingerprint into the database.

        :param fingerprint: Part of a sha1 hash, in hexadecimal format
        :param song_id: Song identifier this fingerprint is off
        :param offset: The offset this fingerprint is from.
        """
        with self.cursor() as cur:
            cur.execute(self.INSERT_FINGERPRINT, (fingerprint, song_id, offset))

    @abc.abstractmethod
    def insert_song(self, song_name: str, file_hash: str, total_hashes: int) -> int:
        """
        Inserts a song name into the database, returns the new
        identifier of the song.

        :param song_name: The name of the song.
        :param file_hash: Hash from the fingerprinted file.
        :param total_hashes: amount of hashes to be inserted on fingerprint table.
        :return: the inserted id.
        """
        pass

    def query(self, fingerprint: str = None) -> List[Tuple]:
        """
        Returns all matching fingerprint entries associated with
        the given hash as parameter, if None is passed it returns all entries.

        :param fingerprint: part of a sha1 hash, in hexadecimal format
        :return: a list of fingerprint records stored in the db.
        """
        with self.cursor() as cur:
            if fingerprint:
                cur.execute(self.SELECT, (fingerprint,))
            else:  # select all if no key
                cur.execute(self.SELECT_ALL)
            return list(cur)

    def get_iterable_kv_pairs(self) -> List[Tuple]:
        """
        Returns all fingerprints in the database.

        :return: a list containing all fingerprints stored in the db.
        """
        return self.query(None)

    def insert_hashes(self, song_id: int, hashes: List[Tuple[str, int]], batch_size: int = 50000) -> None:
        """
        Insert a multitude of fingerprints.

        :param song_id: Song identifier the fingerprints belong to
        :param hashes: A sequence of tuples in the format (hash, offset)
            - hash: Part of a sha1 hash, in hexadecimal format
            - offset: Offset this hash was created from/at.
        :param batch_size: insert batches.
        """
        values = [(song_id, hsh, int(offset)) for hsh, offset in hashes]

        with self.cursor() as cur:
            for index in range(0, len(hashes), batch_size):
                cur.executemany(self.INSERT_FINGERPRINT, values[index: index + batch_size])

    def delete_songs_by_id(self, song_ids: List[int], batch_size: int = 1000) -> None:
        """
        Given a list of song ids it deletes all songs specified and their corresponding fingerprints.

        :param song_ids: song ids to be deleted from the database.
        :param batch_size: number of query's batches.
        """
        with self.cursor() as cur:
            for index in range(0, len(song_ids), batch_size):
                # Create our IN part of the query
                query = self.DELETE_SONGS % ', '.join(['%s'] * len(song_ids[index: index + batch_size]))

                cur.execute(query, song_ids[index: index + batch_size])

    def return_matches(self, hashes: List[Tuple[str, int]], existing_song_id: int = None) \
            -> Tuple[List[Tuple[int, int]], Dict[int, int]]:
        """
        Searches the database for pairs of (hash, offset) values.

        :param hashes: A sequence of tuples in the format (hash, offset)
            - hash: Part of a sha1 hash, in hexadecimal format
            - offset: Offset this hash was created from/at.
        :param existing_song_id: song id to exclude if already fingerprinted.
        :return: a list of (sid, offset_difference) tuples and a
        dictionary with the amount of hashes matched (not considering
        duplicated hashes) in each song.
            - song id: Song identifier
            - offset_difference: (database_offset - sampled_offset)
        """
        # Create a dictionary of hash => offset pairs for later lookups
        mapper = {}
        for hsh, offset in hashes:
            if hsh.upper() in mapper.keys():
                mapper[hsh.upper()].append(offset)
            else:
                mapper[hsh.upper()] = [offset]

        values = list(mapper.keys())
        return self.return_matches_from_db_threaded(values, mapper, existing_song_id)

    def return_matches_from_db_batched(self, values, mapper, existing_song_id: int = None, batch_size: int = 5000):
        dedup_hashes = {}
        results = []
        with self.cursor() as cur:
            for index in range(0, len(values), batch_size):
                q = self.SELECT_FINGERPRINTS_EXCLUDE_ID % ', '.join([self.IN_MATCH] * len(values[index: index + batch_size]))
                cur.execute(q, values[index: index + batch_size], existing_song_id)

                for hsh, sid, cid, offsets in cur:
                    if sid not in dedup_hashes.keys():
                        dedup_hashes[sid] = cid
                    else:
                        dedup_hashes[sid] += cid
                    song_offsets = offsets.split(',')
                    for song_sampled_offset in mapper[hsh]:
                        for offset in song_offsets:
                            if offset != '':
                                results.append((sid, numpy.int64(int(offset)) - song_sampled_offset))
            return results, dedup_hashes

    def return_matches_from_db(self, values, mapper):
        dedup_hashes = {}
        results = []
        with self.cursor() as cur:
            query = self.SELECT_FINGERPRINTS % ', '.join([self.IN_MATCH] * len(values))
            cur.execute(query, values)
            for hsh, sid, cid, offsets in cur:
                if sid not in dedup_hashes.keys():
                    dedup_hashes[sid] = cid
                else:
                    dedup_hashes[sid] += cid
                song_offsets = offsets.split(',')
                for song_sampled_offset in mapper[hsh]:
                    for offset in song_offsets:
                        if offset != '':
                            results.append((sid, numpy.int64(int(offset)) - song_sampled_offset))
            return results, dedup_hashes

    def return_matches_from_db_threaded(self, values, mapper, existing_song_id):
        chunk_size = round(len(values)/12) + 1
        chunks = list(self.chunks(values, chunk_size))
        pool = multiprocessing.Pool(12)

        worker_input = list(zip(chunks, [existing_song_id] * len(chunks)))

        dedup_hashes = {}
        results = []

        iterator = pool.imap_unordered(self.matcher, worker_input)
        while True:
            try:
                tuples = next(iterator)
            except multiprocessing.TimeoutError:
                continue
            except StopIteration:
                break
            except Exception:
                print('Failed fingerprinting {}')
                # Print traceback because we can't reraise it here
                traceback.print_exc(file=sys.stdout)
            else:
                for t in tuples:
                    hsh, sid, cid, offsets = t
                    if sid not in dedup_hashes.keys():
                        dedup_hashes[sid] = cid
                    else:
                        dedup_hashes[sid] += cid
                    song_offsets = offsets.split(',')
                    for song_sampled_offset in mapper[hsh]:
                        for offset in song_offsets:
                            if offset != '':
                                results.append((sid, numpy.int64(int(offset)) - song_sampled_offset))

        pool.close()
        pool.join()

        return results, dedup_hashes

    @staticmethod
    def chunks(lst, n):
        """Yield successive n-sized chunks from lst."""
        for i in range(0, len(lst), n):
            yield lst[i:i + n]

    def matcher(self, arguments):
        values, existing_song_id = arguments
        results = []
        cur_func = self.get_cursor()
        with cur_func() as cur:
            query = self.SELECT_FINGERPRINTS_EXCLUDE_ID % (', '.join([self.IN_MATCH] * len(values)), existing_song_id)
            cur.execute(query, values)

            for hsh, sid, cid, offsets in cur:
                results.append((hsh, sid, cid, offsets))
            return results