from typing import Optional

import pysam
import shortuuid
from tqdm import tqdm

from . import utils


class ReadError(Exception):
    pass


class FastqError(Exception):
    pass


class Quality:
    """Represents a Phred33 quality string.

    Attributes:
        _string: Raw quality string; for internal use only. Use ``string`` instead.
    """

    def __init__(self, qualities: str):
        """
        Args:
            qualities: Raw quality string
        """
        self._string = qualities

    @property
    def string(self):
        """Raw quality string"""
        return self._string

    @property
    def values(self):
        """List of quality values"""
        return list(pysam.qualitystring_to_array(self._string))

    @property
    def probs(self):
        """The quality values converted to probabilities of error"""
        return [10**(-q / 10) for q in self.values]


class Read:
    """Class that represents a FASTQ read. Once the class is initialized, with the
    header, sequence, and quality strings, they should not be changed (hence using @property).
    Only Phred33 quality scores are supported, but there is no check for this.

    Attributes:
        _header: Raw header string, including the ``@``; for internal use only.
            Use ``header`` instead.
        _sequence: Raw sequence string; for internal use only. Use ``sequence`` instead.
        _qualities: :class:`Quality` instance representing the sequence qualities;
            for internal use only. Use ``qualities`` instead.

    Raises:
        ReadError: if the ``header`` does not start with ``@``
    """

    def __init__(self, header: str, sequence: str, qualities: str):
        """
        Args:
            header: Raw header string, including the ``@``
            sequence: Raw sequence string
            qualities: Raw qualities string
        """
        if not header.startswith('@'):
            raise ReadError(f'FASTQ header `{header}` does not start with `@`')

        self._header = header.strip()
        self._sequence = sequence.strip()
        self._qualities = Quality(qualities.strip())

    @property
    def header(self):
        """Raw header string"""
        return self._header

    @property
    def sequence(self):
        """Raw sequence string"""
        return self._sequence

    @property
    def qualities(self):
        """:class:`Quality` instance representing the sequence qualities"""
        return self._qualities

    @property
    def name(self):
        """Name of the sequence, which comes immediately after the ``@``"""
        return self.header[1:self.header.index(' ')
                           ] if ' ' in self.header else self.header[1:]

    @property
    def attributes(self):
        """String of read attributes. This is the substring after the first space in the header."""
        return self.header[self.header.index(' ') +
                           1:] if ' ' in self.header else ''


class Fastq(utils.FileWrapper):
    """Class that represents a FASTQ file.
    """

    def read(self) -> Read:
        """Read a single FASTQ entry.

        Returns:
            The next read as a :class:`Read` instance

        Raises:
            FastqError: If the file was not opened for reading, or the file was closed.
            StopIteration: When there are no more entries to read.
        """
        if self.mode != 'r':
            raise FastqError(f'Can not read from file in mode `{self.mode}`')
        if self.closed:
            raise FastqError('Can not read from closed file')

        header, sequence, _, qualities = [
            next(self.fp).strip() for _ in range(4)
        ]
        return Read(header, sequence, qualities)

    def write(self, entry: Read):
        """Write a single FASTQ entry.

        Args:
            entry: :class:`Read` instance to write to the FASTQ

        Raises:
            FastaError: If the file was not opened for writing, or the file was closed.
        """
        if self.mode != 'w':
            raise FastqError(f'Can not write to file in mode `{self.mode}`')
        if self.closed:
            raise FastqError('Can not write to closed file')

        self.fp.write(f'{entry.header}\n')
        self.fp.write(f'{entry.sequence}\n')
        self.fp.write('+\n')
        self.fp.write(f'{entry.qualities.string}\n')


def fastq_to_bam(
    fastq_path: str,
    bam_path: str,
    name: Optional[str] = None,
    n_threads: int = 1
):
    """Convert a Fastq to unmapped BAM.

    Args:
        fastq_path: Path to the input FASTQ
        bam_path: Path to the output BAM
        name: Name for this set of reads. Defaults to None. If not provided,
            a random string is generated by calling :func:`shortuuid.uuid`. This
            value is added as the read group (RG tag) for all the reads in the BAM.
        n_threads: Number of threads to use. Defaults to 1.
    """
    rg = name or shortuuid.uuid()
    header = pysam.AlignmentHeader.from_dict({
        'HD': {
            'VN': pysam.version.__samtools_version__,
            'SO': 'unsorted'
        },
        'RG': [{
            'ID': rg
        }],
    })
    with Fastq(fastq_path,
               'r') as f_in, pysam.AlignmentFile(bam_path, 'wb', header=header,
                                                 threads=n_threads) as f_out:
        for read in tqdm(f_in, smoothing=0, desc='Writing BAM'):
            al = pysam.AlignedSegment(header)
            al.query_name = read.name
            al.query_sequence = read.sequence
            al.query_qualities = read.qualities.values
            al.flag = 4  # unmapped
            al.tags = [('RG', rg)]
            f_out.write(al)
