"""Input and output functions and classes for FASTA files."""

from typing import TextIO, Union
from collections.abc import Generator, Iterable
from collections import namedtuple
import io
import sys
import textwrap

_FastaSequence = namedtuple('FastaSequence', 
                            'name description sequence')

class FastaSequence(_FastaSequence):
    
    """Named tuple which gives a fasta-formatted sequence when printed.

    Attributes
    ----------
    name : str
        Name of the sequence.
    description : str
        Sequence description string.
    sequence : str
        The actual sequence.

    Methods
    -------
    __str__
        Show the FASTA-formatted sequence.

    Examples
    --------
    >>> s = FastaSequence("example", "This is a description", "ATCG")
    >>> print(s)
    >example This is a description
    ATCG
    <BLANKLINE>
    
    """
    
    def __str__(self) -> str:

        """Show the FASTA-formatted sequence."""

        return ('>' + self.name + ' ' +  self.description + '\n' +
                '\n'.join(textwrap.wrap(self.sequence, width=80)) + 
                '\n')
    


def read_fasta(file: Union[str, TextIO]) -> Generator[FastaSequence]:

    """Stream sequences from a FASTA file.

    Takes a file handle and yields tuple of (name, description, sequence) 
    for each sequence.

    Parameters
    ----------
    file: TextIO
        String or file handle such as on generated by `open(f, mode='r')`.

    Yields
    ------
    FastaSequence
        Named tuple containing name, description, sequence fields.

    Examples
    --------
    >>> list(read_fasta('''
    ... >example This is a description
    ... ATCG
    ... '''))
    [FastaSequence(name='example', description='This is a description', sequence='ATCG')]

    """

    if isinstance(file, str):

        file = io.StringIO(file)

    try:
        file.seek(0)
    except io.UnsupportedOperation:
        pass

    seq_name, seq_desc, seq = None, None, ''

    for line in file:

        line = line.rstrip()

        if line.startswith('>'):

            if seq_name is not None:

                yield FastaSequence(seq_name, seq_desc, seq)

                seq_name, seq_desc, seq = None, None, ''

            seq_header = line.lstrip('>').lstrip().split()
            seq_name, seq_desc = seq_header[0], ' '.join(seq_header[1:])

        elif len(line) > 0:

            seq += line

    if len(seq) > 0:

        yield FastaSequence(seq_name, seq_desc, seq)


def write_fasta(stream: Iterable[FastaSequence],
                file: TextIO = sys.stdout) -> None:
    
    """Stream sequences to a FASTA file.

    Takes an iterable of FastaSequence and writes them to the given file.

    Parameters
    ----------
    stream : Sequence
        Iterable of FastaSequence objects.
    file : TextIO
        File handle such as on generated by `open(f, mode='w')`.

    Returns
    -------
    None

    Examples
    --------
    >>> fasta_stream = [FastaSequence("example", "This is a description", 
    ...                               "ATCG"),
    ...                 FastaSequence("example2", "This is another sequence", 
    ...                               "GGGAAAA")]
    >>> write_fasta(fasta_stream)  # doctest: +SKIP
    >example This is a description
    ATCG
    >example2 This is another sequence
    GGGAAAA
        
    """

    for fasta_seq in stream:

        file.write(str(fasta_seq))

    return None

