"""
Update the FASTQ quality scores from a subsequence FASTQ.

This tool exists because splitcode is unable to retain quality scores when
extracting subsequences from FASTQ files. This script takes a source FASTQ
file and a subsequence FASTQ file (produced by splitcode) and updates the
quality scores in the subsequence file based on the source file.

Upstream issue:

    - https://github.com/pachterlab/splitcode/issues/18

───────
"""

import argparse
import logging
import sys
from gzip import open as gzopen
from logging import Logger
from os import linesep
from pathlib import Path
from typing import cast

from pysam import FastxFile
from pysam import FastxRecord
from syncup import sync

logging.basicConfig(
    level=logging.INFO,
    datefmt="%Y-%m-%dT%H:%M:%S%z",
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
)

LOGGER: Logger = logging.getLogger(Path(__file__).name)
"""A global logger to use for this script."""


def qualrepair(in_source: Path, in_subseq: Path, output: Path) -> None:
    """Update the qualities in `in_subseq` based on `in_source` and write to `output`."""
    with (
        FastxFile(str(in_source), persist=False) as source,
        FastxFile(str(in_subseq), persist=False) as subseq,
        gzopen(output, "wt") if output.suffix == ".gz" else open(output, "w") as writer,
    ):
        for count, (fq, fq_subseq) in enumerate(
            sync(
                iter1=source,
                iter2=subseq,
                key1=lambda rec: cast(str, rec.name),
                key2=lambda rec: cast(str, rec.name),
                cmp_func=lambda x, y: (x == y) - 1,  # only advance iter1 when non-equal
            ),
            start=1,
        ):
            assert fq.name == fq_subseq.name, "Names should be equal!"
            assert fq.quality is not None and fq_subseq.quality is not None
            assert fq.sequence is not None and fq_subseq.sequence is not None

            start = fq.sequence.find(fq_subseq.sequence)
            assert start >= 0, "Subsequence not found in source sequence!"
            quals = fq.quality[start : start + len(fq_subseq.sequence)]

            fastx = FastxRecord(
                name=fq_subseq.name,
                comment=fq_subseq.comment,
                sequence=fq_subseq.sequence,
                quality=quals,
            )

            writer.write(f"{fastx}{linesep}")

            if count % 1_000_000 == 0:
                LOGGER.info(f"Completed writing {count} FASTQ records")

        LOGGER.info(f"Completed writing all {count} FASTQ records")


def cli(cli_args: list[str] | None = None) -> int:
    """Update the FASTQ quality scores from a subsequence FASTQ."""
    cli_args = sys.argv[1:] if cli_args is None else cli_args

    parser = argparse.ArgumentParser(
        description=__doc__,
        add_help=True,
        formatter_class=argparse.RawTextHelpFormatter,
        epilog=r"Copyright © Clint Valentine 2025",
    )
    parser.add_argument(
        "--in-source",
        required=True,
        type=Path,
        help="Source FASTQ file.",
    )
    parser.add_argument(
        "--in-subseq",
        required=True,
        type=Path,
        help="Subsequence FASTQ file.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("/dev/stdout"),
        help="Output repaired FASTQ file.",
    )

    args = parser.parse_args(cli_args)

    try:
        qualrepair(
            in_source=args.in_source,
            in_subseq=args.in_subseq,
            output=args.output,
        )
    except Exception as exception:
        LOGGER.error(f"Error: {exception}")
        return 1

    return 0


if __name__ == "__main__":
    exit(cli())
