import gzip
from pathlib import Path
import re

from Bio import SeqIO
from collections import defaultdict
from utils.search_result_parsing.helper_classes import InputFile


def get_sequences_from_fasta(accs, fasta_fnames):
    """
    Search fasta files for accessions and return dictionary of accessions and sequences
    :param accs: iterable of accession strings
    :param fasta_fnames: filepaths for fasta files
    :return: dictionary, key: accession, value: set of sequences
    """
    if len(accs) == 0:
        return {}
    accs_chars = set()
    for acc in accs:
        accs_chars.update([re.escape(x) for x in list(acc)])
    # regex explanation:
    # group 1: (?<=^>|[\x01|]) - positive lookbehind: content must match before the remainder of remaining regex
    #   either start with '>'
    #   or one of SOH char (unicode: \x01) or pipe char ('|')
    # group 2: ([" + "".join(accs_chars) + "]+) - at least one of the chars in acc_chars (which is the set of all chars
    #   in all 'accs'
    # group 3: (?=[| \x01]|$) - positive lookahead: content must match after the preceding regex
    #   either one of the chars pipe ('|'), space (' '), or SOH ('\x01')
    #   or end of the string
    regex = re.compile("(?<=^>|[\x01|])([" + "".join(accs_chars) + "]+)(?=[| \x01]|$)")
    seqs = defaultdict(set)
    for fname in fasta_fnames:
        open_func = gzip.open if Path(fname).suffix == ".gz" else open
        file_obj = InputFile(fname)
        with open_func(fname, 'rt', encoding=file_obj.encoding) as handle:
            for record in SeqIO.parse(handle, "fasta"):
                for match in accs.intersection(set(regex.findall('>' + record.description))):
                    seqs[match].add(str(record.seq))
    return seqs
