import gzip
import mmap
import re


def readFastaNames(fastaFile, oneSeq = False):
    
    """Read an entire fasta file, and return list of sequence name(s).

    Args:
        fastaFile (str): Name of fasta file. Accepts uncompressed and gzipped files.
        oneSeq (bool): Is there only one sequence present in this file? Defaults to False.

    Returns:
        list: List of sequence names within fasta file.
    """

    assert all([fastaFile.__class__ == str, oneSeq.__class__ == bool]), \
        '`fastaFile` must be a string and `oneSeq` must be boolean.'
    
    seqNames = []
    if fastaFile.endswith('gz'):
        file = gzip.open(fastaFile, 'r')
    else:
        file = open(fastaFile, 'r')

    file_mm = mmap.mmap(file.fileno(), 0, access = mmap.ACCESS_READ)
    
    if oneSeq:
        ind = file_mm.find(b">")
        file_mm.seek(ind)
        seqName = re.sub('(\n|> |>)', '',
                         file_mm.readline().decode()
                         ).split(' ')[0]
        seqNames += [seqName]
    else:
        while True:
            ind = file_mm.find(b">")
            try:
                file_mm.seek(ind)
            except ValueError:
                break
            else:
                seqName = re.sub('(\n|> |>)', '', 
                                 file_mm.readline().decode()
                                 ).split(' ')[0]
                seqNames += [seqName]
            
    file_mm.close()
    file.close()
    
    assert seqNames.__class__ == list, 'Function returned a non-list.'
    assert all([x.__class__ == str for x in seqNames]), \
        'Item(s) in return list are non-strings.'

    return seqNames