# SPDX-FileCopyrightText: 2025 2024 Samir Bertache
#
# SPDX-License-Identifier: AGPL-3.0-or-later

"""
This script is a the Parasplit project, designed to process paired-end FASTQ files by fragmenting DNA sequences at specified restriction enzyme sites.

Copyright © 2024 Samir Bertache

SPDX-License-Identifier: AGPL-3.0-or-later

===============================================================================

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU Affero General Public License as published by the
Free Software Foundation, either version 3 of the License, or (at your option)
any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""

import logging
import re
from typing import Generator, List, Tuple, Union

logging.basicConfig(level=logging.INFO)


#################### Specific part of Borderless option #####################


def find_positions_for_one_site_borderless(
    text: str, Enzyme: Tuple[re.Pattern, int]
) -> List[List[int]]:
    """
    Find all positions of a specific pattern (RESite) in a given text using regular expressions.

    Parameters:
        text (str): The text to search in.
        Enzyme (Tuple[re.Pattern, int]): A compiled regex and an offset.

    Returns:
        List of [position, offset].

    Examples:
        >>> find_positions_for_one_site_borderless("", (re.compile("A"), 1))
        []
        >>> find_positions_for_one_site_borderless("AAAA", (re.compile("AA"), 2))
        [[0, 2], [2, 2]]
        >>> find_positions_for_one_site_borderless("XYZ", (re.compile("A"), 1))
        []
        >>> find_positions_for_one_site_borderless("GATCGATC", (re.compile("GATC"), 4))
        [[0, 4], [4, 4]]
    """
    regex, offset = Enzyme
    return [[m.start(), offset] for m in regex.finditer(text)]


def find_all_pos_borderless(
    text: str,
    ligation_site_list: List[Tuple[re.Pattern, int]],
) -> List[List[int]]:
    """
    Aggregate all borderless site positions plus the end-of-text marker.

    Examples:
        >>> find_all_pos_borderless("", [])
        [[0, 0]]
        >>> find_all_pos_borderless("AAAA", [])
        [[4, 0]]
        >>> find_all_pos_borderless("GAATTC", [(re.compile("GAATTC"), 6)])
        [[0, 6], [6, 0]]
        >>> find_all_pos_borderless("XXGAATTCYYGATCZZ", [(re.compile("GAATTC"), 6), (re.compile("GATC"), 4)])
        [[2, 6], [10, 4], [16, 0]]
    """
    AllSite: List[List[int]] = []
    for enzyme in ligation_site_list:
        AllSite += find_positions_for_one_site_borderless(text, enzyme)
    # always mark end of sequence
    AllSite.append([len(text), 0])
    return sorted(AllSite, key=lambda x: x[0])


def IndexFragList_borderless(
    all_index_list: List[List[int]], seed_size: int
) -> List[List[int]]:
    """
    From a list [[pos,offset],…], build fragments [start,end] discarding
    those shorter than seed_size, and trimming out the enzyme borders.

    Examples:
        >>> # no fragment (text too short)
        >>> IndexFragList_borderless([[3,2],[5,2]], 10)
        []
        >>> # exact seed_size
        >>> IndexFragList_borderless([[0,1],[5,1],[10,0]], 4)
        [[1, 5], [6, 10]]
    """
    ListFragListIndex: List[List[int]] = []
    previous_position = 0
    for i, (current_position, offset) in enumerate(all_index_list):
        prev_offset = all_index_list[i - 1][1] if i > 0 else 0
        start = previous_position + prev_offset
        end = current_position
        if end - start >= seed_size:
            ListFragListIndex.append([start, end])
        previous_position = current_position
    return ListFragListIndex


def index_list_single_borderless(
    sequence: str,
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
) -> List[List[int]]:
    """
    Compute fragment indices for one sequence in borderless mode.

    Examples:
        >>> # no enzyme sites => single fragment from 0 to len
        >>> index_list_single_borderless("AAAA", [], 0)
        [[0, 4]]
    """
    positions = find_all_pos_borderless(sequence, ligation_site_list)
    return IndexFragList_borderless(positions, seed_size)


def index_list_borderless(
    Sequences: List[str],
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
) -> Tuple[List[List[int]], List[List[int]]]:
    """
    Compute fragments for forward and reverse sequences in borderless mode.

    Examples:
        >>> f,r = index_list_borderless(["AAAA","BBBB"], [], 0)
        >>> f, r
        ([[0, 4]], [[0, 4]])
    """
    for_seq = index_list_single_borderless(
        Sequences[0], ligation_site_list, seed_size
    )
    rev_seq = index_list_single_borderless(
        Sequences[1], ligation_site_list, seed_size
    )
    return for_seq, rev_seq


# ------------------- Classic part (simple positions) -----------------------


def find_positions_for_one_site(
    text: str, Enzyme: Tuple[re.Pattern, int]
) -> List[int]:
    """
    Return end-positions (start + offset) for each occurrence.

    Examples:
        >>> find_positions_for_one_site("", (re.compile("A"),1))
        []
        >>> find_positions_for_one_site("AZAA", (re.compile("ZAA"),3))
        [4]
        >>> find_positions_for_one_site("XYZ", (re.compile("A"),1))
        []
    """
    regex, offset = Enzyme
    return [m.start() + offset for m in regex.finditer(text)]


def find_all_pos(
    text: str, ligation_site_list: List[Tuple[re.Pattern, int]]
) -> List[int]:
    """
    Aggregate positions 0, all site-ends, and len(text) sorted.

    Examples:
        >>> find_all_pos("AAAA", [])
        [0, 4]
        >>> find_all_pos("AXA", [(re.compile("XA"),2)])
        [0, 3]
    """
    sites = [0]
    for enzyme in ligation_site_list:
        sites += find_positions_for_one_site(text, enzyme)
    sites.append(len(text))
    return sorted(set(sites))


def IndexFragList(index_list: List[int], seed_size: int) -> List[List[int]]:
    """
    From a sorted list of positions, extract [prev,curr] where size>seed_size.

    Examples:
        >>> IndexFragList([0,5,10], 0)
        [[0, 5], [5, 10]]
        >>> IndexFragList([0,3,5], 4)
        []
    """
    fragments: List[List[int]] = []
    for prev, curr in zip(index_list, index_list[1:]):
        if curr - prev > seed_size:
            fragments.append([prev, curr])
    return fragments


def index_list_single(
    sequence: str,
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
) -> List[List[int]]:
    """
    Wrapper: find_all_pos + IndexFragList.

    Examples:
        >>> index_list_single("AAAA", [], 0)
        [[0, 4]]
        >>> index_list_single("AXA", [(re.compile("XA"),2)], 0)
        [[0, 3]]
    """
    positions = find_all_pos(sequence, ligation_site_list)
    return IndexFragList(positions, seed_size)


def index_list(
    Sequences: List[str],
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
) -> Tuple[List[List[int]], List[List[int]]]:
    """
    Dual wrapper for forward & reverse.

    Examples:
        >>> index_list(["AAAA","BBBB"], [], 0)
        ([[0, 4]], [[0, 4]])
    """
    return (
        index_list_single(Sequences[0], ligation_site_list, seed_size),
        index_list_single(Sequences[1], ligation_site_list, seed_size),
    )


################################ Mode All #####################################


def Create_Pairs_All(
    Sequence, seed_size, ligation_site_list, indexation
) -> Generator:
    """
    Create all possible pairs of fragments from given sequences.

    Parameters:
        Sequence (List[str]): List containing forward and reverse sequences.
        seed_size (int): Minimum size of the fragment to be considered.
        ligation_site_list (List[Tuple[re.Pattern, int]]): List of ligation sites with regex patterns and offsets.

    Yields:
        List: A list containing:
            - A unique identifier for the pair.
            - Information about the first fragment and which sequence it comes from (forward or reverse).
            - Information about the second fragment and which sequence it comes from (forward or reverse).
    """
    ListFragmentFor, ListFragmentRev = indexation(
        Sequence, ligation_site_list, seed_size
    )
    AllFrag = ListFragmentFor + ListFragmentRev
    NbFragFor = len(ListFragmentFor)
    for i, fragI in enumerate(AllFrag):
        for j, fragJ in enumerate(AllFrag):
            if i < j:
                if i < NbFragFor:
                    WhichFragFor = 0
                    SenseFor = 1
                else:
                    WhichFragFor = 1
                    SenseFor = -1
                if j < NbFragFor:
                    WhichFragRev = 0
                    SenseRev = -1
                else:
                    WhichFragRev = 1
                    SenseRev = 1

                Num = str(i) + str(j)
                if ((fragI[1] - fragI[0]) > seed_size) and (
                    (fragJ[1] - fragJ[0]) > seed_size
                ):
                    yield [
                        Num,
                        [fragI, WhichFragFor, SenseFor],
                        [fragJ, WhichFragRev, SenseRev],
                    ]


def build_read_name(
    raw_name: str, pair_id: str | None, has_single_pair: bool
) -> str:
    """
    Construit le nom de lecture.
    - raw_name        : nom tel qu’il apparaît dans le FASTQ (@...).
    - pair_id         : identifiant ij (chaîne "01", "23"...). None si non pertinent.
    - has_single_pair : True s’il n’y a qu’une seule paire à écrire.
    """
    base = raw_name.split(" ")[0]  # garde tout avant le 1ᵉʳ espace
    if has_single_pair or pair_id is None:
        return base
    return f"{base}:{pair_id}"


def processing_All(
    Name: str,
    Sequence: List[str],
    Quality: List[str],
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
    indexation,
) -> Tuple[List[str], List[str]]:
    """
    Process the sequences to generate buffers for forward and reverse reads
    by creating ALL possible fragment pairs. N'ajoute pas de suffixe :ij si
    il n'y a qu'une seule paire.

    Returns:
        bufferF, bufferR: deux listes de lignes FastQ.

    Doctests:
    >>> # Cas simple, une seule paire possible : pas de suffixe
    >>> seqs = ["AAAACCCCGGGG", "TTTTGGGGCCCC"]
    >>> quals = ["IIIIIIIIIIII", "JJJJJJJJJJJJ"]
    >>> # suppose que indexation ne découpe qu'en un seul fragment [0,12]
    >>> def fake_index(Seq, sites, seed): return ([[0,12]], [[0,12]])
    >>> lig = [(re.compile("CCCC"),4)]
    >>> bufF, bufR = processing_All("readX", seqs, quals, lig, 0, fake_index)
    >>> bufF[0].startswith("readX\\n")
    True
    >>> bufR[0].startswith("readX\\n")
    True

    >>> # Cas multiple paires : suffixe ajouté
    >>> def fake_index2(Seq, sites, seed): return ([[0,6],[6,12]], [[0,6],[6,12]])
    >>> bufF2, bufR2 = processing_All("readY", seqs, quals, lig, 0, fake_index2)
    >>> any(":01\\n" in entry for entry in bufF2)
    True
    """
    bufferF = []
    bufferR = []

    # Collecte d'abord toutes les paires
    all_pairs = [
        pair
        for pair in Create_Pairs_All(
            Sequence, seed_size, ligation_site_list, indexation
        )
        if (pair[1][0][1] - pair[1][0][0] > seed_size)
        and (pair[2][0][1] - pair[2][0][0] > seed_size)
    ]
    single_pair = len(all_pairs) == 1

    for BlocDePairs in all_pairs:
        Num, fragA, fragB = BlocDePairs
        Name = build_read_name(Name, Num, single_pair)
        # suffixe uniquement si >1 paire
        suffix = "" if single_pair else f":{BlocDePairs[0]}"
        base_name = Name.split(" ")[0]
        NewName = base_name + suffix

        # reconstruction forward
        frag_for, which_for, dir_for = BlocDePairs[1]
        seq_f = Sequence[which_for][frag_for[0] : frag_for[1]][::dir_for]
        qual_f = Quality[which_for][frag_for[0] : frag_for[1]][::dir_for]
        bufferF.append(f"{NewName}\n{seq_f}\n+\n{qual_f}\n")

        # reconstruction reverse
        frag_rev, which_rev, dir_rev = BlocDePairs[2]
        seq_r = Sequence[which_rev][frag_rev[0] : frag_rev[1]][::dir_rev]
        qual_r = Quality[which_rev][frag_rev[0] : frag_rev[1]][::dir_rev]
        bufferR.append(f"{NewName}\n{seq_r}\n+\n{qual_r}\n")

    return bufferF, bufferR


"""
    For comprehensive explanation :

    BlocDePairs =   [Num,
                    [[StartIndex, EndIndex], CommingFromForwardOrReverse, ReverseOrNot],
                    [[StartIndex, EndIndex], CommingFCommingFromForwardOrReverseromReverse, ReverseOrNot]]


    bufferF = NewName
              Sequence[CommingFromForwardOrReverse][Index_Fragment][ReverseOrNot]
              +
              Qualitie[CommingFromForwardOrReverse][Index_Fragment][ReverseOrNot]
"""


################################ Mode FR #####################################


def Create_Pairs_Fr(
    fragments: List[List[List[int]]],
) -> Generator[List[Union[str, List[int]]], None, None]:
    """
    Create pairs of forward and reverse fragments.

    Parameters:
        fragments (List[List[List[int]]]): List containing forward and reverse fragments. Each fragment is a list of start and end indices.

    Yields:
        List: A list containing:
            - A unique identifier for the pair.
            - The forward fragment indices.
            - The reverse fragment indices.

    """
    forward_fragments = fragments[0]
    reverse_fragments = fragments[1]

    for i, index_f_frag in enumerate(forward_fragments):
        for j, index_r_frag in enumerate(reverse_fragments):
            Num = str(i) + str(j)
            yield [Num, index_f_frag, index_r_frag]


def processing_Fr(
    TNom: str,
    TSeq: List[str],
    TQual: List[str],
    ligation_site_list: List[Tuple[re.Pattern, int]],
    seed_size: int,
    indexation,
) -> Tuple[List[str], List[str]]:
    """
    Process the sequences to generate buffers for forward and reverse reads
    selon le mode FR (un fragment forward + un fragment reverse).
    N'ajoute pas de suffixe :ij si une seule paire.

    Doctests:
    >>> seqs = ["AAAACCCCGGGG", "TTTTGGGGCCCC"]
    >>> quals = ["IIIIIIIIIIII", "JJJJJJJJJJJJ"]
    >>> # Un seul fragment forward et reverse => une seule paire
    >>> def idx_fr(seq, lig, sd): return ([[0,6]], [[6,12]])
    >>> bufF, bufR = processing_Fr("readZ", seqs, quals, [], 0, idx_fr)
    >>> bufF[0].startswith("readZ\\n")
    True
    >>> bufR[0].startswith("readZ\\n")
    True

    >>> # Plusieurs fragments => plusieurs paires et suffixes
    >>> def idx_fr2(seq, lig, sd): return ([[0,6],[6,12]], [[0,6],[6,12]])
    >>> bufF2, bufR2 = processing_Fr("readW", seqs, quals, [], 0, idx_fr2)
    >>> any(":01\\n" in entry for entry in bufF2)
    True
    """
    bufferF = []
    bufferR = []

    # liste des fragments forward et reverse
    ListFragFor, ListFragRev = indexation(TSeq, ligation_site_list, seed_size)
    all_pairs = list(Create_Pairs_Fr([ListFragFor, ListFragRev]))
    single_pair = len(all_pairs) == 1

    for Num, index_f_frag, index_r_frag in all_pairs:
        suffix = "" if single_pair else f":{Num}"
        base_name = TNom.split(" ")[0]
        NewName = base_name + suffix

        seq_f = TSeq[0][index_f_frag[0] : index_f_frag[1]]
        qual_f = TQual[0][index_f_frag[0] : index_f_frag[1]]
        bufferF.append(f"{NewName}\n{seq_f}\n+\n{qual_f}\n")

        seq_r = TSeq[1][index_r_frag[0] : index_r_frag[1]]
        qual_r = TQual[1][index_r_frag[0] : index_r_frag[1]]
        bufferR.append(f"{NewName}\n{seq_r}\n+\n{qual_r}\n")

    return bufferF, bufferR


################################ Common #####################################


def process_items(
    Input_Buffer,
    Output_buffer,
    ligation_site_list,
    seed_size,
    buffer_size,
    mode,
    borderless,
):
    """
    _summary_ : Process the sequences to generate FastQ sequences paired
    """
    BigBufferF = []
    BigBufferR = []

    if borderless:
        indexation = index_list_borderless
    else:
        indexation = index_list

    if mode == "all":
        Treatment = processing_All
    else:
        Treatment = processing_Fr

    while True:
        try:
            Items = Input_Buffer.get()
            if Items is None:
                break

            for item in Items:
                # Run the synchronous processing in a separate process
                bufferF, bufferR = Treatment(
                    item[0][0],
                    item[1],
                    item[2],
                    ligation_site_list,
                    seed_size,
                    indexation,
                )

                BigBufferF.extend(bufferF)
                BigBufferR.extend(bufferR)

                if len(BigBufferF) > buffer_size:
                    Output_buffer.put([BigBufferF, BigBufferR])
                    BigBufferF = []
                    BigBufferR = []

        except Exception as e:
            logging.error(f"Error in process items : {e}")

    if BigBufferF or BigBufferR:
        Output_buffer.put([BigBufferF, BigBufferR])
    Output_buffer.put(None)
