# Copyright 2009-2011 by Peter Cock.  All rights reserved.
#
# This file is part of the Biopython distribution and governed by your
# choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
# Please see the LICENSE file that should have been included as part of this
# package.
"""Dictionary like indexing of sequence files (PRIVATE).

You are not expected to access this module, or any of its code, directly. This
is all handled internally by the Bio.SeqIO.index(...) and index_db(...)
functions which are the public interface for this functionality.

The basic idea is that we scan over a sequence file, looking for new record
markers. We then try to extract the string that Bio.SeqIO.parse/read would
use as the record id, ideally without actually parsing the full record. We
then use a subclassed Python dictionary to record the file offset for the
record start against the record id.

Note that this means full parsing is on demand, so any invalid or problem
record may not trigger an exception until it is accessed. This is by design.

This means our dictionary like objects have in memory ALL the keys (all the
record identifiers), which shouldn't be a problem even with second generation
sequencing. If memory is an issue, the index_db(...) interface stores the
keys and offsets in an SQLite database - which can be re-used to avoid
re-indexing the file for use another time.
"""

from __future__ import print_function

import re
from io import BytesIO
from Bio._py3k import StringIO
from Bio._py3k import _bytes_to_string

from Bio import SeqIO
from Bio import Alphabet
from Bio.File import _IndexedSeqFileProxy, _open_for_random_access


class SeqFileRandomAccess(_IndexedSeqFileProxy):
    """Base class for defining random access to sequence files."""

    def __init__(self, filename, format, alphabet):
        """Initialize the class."""
        self._handle = _open_for_random_access(filename)
        self._alphabet = alphabet
        self._format = format
        # Load the parser class/function once an avoid the dict lookup in each
        # __getitem__ call:
        i = SeqIO._FormatToIterator[format]
        # The following alphabet code is a bit nasty... duplicates logic in
        # Bio.SeqIO.parse()
        if alphabet is None:

            def _parse(handle):
                """Dynamically generated parser function (PRIVATE)."""
                return next(i(handle))

        else:
            # TODO - Detect alphabet support ONCE at __init__
            def _parse(handle):
                """Dynamically generated parser function (PRIVATE)."""
                try:
                    return next(i(handle, alphabet=alphabet))
                except TypeError:
                    return next(SeqIO._force_alphabet(i(handle), alphabet))

        self._parse = _parse

    def get(self, offset):
        """Return SeqRecord."""
        # Should be overridden for binary file formats etc:
        return self._parse(StringIO(_bytes_to_string(self.get_raw(offset))))


####################
# Special indexers #
####################
# Anything where the records cannot be read simply by parsing from
# the record start. For example, anything requiring information from
# a file header - e.g. SFF files where we would need to know the
# number of flows.
class SffRandomAccess(SeqFileRandomAccess):
    """Random access to a Standard Flowgram Format (SFF) file."""

    def __init__(self, filename, format, alphabet):
        """Initialize the class."""
        SeqFileRandomAccess.__init__(self, filename, format, alphabet)
        header_length, index_offset, index_length, number_of_reads, self._flows_per_read, self._flow_chars, self._key_sequence = SeqIO.SffIO._sff_file_header(
            self._handle
        )

    def __iter__(self):
        """Load any index block in the file, or build it the slow way (PRIVATE)."""
        if self._alphabet is None:
            self._alphabet = Alphabet.generic_dna
        handle = self._handle
        handle.seek(0)
        # Alread did this in __init__ but need handle in right place
        header_length, index_offset, index_length, number_of_reads, self._flows_per_read, self._flow_chars, self._key_sequence = SeqIO.SffIO._sff_file_header(
            handle
        )
        if index_offset and index_length:
            # There is an index provided, try this the fast way:
            count = 0
            max_offset = 0
            try:
                for name, offset in SeqIO.SffIO._sff_read_roche_index(handle):
                    max_offset = max(max_offset, offset)
                    yield name, offset, 0
                    count += 1
                if count != number_of_reads:
                    raise ValueError(
                        "Indexed %i records, expected %i" % (count, number_of_reads)
                    )
                # If that worked, call _check_eof ...
            except ValueError as err:
                import warnings
                from Bio import BiopythonParserWarning

                warnings.warn(
                    "Could not parse the SFF index: %s" % err, BiopythonParserWarning
                )
                assert count == 0, "Partially populated index"
                handle.seek(0)
                # Drop out to the slow way...
            else:
                # Fast way worked, check EOF
                if index_offset + index_length <= max_offset:
                    # Can have an index at start (or mid-file)
                    handle.seek(max_offset)
                    # Parse the final read,
                    SeqIO.SffIO._sff_read_raw_record(handle, self._flows_per_read)
                    # Should now be at the end of the file!
                SeqIO.SffIO._check_eof(handle, index_offset, index_length)
                return
        # We used to give a warning in this case, but Ion Torrent's
        # SFF files don't have an index so that would be annoying.
        # Fall back on the slow way!
        count = 0
        for name, offset in SeqIO.SffIO._sff_do_slow_index(handle):
            yield name, offset, 0
            count += 1
        if count != number_of_reads:
            raise ValueError(
                "Indexed %i records, expected %i" % (count, number_of_reads)
            )
        SeqIO.SffIO._check_eof(handle, index_offset, index_length)

    def get(self, offset):
        """Return the SeqRecord starting at the given offset."""
        handle = self._handle
        handle.seek(offset)
        return SeqIO.SffIO._sff_read_seq_record(
            handle,
            self._flows_per_read,
            self._flow_chars,
            self._key_sequence,
            self._alphabet,
        )

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        handle = self._handle
        handle.seek(offset)
        return SeqIO.SffIO._sff_read_raw_record(handle, self._flows_per_read)


class SffTrimedRandomAccess(SffRandomAccess):
    """Random access to an SFF file with defined trimming applied to each sequence."""

    def get(self, offset):
        """Return the SeqRecord starting at the given offset."""
        handle = self._handle
        handle.seek(offset)
        return SeqIO.SffIO._sff_read_seq_record(
            handle,
            self._flows_per_read,
            self._flow_chars,
            self._key_sequence,
            self._alphabet,
            trim=True,
        )


###################
# Simple indexers #
###################


class SequentialSeqFileRandomAccess(SeqFileRandomAccess):
    """Random access to a simple sequential sequence file."""

    def __init__(self, filename, format, alphabet):
        """Initialize the class."""
        SeqFileRandomAccess.__init__(self, filename, format, alphabet)
        marker = {
            "ace": b"CO ",
            "embl": b"ID ",
            "fasta": b">",
            "genbank": b"LOCUS ",
            "gb": b"LOCUS ",
            "imgt": b"ID ",
            "phd": b"BEGIN_SEQUENCE",
            "pir": b">..;",
            "qual": b">",
            "swiss": b"ID ",
            "uniprot-xml": b"<entry ",
        }[format]
        self._marker = marker
        self._marker_re = re.compile(b"^" + marker)

    def __iter__(self):
        """Return (id, offset, length) tuples."""
        marker_offset = len(self._marker)
        marker_re = self._marker_re
        handle = self._handle
        handle.seek(0)
        # Skip any header before first record
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if marker_re.match(line) or not line:
                break
        # Should now be at the start of a record, or end of the file
        while marker_re.match(line):
            # Here we can assume the record.id is the first word after the
            # marker. This is generally fine... but not for GenBank, EMBL, Swiss
            id = line[marker_offset:].strip().split(None, 1)[0]
            length = len(line)
            while True:
                end_offset = handle.tell()
                line = handle.readline()
                if marker_re.match(line) or not line:
                    yield _bytes_to_string(id), start_offset, length
                    start_offset = end_offset
                    break
                else:
                    # Track this explicitly as can't do file offset difference on BGZF
                    length += len(line)
        assert not line, repr(line)

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        # For non-trivial file formats this must be over-ridden in the subclass
        handle = self._handle
        marker_re = self._marker_re
        handle.seek(offset)
        lines = [handle.readline()]
        while True:
            line = handle.readline()
            if marker_re.match(line) or not line:
                # End of file, or start of next record => end of this record
                break
            lines.append(line)
        return b"".join(lines)


#######################################
# Fiddly indexers: GenBank, EMBL, ... #
#######################################


class GenBankRandomAccess(SequentialSeqFileRandomAccess):
    """Indexed dictionary like access to a GenBank file."""

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        marker_re = self._marker_re
        accession_marker = b"ACCESSION "
        version_marker = b"VERSION "
        # Skip and header before first record
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if marker_re.match(line) or not line:
                break
        # Should now be at the start of a record, or end of the file
        while marker_re.match(line):
            # We cannot assume the record.id is the first word after LOCUS,
            # normally the first entry on the VERSION or ACCESSION line is used.
            # However if both missing, GenBank parser falls back on LOCUS entry.
            try:
                key = line[5:].split(None, 1)[0]
            except ValueError:
                # Warning?
                # No content in LOCUS line
                key = None
            length = len(line)
            while True:
                end_offset = handle.tell()
                line = handle.readline()
                if marker_re.match(line) or not line:
                    if not key:
                        raise ValueError(
                            "Did not find usable ACCESSION/VERSION/LOCUS lines"
                        )
                    yield _bytes_to_string(key), start_offset, length
                    start_offset = end_offset
                    break
                elif line.startswith(accession_marker):
                    try:
                        key = line.rstrip().split()[1]
                    except IndexError:
                        # No content in ACCESSION line
                        pass
                elif line.startswith(version_marker):
                    try:
                        version_id = line.rstrip().split()[1]
                        if (
                            version_id.count(b".") == 1
                            and version_id.split(b".")[1].isdigit()
                        ):
                            # This should mimic the GenBank parser...
                            key = version_id
                    except IndexError:
                        # No content in VERSION line
                        pass

                length += len(line)
        assert not line, repr(line)


class EmblRandomAccess(SequentialSeqFileRandomAccess):
    """Indexed dictionary like access to an EMBL file."""

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        marker_re = self._marker_re
        semi_char = b";"
        sv_marker = b"SV "
        ac_marker = b"AC "
        # Skip any header before first record
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if marker_re.match(line) or not line:
                break
        # Should now be at the start of a record, or end of the file
        while marker_re.match(line):
            # We cannot assume the record.id is the first word after ID,
            # normally the SV line is used.
            setbysv = False  # resets sv as false
            length = len(line)
            if line[2:].count(b";") in [5, 6]:
                # Looks like the semi colon separated style introduced in 2006
                # Or style from IPD-IMGT/HLA after their v3.16.0 release
                parts = line[3:].rstrip().split(b";")
                if parts[1].strip().startswith(sv_marker):
                    # The SV bit gives the version
                    key = parts[0].strip() + b"." + parts[1].strip().split()[1]
                    setbysv = True
                else:
                    key = parts[0].strip()
            elif line[2:].count(b";") in [2, 3]:
                # Looks like the pre 2006 style, take first word only
                # Or, with two colons, the KIPO patent variantion
                key = line[3:].strip().split(None, 1)[0]
                if key.endswith(b";"):
                    key = key[:-1]
            else:
                raise ValueError("Did not recognise the ID line layout:\n%r" % line)
            while True:
                end_offset = handle.tell()
                line = handle.readline()
                if marker_re.match(line) or not line:
                    end_offset = handle.tell() - len(line)
                    yield _bytes_to_string(key), start_offset, length
                    start_offset = end_offset
                    break
                elif line.startswith(ac_marker) and not setbysv:
                    key = line.rstrip().split()[1]
                    if key.endswith(b";"):
                        key = key[:-1]
                elif line.startswith(sv_marker):
                    key = line.rstrip().split()[1]
                    setbysv = True
                length += len(line)
        assert not line, repr(line)


class SwissRandomAccess(SequentialSeqFileRandomAccess):
    """Random access to a SwissProt file."""

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        marker_re = self._marker_re
        # Skip any header before first record
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if marker_re.match(line) or not line:
                break
        # Should now be at the start of a record, or end of the file
        while marker_re.match(line):
            length = len(line)
            # We cannot assume the record.id is the first word after ID,
            # normally the following AC line is used.
            line = handle.readline()
            length += len(line)
            assert line.startswith(b"AC ")
            key = line[3:].strip().split(b";")[0].strip()
            while True:
                end_offset = handle.tell()
                line = handle.readline()
                if marker_re.match(line) or not line:
                    yield _bytes_to_string(key), start_offset, length
                    start_offset = end_offset
                    break
                length += len(line)
        assert not line, repr(line)


class UniprotRandomAccess(SequentialSeqFileRandomAccess):
    """Random access to a UniProt XML file."""

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        marker_re = self._marker_re
        start_acc_marker = b"<accession>"
        end_acc_marker = b"</accession>"
        end_entry_marker = b"</entry>"
        # Skip any header before first record
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if marker_re.match(line) or not line:
                break
        # Should now be at the start of a record, or end of the file
        while marker_re.match(line):
            length = len(line)
            # We expect the next line to be <accession>xxx</accession>
            # (possibly with leading spaces)
            # but allow it to be later on within the <entry>
            key = None
            while True:
                line = handle.readline()
                if key is None and start_acc_marker in line:
                    assert end_acc_marker in line, line
                    key = line[line.find(start_acc_marker) + 11 :].split(b"<", 1)[0]
                    length += len(line)
                elif end_entry_marker in line:
                    length += line.find(end_entry_marker) + 8
                    end_offset = (
                        handle.tell() - len(line) + line.find(end_entry_marker) + 8
                    )
                    assert start_offset + length == end_offset
                    break
                elif marker_re.match(line) or not line:
                    # Start of next record or end of file
                    raise ValueError("Didn't find end of record")
                else:
                    length += len(line)
            if not key:
                raise ValueError(
                    "Did not find <accession> line in bytes %i to %i"
                    % (start_offset, start_offset + length)
                )
            yield _bytes_to_string(key), start_offset, length
            # Find start of next record
            while not marker_re.match(line) and line:
                start_offset = handle.tell()
                line = handle.readline()
        assert not line, repr(line)

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        handle = self._handle
        marker_re = self._marker_re
        end_entry_marker = b"</entry>"
        handle.seek(offset)
        data = [handle.readline()]
        while True:
            line = handle.readline()
            i = line.find(end_entry_marker)
            if i != -1:
                data.append(line[: i + 8])
                break
            if marker_re.match(line) or not line:
                # End of file, or start of next record
                raise ValueError("Didn't find end of record")
            data.append(line)
        return b"".join(data)

    def get(self, offset):
        """Return the SeqRecord starting at the given offset."""
        # TODO - Can we handle this directly in the parser?
        # This is a hack - use get_raw for <entry>...</entry> and wrap it with
        # the apparently required XML header and footer.
        data = (
            b"""<?xml version='1.0' encoding='UTF-8'?>
        <uniprot xmlns="http://uniprot.org/uniprot"
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:schemaLocation="http://uniprot.org/uniprot
        http://www.uniprot.org/support/docs/uniprot.xsd">
        """
            + self.get_raw(offset)
            + b"</uniprot>"
        )
        return next(SeqIO.UniprotIO.UniprotIterator(BytesIO(data)))


class IntelliGeneticsRandomAccess(SeqFileRandomAccess):
    """Random access to a IntelliGenetics file."""

    def __init__(self, filename, format, alphabet):
        """Initialize the class."""
        SeqFileRandomAccess.__init__(self, filename, format, alphabet)
        self._marker_re = re.compile(b"^;")

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        # Skip any header
        offset = 0
        line = ""
        while True:
            offset += len(line)
            line = handle.readline()
            if not line:
                break  # Premature end of file, or just empty?
            if not line.startswith(b";;"):
                break
        while line:
            length = 0
            assert offset + len(line) == handle.tell()
            if not line.startswith(b";"):
                raise ValueError("Records should start with ';' and not:\n%r" % line)
            while line.startswith(b";"):
                length += len(line)
                line = handle.readline()
            key = line.rstrip()
            # Now look for the first line which starts ";"
            while line and not line.startswith(b";"):
                length += len(line)
                line = handle.readline()
            yield _bytes_to_string(key), offset, length
            offset += length
            assert offset + len(line) == handle.tell()

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        handle = self._handle
        handle.seek(offset)
        marker_re = self._marker_re
        lines = []
        line = handle.readline()
        while line.startswith(b";"):
            lines.append(line)
            line = handle.readline()
        while line and not line.startswith(b";"):
            lines.append(line)
            line = handle.readline()
        return b"".join(lines)


class TabRandomAccess(SeqFileRandomAccess):
    """Random access to a simple tabbed file."""

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        tab_char = b"\t"
        while True:
            start_offset = handle.tell()
            line = handle.readline()
            if not line:
                break  # End of file
            try:
                key = line.split(tab_char)[0]
            except ValueError as err:
                if not line.strip():
                    # Ignore blank lines
                    continue
                else:
                    raise err
            else:
                yield _bytes_to_string(key), start_offset, len(line)

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        handle = self._handle
        handle.seek(offset)
        return handle.readline()


##########################
# Now the FASTQ indexers #
##########################


class FastqRandomAccess(SeqFileRandomAccess):
    """Random access to a FASTQ file (any supported variant).

    With FASTQ the records all start with a "@" line, but so can quality lines.
    Note this will cope with line-wrapped FASTQ files.
    """

    def __iter__(self):
        """Iterate over the sequence records in the file."""
        handle = self._handle
        handle.seek(0)
        id = None
        start_offset = handle.tell()
        line = handle.readline()
        if not line:
            # Empty file!
            return
        if line[0:1] != b"@":
            raise ValueError("Problem with FASTQ @ line:\n%r" % line)
        while line:
            # assert line[0]=="@"
            # This record seems OK (so far)
            id = line[1:].rstrip().split(None, 1)[0]
            # Find the seq line(s)
            seq_len = 0
            length = len(line)
            while line:
                line = handle.readline()
                length += len(line)
                if line.startswith(b"+"):
                    break
                seq_len += len(line.strip())
            if not line:
                raise ValueError("Premature end of file in seq section")
            # assert line[0]=="+"
            # Find the qual line(s)
            qual_len = 0
            while line:
                if seq_len == qual_len:
                    if seq_len == 0:
                        # Special case, quality line should be just "\n"
                        line = handle.readline()
                        if line.strip():
                            raise ValueError(
                                "Expected blank quality line, not %r" % line
                            )
                        length += len(line)  # Need to include the blank ling
                    # Should be end of record...
                    end_offset = handle.tell()
                    line = handle.readline()
                    if line and line[0:1] != b"@":
                        raise ValueError("Problem with line %r" % line)
                    break
                else:
                    line = handle.readline()
                    qual_len += len(line.strip())
                    length += len(line)
            if seq_len != qual_len:
                raise ValueError("Problem with quality section")
            yield _bytes_to_string(id), start_offset, length
            start_offset = end_offset
        # print("EOF")

    def get_raw(self, offset):
        """Return the raw record from the file as a bytes string."""
        # TODO - Refactor this and the __init__ method to reduce code duplication?
        handle = self._handle
        handle.seek(offset)
        line = handle.readline()
        data = line
        if line[0:1] != b"@":
            raise ValueError("Problem with FASTQ @ line:\n%r" % line)
        # Find the seq line(s)
        seq_len = 0
        while line:
            line = handle.readline()
            data += line
            if line.startswith(b"+"):
                break
            seq_len += len(line.strip())
        if not line:
            raise ValueError("Premature end of file in seq section")
        assert line[0:1] == b"+"
        # Find the qual line(s)
        qual_len = 0
        while line:
            if seq_len == qual_len:
                if seq_len == 0:
                    # Special case, quality line should be just "\n"
                    line = handle.readline()
                    if line.strip():
                        raise ValueError("Expected blank quality line, not %r" % line)
                    data += line
                # Should be end of record...
                line = handle.readline()
                if line and line[0:1] != b"@":
                    raise ValueError("Problem with line %r" % line)
                break
            else:
                line = handle.readline()
                data += line
                qual_len += len(line.strip())
        if seq_len != qual_len:
            raise ValueError("Problem with quality section")
        return data


###############################################################################

_FormatToRandomAccess = {
    "ace": SequentialSeqFileRandomAccess,
    "embl": EmblRandomAccess,
    "fasta": SequentialSeqFileRandomAccess,
    "fastq": FastqRandomAccess,  # Class handles all three variants
    "fastq-sanger": FastqRandomAccess,  # alias of the above
    "fastq-solexa": FastqRandomAccess,
    "fastq-illumina": FastqRandomAccess,
    "genbank": GenBankRandomAccess,
    "gb": GenBankRandomAccess,  # alias of the above
    "ig": IntelliGeneticsRandomAccess,
    "imgt": EmblRandomAccess,
    "phd": SequentialSeqFileRandomAccess,
    "pir": SequentialSeqFileRandomAccess,
    "sff": SffRandomAccess,
    "sff-trim": SffTrimedRandomAccess,
    "swiss": SwissRandomAccess,
    "tab": TabRandomAccess,
    "qual": SequentialSeqFileRandomAccess,
    "uniprot-xml": UniprotRandomAccess,
}