# ----------------------------------------------------------------------------
# Copyright (c) 2019-2026, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

from typing import Union

import qiime2
from q2_types.feature_data import (
    AlignedDNAFASTAFormat, DNAFASTAFormat,
    DNAIterator, AlignedDNAIterator)
from skbio import DNA


def _trim_sequence(sequence: DNA,
                   position_start: int,
                   position_end: int) -> DNA:
    return DNA(sequence[position_start:position_end])


def _trim_all_sequences(aligned_sequences: AlignedDNAFASTAFormat,
                        trim_positions: dict) -> AlignedDNAFASTAFormat:
    """
    Trim all sequences within given alignment based on provided positions.

    Arguments:
        aligned_sequences (AlignedDNAFASTAFormat): original, aligned sequences
        trim_positions (dict): dictionary containing positions for trimming

    Returns:
        result (AlignedDNAFASTAFormat): trimmed aligned sequences
    """

    result = AlignedDNAFASTAFormat()
    with result.open() as out_fasta:
        for seq in aligned_sequences.view(AlignedDNAIterator):
            seq_trimmed = _trim_sequence(
                seq,
                trim_positions["start"],
                trim_positions["end"])
            seq_trimmed.write(out_fasta)
    return result


def _find_terminal_positions(primer_positions: dict,
                             keep_primer_location=False) -> (int, int):
    """
    Identify left- (5') and rightmost (3') trimming position. If both primers
    were used return: (start index of the forward primer, end index of
    the reverse primer). If only forward primer was found return
    (its start index, None). If only reverse primer was found return
    (None, its end index).

    Arguments:
        primer_positions (dict): dictionary containing start and end positions
                                    of forward and reverse primers

    Returns:
        pos_start (int): leftmost (5') position index for trimming
        pos_end (int): rightmost (3') position index for trimming
    """

    pos_start, pos_end = None, None
    # when we found both primers
    if len(primer_positions) == 2:
        if primer_positions["reverse"]["start"] < \
                primer_positions["forward"]["end"]:
            raise ValueError("Reverse primer overlaps or aligned upstream the "
                             "forward primer. Are the primers correct?")
        if keep_primer_location:
            pos_start = min([x["start"] for x in primer_positions.values()])
            pos_end = max([x["end"] for x in primer_positions.values()])
        else:
            pos_start = primer_positions["forward"]["end"]
            pos_end = primer_positions["reverse"]["start"]
    # when only fwd primer was used
    elif "forward" in primer_positions.keys():
        if keep_primer_location:
            pos_start = primer_positions["forward"]["start"]
        else:
            pos_start = primer_positions["forward"]["end"]
    # when only rev primer was used
    elif "reverse" in primer_positions.keys():
        if keep_primer_location:
            pos_end = primer_positions["reverse"]["end"]
        else:
            pos_end = primer_positions["reverse"]["start"]
    return pos_start, pos_end


def _locate_primer_positions(
        alignment_with_primers: AlignedDNAFASTAFormat,
        keep_primer_location) -> dict:
    """
    Identify position of each primer within the alignment.

    Arguments:
        alignment_with_primers (AlignedDNAFASTAFormat): sequence alignment
                        containing at least one aligned primer

    Returns:
        (dict): dictionary containing trimming positions
                using 0-based indexing
    """

    primers_aligned = dict()
    for aln_seq in alignment_with_primers.view(DNAIterator):
        if aln_seq.metadata["id"] in ["forward", "reverse"]:
            primers_aligned[aln_seq.metadata["id"]] = (str(aln_seq))

    primer_positions = dict()
    for primer_id, primer_seq in primers_aligned.items():
        primer_positions[primer_id] = {
            'start': next(
                (i for i, nt in enumerate(primer_seq) if nt != "-")),
            'end': len(primer_seq) - next(
                (i for i, nt in enumerate(primer_seq[::-1]) if nt != "-"))
        }

    pos_start, pos_end = _find_terminal_positions(primer_positions,
                                                  keep_primer_location)

    # not doing any validation like in _prepare_positions since none of
    # the conditions checked there are possible to run into with the way
    # _locate_primer_positions and _find_terminal_positions are implemented

    return {"start": pos_start, "end": pos_end}


def _prepare_positions(position_start: Union[int, None],
                       position_end: Union[int, None],
                       aln_len: int) -> dict:
    """
    Prepare a dictionary with start and end position for sequence trimming.

    Arguments:
        position_start (int, None): leftmost (5') position; 1-based indexing
        position_end (int, None): rightmost (3') position; 1-based indexing
        aln_len (int): length of the alignment

    Returns:
        (dict): dictionary containing trimming positions
                using 0-based indexing
    """

    if not position_start and not position_end:
        raise ValueError("Neither primers nor sequence positions were "
                         "provided - nothing to be trimmed.")
    if position_start:
        if position_end and position_start >= position_end:
            raise ValueError("Start position should be smaller than end "
                             "position while trimming.")
        elif position_start >= aln_len:
            raise ValueError("Start position should be smaller than "
                             "alignment length.")
        else:
            position_start -= 1
    if position_end and position_end > aln_len:
        position_end = aln_len

    return {'start': position_start, 'end': position_end}


def _process_primers(
        primer_fwd: Union[str, None],
        primer_rev: Union[str, None]) -> DNAFASTAFormat:
    """
    Convert provided primers into skbio DNA format. Will reverse complement
    the reverse primer, if provided.

    Arguments:
        primer_fwd (str, None): forward primer
        primer_rev (str, None): reverse primer

    Returns:
        primers_fasta (DNAFASTAFormat): primers in FASTA format
    """

    primers = {
        'forward': DNA(
            primer_fwd, metadata={
                'id': 'forward'})
        if primer_fwd else None,
        'reverse': DNA(
            primer_rev, metadata={
                'id': 'reverse'}).reverse_complement()
        if primer_rev else None
    }

    # save primers in that format to pass them to mafft_add
    primers_fasta = DNAFASTAFormat()
    with primers_fasta.open() as out:
        [primer.write(out) for primer in primers.values() if primer]

    return primers_fasta


def _trim_alignment(expand_alignment_action,
                    aligned_sequences,
                    primer_fwd=None,
                    primer_rev=None,
                    position_start=None,
                    position_end=None,
                    keep_primer_location=False,
                    n_threads=1,
                    keeplength=True) -> AlignedDNAFASTAFormat:
    """
    Trim alignment based on primer alignment or explicitly specified
    positions. When at least one primer sequence is given, primer-based
    trimming will be performed via the `addfragments` option of mafft,
    otherwise position-based trimming is done.

    Arguments:
        expand_alignment_action: qiime action for multiple seq. alignment
        aligned_sequences (AlignedDNAFASTAFormat): alignment to be trimmed
        primer_fwd (str, None): forward primer used to find the leftmost (5')
                            trimming position
        primer_rev (str, None): reverse primer used to find the rightmost (3')
                            trimming position
        position_start (int, None): leftmost (5') trimming position
        position_end (int, None): rightmost (3') trimming position

    Returns:
        result (AlignedFASTAFormat): trimmed alignment

    """

    # if primers are provided use them
    # otherwise fall back to position-based slicing
    if primer_fwd or primer_rev:
        primers_fasta = _process_primers(primer_fwd, primer_rev)
        primers = qiime2.Artifact.import_data(
            'FeatureData[Sequence]', primers_fasta)

        # expand the existing alignment by addition of primers
        # Use 'keeplength' to make sure original alignment length
        # does not change. Otherwise there is potential for the incorrect
        # region to be extracted from the alignment. This is important
        # to make sure alignment positions can easily be referenced
        # to the original alignment. See the 'q2-alignment' plugin,
        # specifically the 'mafft-add' test cases for examples.
        alignment_with_primers, = expand_alignment_action(
            alignment=aligned_sequences,
            sequences=primers,
            addfragments=True,
            n_threads=n_threads,
            keeplength=keeplength)

        # find trim positions based on primer positions within alignment
        trim_positions = _locate_primer_positions(alignment_with_primers,
                                                  keep_primer_location)
    else:
        # find length of the alignment
        seq_iter = aligned_sequences.view(AlignedDNAIterator).generator
        aln_len = len(next(seq_iter))

        trim_positions = _prepare_positions(
            position_start, position_end, aln_len)

    result = _trim_all_sequences(aligned_sequences, trim_positions)

    return result


def trim_alignment(ctx,
                   aligned_sequences,
                   primer_fwd=None,
                   primer_rev=None,
                   position_start=None,
                   position_end=None,
                   keep_primer_location=False,
                   n_threads=1):
    """
    Trim an existing alignment based on provided primers or specific,
    pre-defined positions. Primers take precedence over the positions,
    i.e. if both are provided, positions will be ignored.

    This action uses mafft_add action to add provided primers (if any)
    to the existing alignment. Subsequently, start (5'-most) and end
    (3'-most) position from fwd and rev primer located within the new
    alignment is identified and used for slicing the original alignment.
    """
    expand_alignment_action = ctx.get_action('alignment', 'mafft_add')

    result = _trim_alignment(
        expand_alignment_action,
        aligned_sequences,
        primer_fwd,
        primer_rev,
        position_start,
        position_end,
        keep_primer_location=keep_primer_location,
        n_threads=n_threads)

    return qiime2.Artifact.import_data('FeatureData[AlignedSequence]', result)
