# ----------------------------------------------------------------------------
# Copyright (c) 2025, 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.
# ----------------------------------------------------------------------------
import gzip
import os
from itertools import zip_longest
from pathlib import Path

import pandas as pd
from q2_types.sample_data import SampleData

from q2_annotate.kraken2.select import _get_indentation
from q2_types.kraken2 import (
    Kraken2OutputDirectoryFormat,
    Kraken2ReportDirectoryFormat,
    Kraken2ReportFormat,
)
from q2_types.per_sample_sequences import (
    CasavaOneEightSingleLanePerSampleDirFmt,
    SequencesWithQuality,
    JoinedSequencesWithQuality,
    PairedEndSequencesWithQuality,
)


def _normalize_taxon_id(taxon_id) -> str:
    """
    Normalize a taxon ID to a string and remove any trailing '.0'.

    Args:
        taxon_id (any): The taxon ID to normalize.

    Returns:
        str: The normalized taxon ID.
    """
    taxon_id = str(taxon_id).strip()
    if taxon_id.endswith(".0"):
        taxon_id = taxon_id[:-2]
    return taxon_id


def _normalize_read_id(read_id: str) -> str:
    """
    Normalize a read ID by stripping whitespace, removing leading '@',
    taking the first part of a space-separated string, and removing
    trailing '/1' or '/2'.

    Args:
        read_id (str): The read ID to normalize.

    Returns:
        str: The normalized read ID.
    """
    read_id = str(read_id).strip()
    if read_id.startswith("@"):
        read_id = read_id[1:]
    read_id = read_id.split()[0]

    if read_id.endswith("/1") or read_id.endswith("/2"):
        read_id = read_id[:-2]

    return read_id


def _is_gzip_file(filepath: Path) -> bool:
    """
    Check if a file is GZIP compressed by inspecting its magic number.

    Args:
        filepath (Path): The path to the file to check.

    Returns:
        bool: True if the file is GZIP compressed, False otherwise.
    """
    with open(filepath, "rb") as fh:
        return fh.read(2) == b"\x1f\x8b"


def _open_fastq(filepath: Path, mode: str, compress: bool = False):
    """
    Open a FASTQ file in the desired mode, optionally using GZIP compression.

    Args:
        filepath (Path): The path to the output FASTQ file.
        mode (str): The mode to open the file in (e.g., "w", "a").
        compress (bool): If True, the output will be GZIP compressed.

    Returns:
        file object: A file object opened for writing in text mode.
    """
    if compress:
        return gzip.open(filepath, mode=f"{mode}t")
    return open(filepath, mode=mode)


def _assert_distinct_input_output_paths(input_fp: Path, output_fp: Path):
    """
    Ensure that the input and output paths do not refer to the same file.

    Args:
        input_fp (Path): The input file path.
        output_fp (Path): The output file path.

    Raises:
        ValueError: If the input and output paths are identical or refer to
            the same inode.
    """
    if input_fp.resolve() == output_fp.resolve():
        raise ValueError(
            "Input and output FASTQ paths are identical. "
            "Refusing in-place overwrite to prevent data loss."
        )

    # Hard links look like distinct paths to resolve(), but share an inode.
    # Opening output for write truncates that inode, which would destroy the
    # input we are still reading. Unlikely in _filter_reads_kraken2 (fresh
    # output artifact), but required for these stream-and-write helpers.
    if output_fp.exists() and os.path.samefile(str(input_fp), str(output_fp)):
        raise ValueError(
            "Input and output FASTQ paths refer to the same inode (hardlink). "
            "Refusing in-place overwrite to prevent data loss."
        )


def _iter_fastq_records(fh):
    """
    Iterate over FASTQ records in a file handle.

    Args:
        fh (file object): An open file handle for a FASTQ file.

    Yields:
        tuple: A tuple containing (header, sequence, separator, quality).
    """
    while True:
        header = fh.readline()
        if header == "":
            return

        sequence = fh.readline()
        separator = fh.readline()
        quality = fh.readline()

        yield header, sequence, separator, quality


def _extract_matching_read_ids_from_output(
    output_fp: Path, taxon_ids: set[str]
) -> set[str]:
    """
    Extract read IDs from a Kraken2 output file that match any of the
    provided taxon IDs.

    Args:
        output_fp (Path): The path to the Kraken2 output file.
        taxon_ids (set): The set of taxon IDs (as strings) to match against.

    Returns:
        set: The set of read IDs that matched the taxon IDs.
    """
    if not taxon_ids:
        return set()

    matched_read_ids = set()
    with open(output_fp, "r") as fh:
        for line in fh:
            columns = line.rstrip("\n").split("\t")
            if len(columns) < 3:
                continue

            taxon_id = _normalize_taxon_id(columns[2])
            if taxon_id in taxon_ids:
                matched_read_ids.add(_normalize_read_id(columns[1]))

    return matched_read_ids


def _collect_matching_taxon_ids(
    report: Path,
    taxonomy: str,
    include_descendants: bool = True,
    contains: bool = False,
) -> set[str]:
    """
    Collect taxon IDs from a Kraken2 report that match a taxonomy query.

    Args:
        report (Path): The path to the Kraken2 report file.
        taxonomy (str): The taxonomy query (name or ID).
        include_descendants (bool): If True, include all descendant taxon IDs.
            Defaults to True.
        contains (bool): If True, perform substring matching on taxon names.
            Defaults to False.

    Returns:
        set: The set of matching taxon IDs (as strings).
    """
    report_df = Kraken2ReportFormat(report, "r").view(pd.DataFrame)
    names = report_df["name"].astype(str)
    names_clean = names.str.strip()
    taxon_ids = report_df["taxon_id"].map(_normalize_taxon_id)

    if taxonomy.isdigit():
        match_mask = taxon_ids == _normalize_taxon_id(taxonomy)
    else:
        normalized_query = taxonomy.casefold()
        if contains:
            match_mask = names_clean.str.casefold().str.contains(
                normalized_query, regex=False
            )
        else:
            match_mask = names_clean.str.casefold() == normalized_query

    matched_positions = list(report_df.index[match_mask])
    if not matched_positions:
        return set()

    matched_taxa = set()
    indentations = names.map(_get_indentation)

    for position in matched_positions:
        matched_taxa.add(taxon_ids.iloc[position])

        if not include_descendants:
            continue

        parent_indent = indentations.iloc[position]
        for child_position in range(position + 1, len(report_df)):
            if indentations.iloc[child_position] <= parent_indent:
                break
            matched_taxa.add(taxon_ids.iloc[child_position])

    return matched_taxa


def _filter_single_end_fastq(
    input_fp: Path,
    output_fp: Path,
    matched_read_ids: set[str],
    exclude: bool = False,
):
    """
    Filter a single-end FASTQ file based on a set of matched read IDs.

    Args:
        input_fp (Path): The input FASTQ file path.
        output_fp (Path): The output FASTQ file path.
        matched_read_ids (set): The set of read IDs to filter by.
        exclude (bool): If True, exclude the matched read IDs instead of
            retaining them. Defaults to False.
    """
    _assert_distinct_input_output_paths(input_fp, output_fp)

    gzip_input = _is_gzip_file(input_fp)
    with (
        _open_fastq(input_fp, mode="r", compress=gzip_input) as in_fh,
        _open_fastq(output_fp, mode="w", compress=gzip_input) as out_fh,
    ):
        for record in _iter_fastq_records(in_fh):
            read_id = _normalize_read_id(record[0])
            should_keep = read_id in matched_read_ids
            if exclude:
                should_keep = not should_keep

            if should_keep:
                out_fh.writelines(record)


def _filter_paired_end_fastq(
    forward_input_fp: Path,
    reverse_input_fp: Path,
    forward_output_fp: Path,
    reverse_output_fp: Path,
    matched_read_ids: set[str],
    exclude: bool = False,
):
    """
    Filter paired-end FASTQ files based on a set of matched read IDs.

    Args:
        forward_input_fp (Path): The forward input FASTQ file path.
        reverse_input_fp (Path): The reverse input FASTQ file path.
        forward_output_fp (Path): The forward output FASTQ file path.
        reverse_output_fp (Path): The reverse output FASTQ file path.
        matched_read_ids (set): The set of read IDs to filter by.
        exclude (bool): If True, exclude the matched read IDs instead of
            retaining them. Defaults to False.

    Raises:
        ValueError: If the forward and reverse files have different record
            counts or are out of sync.
    """
    _assert_distinct_input_output_paths(forward_input_fp, forward_output_fp)
    _assert_distinct_input_output_paths(reverse_input_fp, reverse_output_fp)

    gzip_forward = _is_gzip_file(forward_input_fp)
    gzip_reverse = _is_gzip_file(reverse_input_fp)

    with (
        _open_fastq(forward_input_fp, mode="r", compress=gzip_forward) as fwd_in,
        _open_fastq(reverse_input_fp, mode="r", compress=gzip_reverse) as rev_in,
        _open_fastq(forward_output_fp, mode="w", compress=gzip_forward) as fwd_out,
        _open_fastq(reverse_output_fp, mode="w", compress=gzip_reverse) as rev_out,
    ):
        for fwd_record, rev_record in zip_longest(
            _iter_fastq_records(fwd_in), _iter_fastq_records(rev_in)
        ):
            fwd_id = _normalize_read_id(fwd_record[0])
            rev_id = _normalize_read_id(rev_record[0])
            if fwd_id != rev_id:
                raise ValueError(
                    "Forward and reverse FASTQ files are out of sync. "
                    f"Found '{fwd_id}' and '{rev_id}'."
                )

            should_keep = fwd_id in matched_read_ids
            if exclude:
                should_keep = not should_keep

            if should_keep:
                fwd_out.writelines(fwd_record)
                rev_out.writelines(rev_record)


def _validate_read_sample_ids(
    read_sample_ids: set[str],
    report_sample_ids: set[str],
    output_sample_ids: set[str],
) -> None:
    """
    Validate that the sample IDs match across reads, reports, and outputs.

    Args:
        read_sample_ids (set): Sample IDs found in the input reads.
        report_sample_ids (set): Sample IDs found in the Kraken2 reports.
        output_sample_ids (set): Sample IDs found in the Kraken2 outputs.

    Raises:
        ValueError: If the sample IDs do not match across all inputs.
    """
    if report_sample_ids != output_sample_ids:
        missing_in_reports = sorted(output_sample_ids - report_sample_ids)
        missing_in_outputs = sorted(report_sample_ids - output_sample_ids)
        raise ValueError(
            "Sample IDs in Kraken2 reports and outputs do not match. "
            f"Missing in reports: {missing_in_reports}. "
            f"Missing in outputs: {missing_in_outputs}."
        )

    missing_in_reads = sorted(report_sample_ids - read_sample_ids)
    if missing_in_reads:
        raise ValueError(
            "Some Kraken2 classification sample IDs are missing from reads. "
            f"Missing in reads: {missing_in_reads}. "
        )

    missing_in_classifications = sorted(read_sample_ids - report_sample_ids)
    if missing_in_classifications:
        raise ValueError(
            "Some read sample IDs are missing from Kraken2 classifications. "
            f"Missing in Kraken2 classifications: {missing_in_classifications}. "
        )


def _filter_reads_kraken2(
    reads: CasavaOneEightSingleLanePerSampleDirFmt,
    reports: Kraken2ReportDirectoryFormat,
    outputs: Kraken2OutputDirectoryFormat,
    taxonomy: str,
    include_descendants: bool = True,
    contains: bool = False,
    exclude: bool = False,
) -> CasavaOneEightSingleLanePerSampleDirFmt:
    """
    Filter reads based on Kraken2 classifications that match a taxonomy query.

    Args:
        reads (CasavaOneEightSingleLanePerSampleDirFmt): Reads used as input
            for Kraken2 classification.
        reports (Kraken2ReportDirectoryFormat): Kraken2 reports generated
            for `reads`.
        outputs (Kraken2OutputDirectoryFormat): Kraken2 output files generated
            for `reads`.
        taxonomy (str): Taxonomy query. This can be a Kraken2 taxon name
            (e.g., "Bacteria") or a taxon ID (e.g., "2").
        include_descendants (bool): If True, include all descendant taxa under
            each matching taxon. Defaults to True.
        contains (bool): If True, perform case-insensitive substring matching
            on taxon names instead of exact matching. Defaults to False.
        exclude (bool): If False (default), retain only reads matching the
            taxonomy query. If True, discard matching reads and retain
            everything else.

    Returns:
        CasavaOneEightSingleLanePerSampleDirFmt: The filtered reads.
    """
    taxonomy = taxonomy.strip()
    if not taxonomy:
        raise ValueError("`taxonomy` cannot be an empty string.")

    sample_fastqs = reads.manifest.to_dict(orient="index")

    report_map = reports.file_dict()
    output_map = outputs.file_dict()

    read_sample_ids = set(sample_fastqs.keys())
    report_sample_ids = set(report_map.keys())
    output_sample_ids = set(output_map.keys())
    _validate_read_sample_ids(read_sample_ids, report_sample_ids, output_sample_ids)

    filtered_reads = CasavaOneEightSingleLanePerSampleDirFmt()

    taxonomy_found = False
    paired = all(fastqs.get("reverse") for fastqs in sample_fastqs.values())

    for sample_id, sample_fps in sample_fastqs.items():
        forward_input = Path(sample_fps["forward"])
        forward_output = filtered_reads.path / forward_input.name

        matched_taxon_ids = _collect_matching_taxon_ids(
            report_map[sample_id],
            taxonomy=taxonomy,
            include_descendants=include_descendants,
            contains=contains,
        )
        if matched_taxon_ids:
            taxonomy_found = True

        matched_read_ids = _extract_matching_read_ids_from_output(
            output_map[sample_id], matched_taxon_ids
        )

        if paired:
            reverse_input = Path(sample_fps["reverse"])
            reverse_output = filtered_reads.path / reverse_input.name
            _filter_paired_end_fastq(
                forward_input_fp=forward_input,
                reverse_input_fp=reverse_input,
                forward_output_fp=forward_output,
                reverse_output_fp=reverse_output,
                matched_read_ids=matched_read_ids,
                exclude=exclude,
            )
        else:
            _filter_single_end_fastq(
                input_fp=forward_input,
                output_fp=forward_output,
                matched_read_ids=matched_read_ids,
                exclude=exclude,
            )

    if not taxonomy_found:
        raise ValueError(
            f"Taxonomy query '{taxonomy}' was not found in any Kraken2 report."
        )

    return filtered_reads


def filter_reads_kraken2(
    ctx,
    reads,
    reports,
    outputs,
    taxonomy,
    include_descendants=True,
    contains=False,
    exclude=False,
    num_partitions=1,
):
    """
    Filter reads based on Kraken2 classifications that match a taxonomy query.

    Args:
        reads (CasavaOneEightSingleLanePerSampleDirFmt): Reads used as input
            for Kraken2 classification.
        reports (Kraken2ReportDirectoryFormat): Kraken2 reports generated
            for `reads`.
        outputs (Kraken2OutputDirectoryFormat): Kraken2 output files generated
            for `reads`.
        taxonomy (str): Taxonomy query. This can be a Kraken2 taxon name
            (e.g., "Bacteria") or a taxon ID (e.g., "2").
        include_descendants (bool): If True, include all descendant taxa under
            each matching taxon. Defaults to True.
        contains (bool): If True, perform case-insensitive substring matching
            on taxon names instead of exact matching. Defaults to False.
        exclude (bool): If False (default), retain only reads matching the
            taxonomy query. If True, discard matching reads and retain
            everything else.
        num_partitions (int): Number of partitions to use by parsl.
            Defaults to 1.

    Returns:
        CasavaOneEightSingleLanePerSampleDirFmt: The filtered reads.
    """
    kwargs = {
        k: v
        for k, v in locals().items()
        if k not in ["ctx", "reads", "reports", "outputs", "num_partitions"]
    }

    if reads.type <= SampleData[SequencesWithQuality | JoinedSequencesWithQuality]:
        _partition_reads = ctx.get_action("demux", "partition_samples_single")
    elif reads.type <= SampleData[PairedEndSequencesWithQuality]:
        _partition_reads = ctx.get_action("demux", "partition_samples_paired")
    else:
        raise NotImplementedError()

    _partition_reports = ctx.get_action("types", "partition_kraken2_reports")
    _partition_outputs = ctx.get_action("types", "partition_kraken2_outputs")
    _filter_reads_kraken2 = ctx.get_action("annotate", "_filter_reads_kraken2")
    _collate_reads = ctx.get_action("fondue", "combine_seqs")

    (partitioned_reads,) = _partition_reads(reads, num_partitions)
    (partitioned_reports,) = _partition_reports(reports, num_partitions)
    (partitioned_outputs,) = _partition_outputs(outputs, num_partitions)

    results = []
    for _reads, _reports, _outputs in zip(
        partitioned_reads.values(),
        partitioned_reports.values(),
        partitioned_outputs.values(),
    ):
        (result,) = _filter_reads_kraken2(
            reads=_reads, reports=_reports, outputs=_outputs, **kwargs
        )
        results.append(result)

    (combined_reads,) = _collate_reads(results)

    return combined_reads
