import sqlite3
import sys
from collections import defaultdict
from os.path import isfile, getsize
from pathlib import Path
import yaml


class ProteinGroupSqlDbInteractor:
    def __init__(self, path_to_pg_db):
        """
        database schema of protein_groups_db.sql:
        tables: ProteinGroupsQuant, ProteinGroupsProteins, ProteinGroupsInformation, Samples, ProteinAccSpectra
        columns ProteinGroupsQuant: "ID" (INT NOT NULL) | "quant" (FLOAT NOT NULL) | "sample_id" (NOT NULL), KEY: ID + sample_ID
        columns ProteinGroupsProteins: "ID" (INT PRIMARY KEY NOT NULL) | "proteins" (varchar NOT NULL) -> string of concat protein accs
        columns ProteinGroupsInformation: "ID" (varchar(20) PRIMARY KEY NOT NULL), "information" (varchar NOT NULL)
                                        "ID" = 'protein_sep', 'quant_type', 'with_spectra'
        columns Samples: "sample_id" (INT PRIMARY KEY NOT NULL)  | "sample_name"  (varchar NOT NULL) | "sample_group" (varchar not filled yet)
        columns ProteinAccSpectra: pgid | acc | spectra | sample_id) (filled if spectra available, else empty)
        :param path_to_pg_db: path to sql database generated by parsing
        """
        if 'protein_groups.yaml' in str(path_to_pg_db):
            # automatic conversion from protein_group.yaml to protein_groups_db.sql, saved in same directory
            self.path_to_pg_db = Path(path_to_pg_db).parent/ 'protein_groups_db.sql'
            self.convert_pg_yaml_to_pg_sql_db(path_to_pg_db, self.path_to_pg_db)
        else:
            self.path_to_pg_db = Path(path_to_pg_db)

    @classmethod
    def create_protein_group_database_from_data(cls, path_to_db, protein_group_data, protein_info_data):
        # create and fill SQL tables, return initialized ProteinGroupSqlDbInteractor object
        # used by search_result_parsing/api_class.py or by pg_yaml converting
        protein_sep = protein_info_data['protein_sep']
        db_obj = cls(path_to_db)
        db_obj._create_protein_group_quant_table()
        db_obj._create_protein_group_proteins_table()
        db_obj._create_protein_accs_spectra_table()
        db_obj._fill_pg_tables(protein_group_data, protein_info_data['samples'], protein_sep, protein_info_data['with_spectra'])

        db_obj._create_protein_information_table()
        db_obj._fill_pg_info_table(protein_info_data)

        db_obj._create_samples_table()
        db_obj._fill_samples_table(protein_info_data['samples'])
        return db_obj

    def isSQLite3(self, path_to_file):
        if not isfile(path_to_file):
            return False
        if getsize(path_to_file) < 100: # SQLite database file header is 100 bytes
            return False
        with open(path_to_file, 'rb') as fd:
            header = fd.read(100)
        return header[:16] == b'SQLite format 3\x00'

    def create_connection(self):
        conn = None
        try:
            conn = sqlite3.connect(self.path_to_pg_db)
        except Exception as e:
            print(e, file=sys.stderr)
            print(f'No connection possible to protein group sql database {self.path_to_pg_db}. Prophane quits.')
            exit(1)
        return conn

    def _create_protein_group_quant_table(self):
        with self.create_connection() as conn:
            conn.execute('''CREATE TABLE ProteinGroupsQuant
             (ID INT NOT NULL,
             quant FLOAT NOT NULL,
             sample_id INT NOT NULL,
             PRIMARY KEY (ID, sample_id));''')

    def _create_protein_accs_spectra_table(self):
        with self.create_connection() as conn:
            conn.execute('''CREATE TABLE ProteinAccSpectra
             (pgid INT NOT NULL,
             acc varchar,
             spectra varchar, 
             sample_id INT NOT NULL,
             PRIMARY KEY (pgid, acc, sample_id));''')


    def _create_protein_group_proteins_table(self):
        with self.create_connection() as conn:
            conn.execute('''CREATE TABLE ProteinGroupsProteins
             (ID INT PRIMARY KEY NOT NULL,
             proteins varchar NOT NULL);''')

    def _create_samples_table(self):
        with self.create_connection() as conn:
            conn.execute('''CREATE TABLE SAMPLES
             (sample_id INT PRIMARY KEY NOT NULL,
             sample_name varchar NOT NULL,
             sample_group varchar);''')


    def _create_protein_information_table(self):
        with self.create_connection() as conn:
            conn.execute('''CREATE TABLE ProteinGroupsInformation
                     (ID varchar(20) PRIMARY KEY NOT NULL,
                     information varchar NOT NULL);''')

    def _insert_into_protein_group_table(self, conn, cursor, with_spectra, records_to_insert_into_proteins_table,
                                        records_to_insert_into_quant_table, records_to_insert_into_spectra_table):
        sqlite_insert_query = """INSERT INTO ProteinGroupsProteins
                              (ID,proteins) 
                              VALUES (?, ?);"""
        cursor.executemany(sqlite_insert_query, records_to_insert_into_proteins_table)

        sqlite_insert_query = """INSERT INTO ProteinGroupsQuant
                          (ID,quant,sample_id) 
                          VALUES (?, ?, ?);"""
        cursor.executemany(sqlite_insert_query, records_to_insert_into_quant_table)
        if with_spectra:
            sqlite_insert_query = """INSERT INTO ProteinAccSpectra
                              (pgid,acc,spectra,sample_id) 
                              VALUES (?, ?, ?, ?);"""
            cursor.executemany(sqlite_insert_query, records_to_insert_into_spectra_table)
        conn.commit()

    def _fill_pg_tables(self, list_of_pg_dict, sample_list, protein_sep, with_spectra):
        records_to_insert_into_quant_table = []
        records_to_insert_into_proteins_table = []
        records_to_insert_into_spectra_table = []
        with self.create_connection() as conn:
            cursor = conn.cursor()
            for pgid, pg in enumerate(list_of_pg_dict):
                accs_str = pg['proteins']
                if with_spectra:
                    accs = accs_str.split(protein_sep)
                records_to_insert_into_proteins_table.append((pgid, accs_str))
                for sample_id in range(len(sample_list)):
                    quant = pg['quant'][sample_id]
                    records_to_insert_into_quant_table.append((pgid, quant, sample_id))
                    if with_spectra:
                        spectra_list = self.get_spectra_as_list(pg['spectra'][sample_id], len(accs), protein_sep)
                        for i, acc in enumerate(accs):
                            records_to_insert_into_spectra_table.append((pgid, acc, spectra_list[i], sample_id))
                    if len(records_to_insert_into_proteins_table) > 499:
                        self._insert_into_protein_group_table(conn, cursor, with_spectra,
                                                              records_to_insert_into_proteins_table,
                                                              records_to_insert_into_quant_table,
                                                              records_to_insert_into_spectra_table)
                        records_to_insert_into_proteins_table,  records_to_insert_into_quant_table, \
                        records_to_insert_into_spectra_table = [], [], []
            self._insert_into_protein_group_table(conn, cursor, with_spectra, records_to_insert_into_proteins_table,
                                              records_to_insert_into_quant_table, records_to_insert_into_spectra_table)

    def _fill_samples_table(self, sample_list):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            for id, sample_name in enumerate(sample_list):
                cursor.execute(f"INSERT INTO Samples (sample_id, sample_name) \
                    VALUES ('{id}', '{sample_name}');")
            conn.commit()

    def _fill_pg_info_table(self, protein_info_data):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"INSERT INTO ProteinGroupsInformation (ID,information) \
                VALUES ('protein_sep', '{protein_info_data['protein_sep']}');")
            cursor.execute(f"INSERT INTO ProteinGroupsInformation (ID,information) \
                    VALUES ('quant_type', '{protein_info_data['quant_type']}');")
            cursor.execute(f"INSERT INTO ProteinGroupsInformation (ID,information) \
                    VALUES ('with_spectra', '{protein_info_data['with_spectra']}');")
            cursor.execute("INSERT INTO ProteinGroupsInformation (ID,information)"
                           " VALUES ('db_schema_version', 4);")
            conn.commit()

    def get_spectra_as_list(self, spectra_pg, accs_nb, protein_sep):
        spectra_list = spectra_pg.split(f']{protein_sep}[')
        if spectra_list == ['[[]]'] or spectra_list == [''] or spectra_list == ['-']:
            spectra_list = [None] * accs_nb
        else:
            for i, spectra_str in enumerate(spectra_list):
                while spectra_str.startswith('['):
                    spectra_str =spectra_str[1:]
                while spectra_str.endswith(']'):
                    spectra_str =spectra_str[:-1]
                spectra_list[i] = spectra_str
            # None for no spectra
            spectra_list = [None if spectra == '' else spectra for spectra in spectra_list]
            # remove " and ' from spectra string
            spectra_list = [spectra.replace("'", "").replace('"', "") if spectra is not None else None for spectra in spectra_list ]
        return spectra_list

    def get_pg_information(self):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"SELECT * from ProteinGroupsInformation;")
            rows = cursor.fetchall()
            inf_dict = {row[0]: row[1] for row in rows}
        return inf_dict

    def get_pg_id_to_accs_dict(self, pg_list=None):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            selected_ids_part = f" WHERE ID IN ({','.join([str(pg) for pg in pg_list])});" if pg_list else ";"
            cursor.execute(f"SELECT ID, proteins from ProteinGroupsProteins " + selected_ids_part)
            rows = cursor.fetchall()
            pg_dict = {row[0]: row[1] for row in rows}
        return pg_dict

    def convert_pg_yaml_to_pg_sql_db(self, path_to_pg_yaml, path_to_sql_db):
        """
        converts the previous protein group information file protein_group.yaml into sql database containing the same
        information in path_to_sql_db.
        """
        if not self.isSQLite3(path_to_pg_yaml) and not self.isSQLite3(path_to_sql_db):
            with open(path_to_pg_yaml, "r") as handle:
                pg_dict_complete = yaml.safe_load(handle)
            samples = pg_dict_complete['samples']
            samples.sort()
            with_spectra = True if 'spectra' in pg_dict_complete['protein_groups'][0].keys() else False
            pg_dict_info = {'samples': samples, 'protein_sep': pg_dict_complete['style']['protein_sep'],
                            'with_spectra': with_spectra, 'quant_type': pg_dict_complete['style']['quant_type']}
            self.create_protein_group_database_from_data(path_to_sql_db, pg_dict_complete['protein_groups'], pg_dict_info)

    def get_number_of_pgs(self):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"SELECT COUNT (DISTINCT ID) from ProteinGroupsProteins;")
            rows = cursor.fetchall()
        return rows[0][0]

    def get_spectra_per_acc(self, accs, pgid):
        acc_to_sample_to_spectra_dict = defaultdict(list)
        accs_str = ','.join([f"'{acc}'" for acc in accs])
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"SELECT acc,sample_id,spectra from ProteinAccSpectra WHERE acc IN ({accs_str}) AND pgid={pgid};")
            rows = cursor.fetchall()
            for row in rows:
                spectra = '-' if row[2] is None else row[2]
                acc_to_sample_to_spectra_dict[row[0]].append(';'.join(spectra.split(self.get_protein_sep())))
        return acc_to_sample_to_spectra_dict

    def get_quant_values_of_pgs(self, pgid_list=None):
        pgid_to_quant_dict = defaultdict(list)
        select_str = f"SELECT ID, quant, sample_id from ProteinGroupsQuant "
        where_part = f"WHERE ProteinGroupsQuant.id IN ({','.join([str(pg) for pg in pgid_list])});" if pgid_list else ";"
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(select_str+where_part)
            rows = cursor.fetchall()
            # fetch results are sorted by sample_id -> values in same order like in sample table
            for row in rows:
                pgid_to_quant_dict[row[0]].append(row[1])
        return pgid_to_quant_dict

    def get_spectra_header(self):
        """
        param pg_db: protein group sql db
        return spectra_header:[string list], spectra_IDs_per_accession_dict[acc]=[spectra]
        """
        spectra_header = []
        for sample in self.get_samples():
            spectra_header.extend([sample + ":: spectra_IDs"])
        return spectra_header

    def get_set_all_spectra_per_sample_for_pg(self, pgid):
        # columns ProteinAccSpectra: pgid | acc | spectra | sample_id) (filled if spectra available)
        sample_to_spectra_set = defaultdict(set)
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"SELECT sample_name, spectra from ProteinAccSpectra INNER Join Samples ON "
                           f"ProteinAccSpectra.sample_id = Samples.sample_id WHERE ProteinAccSpectra.pgid={pgid};")
            rows = cursor.fetchall()
            for row in rows:
                if row[1] is not None:
                    sample_to_spectra_set[row[0]].update(row[1].split(self.get_protein_sep()))
                elif row[0] not in sample_to_spectra_set.keys():
                    sample_to_spectra_set[row[0]] = set()
        return sample_to_spectra_set

    def get_acc_to_spectra_dict(self, pgids=None):
        """
        for lca determination
        :return acc_to_spectra_dict: protein_acc to list of spectra IDs dict of all samples
        """
        acc_to_spectra_dict = defaultdict(set)
        with self.create_connection() as conn:
            cursor = conn.cursor()
            where_part = f" WHERE pgid IN ({','.join([str(pg) for pg in pgids])});" if pgids else ";"
            cursor.execute(f"SELECT acc, spectra from ProteinAccSpectra" + where_part)
            rows = cursor.fetchall()
            for row in rows:
                if row[1] is not None:
                    acc_to_spectra_dict[row[0]].update(row[1].split(self.get_protein_sep()))
                elif row[1] is None and row[0] not in acc_to_spectra_dict.keys():
                    acc_to_spectra_dict[row[0]] = set()
        return acc_to_spectra_dict

    def get_acc_list_per_pg(self, pgids=None):
        """
        :return pg_to_accs_dict, {int pgid: [list of acc in same order as accs string]}
        """
        sep = self.get_pg_information()['protein_sep']
        pg_to_accs_str_dict = self.get_pg_id_to_accs_dict(pgids)
        pg_to_accs_dict= {pg: accs_str.split(sep) for pg, accs_str in pg_to_accs_str_dict.items()}
        return pg_to_accs_dict

    def get_all_accs_from_pg_db(self):
        """get set of all protein accs """
        pg_to_accs_dict= self.get_acc_list_per_pg()
        all_accs = set()
        for accs in pg_to_accs_dict.values():
            all_accs.update(set(accs))
        return all_accs

    def get_groups_to_align_from_pg_db(self):
        """
        return list of protein_group IDs with more than one accs
        """
        sep = self.get_pg_information()['protein_sep']
        pg_to_accs_dict= self.get_pg_id_to_accs_dict()
        group_ids = [pgid for pgid, acc_str in pg_to_accs_dict.items() if sep in acc_str]
        return group_ids

    def get_protein_sep(self):
        sep = self.get_pg_information()['protein_sep']
        return sep

    def is_with_spectra(self):
        """
        determines if spectra information are present, improves LCA determination
        """
        with_spectra = True if self.get_pg_information()['with_spectra']=='True' else False
        return with_spectra

    def get_samples(self):
        with self.create_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(f"SELECT DISTINCT (sample_name) from Samples;")
            rows = cursor.fetchall()
            samples = [row[0] for row in rows]
        samples.sort()
        return samples


def create_protein_group_database(path_to_db, protein_group_data, protein_info_data):
    return ProteinGroupSqlDbInteractor.create_protein_group_database_from_data(path_to_db, protein_group_data,
                                                                               protein_info_data)