twitcher 0.7.0

use std::{
    ops::Range,
    path::{Path, PathBuf},
};

use anyhow::Context;
use rust_htslib::faidx;
use tokio::sync::Mutex;

use crate::common::{
    ImmutableSequence, MutableSequence,
    coords::{GenomePosition, GenomeRegion},
};

pub struct ReferenceReader {
    name: String,
    inner: Mutex<faidx::Reader>,
    soft_mask: bool,
}

#[derive(Debug)]
pub struct ReferenceQueryResult {
    pub region: GenomeRegion,
    pub sequence: ImmutableSequence,
    pub range_in_sequence: Range<usize>,
}

impl ReferenceReader {
    pub fn get_name(&self) -> &str {
        &self.name
    }

    /// Length of `contig` in the reference index, or an error when it is absent.
    ///
    /// `faidx_seq_len` reports an unknown sequence as `-1`, which rust-htslib hands back as
    /// `u64::MAX`. That value must never reach [`faidx::Reader::fetch_seq`]: htslib returns a
    /// null pointer for an unknown sequence, which rust-htslib wraps into a `Vec` without
    /// checking, aborting the process with a non-unwinding panic that no caller can catch.
    /// Rejecting it here keeps a contig missing from the reference to a single failed query.
    fn contig_len(reader: &faidx::Reader, contig: &str) -> anyhow::Result<usize> {
        let seq_len = tokio::task::block_in_place(|| reader.fetch_seq_len(contig));
        if i64::try_from(seq_len).is_err() {
            anyhow::bail!("Contig {contig} is not present in the reference");
        }
        Ok(usize::try_from(seq_len)?)
    }

    /// The inclusive end [`faidx::Reader::fetch_seq`] expects, i.e. one before `end_excl`.
    ///
    /// An empty region has no such position: `end_excl - 1` would underflow at the start of a
    /// contig, and `fetch_seq(contig, n, n - 1)` returns a null pointer that rust-htslib wraps
    /// into a `Vec` without checking, aborting the process. A zero-width cluster (a lone
    /// insertion) reaches this only when it is queried without any padding.
    fn last_position(
        start: usize,
        end_excl: usize,
        region: &GenomeRegion,
    ) -> anyhow::Result<usize> {
        if end_excl <= start {
            anyhow::bail!(
                "Cannot read the empty region {region} from the reference; \
                 query it with padding to get a non-empty sequence"
            );
        }
        Ok(end_excl - 1)
    }

    pub async fn get_seq_exact_unmasked(
        &self,
        query: GenomeRegion,
    ) -> anyhow::Result<ImmutableSequence> {
        let contig = query.contig().clone();
        let contig_name = contig.as_str()?;
        let reader = self.inner.lock().await;
        let seq_len = Self::contig_len(&reader, contig_name)?;
        let this_contig =
            GenomeRegion::new_bounded(GenomePosition::new_0(contig.clone(), 0), seq_len);

        if !this_contig.contains(&query) {
            anyhow::bail!("Coordinates out of bounds: {query}");
        }
        let (q_start, q_end_excl) = (
            query.start().position_0(),
            query
                .end_excl()
                .context("query has no end coordinate")?
                .position_0(),
        );
        let q_end = Self::last_position(q_start, q_end_excl, &query)?;

        let mut seq =
            tokio::task::block_in_place(|| reader.fetch_seq(contig_name, q_start, q_end))?;
        drop(reader);

        seq.make_ascii_uppercase();
        Ok(seq.into())
    }

    pub async fn get_seq(
        &self,
        query: GenomeRegion,
        padding_left: usize,
        padding_right: usize,
    ) -> anyhow::Result<Option<ReferenceQueryResult>> {
        let contig = query.contig().clone();
        let contig_name = contig.as_str()?;
        let reader = self.inner.lock().await;
        let seq_len = Self::contig_len(&reader, contig_name)?;
        let this_contig =
            GenomeRegion::new_bounded(GenomePosition::new_0(contig.clone(), 0), seq_len);

        if !this_contig.contains(&query) {
            anyhow::bail!("Coordinates out of bounds: {query}");
        }

        let Some(query) = query.intersection(&this_contig) else {
            anyhow::bail!("no intersection even though contained?");
        };

        let reader_query = {
            let reader_query_start =
                query.start().clone() - padding_left.min(query.start().position_0());
            let reader_query_end = query.end_excl().map(|end| end + padding_right);
            let reader_query = GenomeRegion::from_incl_excl(reader_query_start, reader_query_end)?;
            reader_query
                .intersection(&this_contig)
                .context("no intersection?")? // might cut off at the end
        };

        let (rq_start, rq_end_excl) = (
            reader_query.start().position_0(),
            reader_query
                .end_excl()
                .context("query has no end coordinate")?
                .position_0(),
        );
        let rq_end = Self::last_position(rq_start, rq_end_excl, &reader_query)?;

        let mut seq =
            tokio::task::block_in_place(|| reader.fetch_seq(contig_name, rq_start, rq_end))?;
        drop(reader);

        // Now we have query region and sequence, only need to check for soft_mask
        if self.soft_mask {
            let result = Self::exclude_lowercase(seq, &reader_query, &query);
            Ok(result)
        } else {
            seq.make_ascii_uppercase();
            let rs = reader_query.start().abs_diff(query.start())
                ..reader_query
                    .start()
                    .abs_diff(&query.end_excl().context("no query end")?);
            Ok(Some(ReferenceQueryResult {
                region: reader_query,
                sequence: seq.into(),
                range_in_sequence: rs,
            }))
        }
    }

    /// Drop the result when the queried region itself is soft-masked. Only the query's own bases
    /// are checked, never the padding around it — and a query that covers no base at all, such as
    /// the region of a lone insertion, is therefore never masked out.
    fn exclude_lowercase(
        mut sequence: MutableSequence,
        sequence_coordinates: &GenomeRegion,
        query_coordinates: &GenomeRegion,
    ) -> Option<ReferenceQueryResult> {
        let to_offset =
            |pos: &GenomePosition| pos.position_0() - sequence_coordinates.start().position_0();

        let (query_start_offset, query_end_excl_offset) = (
            to_offset(query_coordinates.start()),
            to_offset(&query_coordinates.end_excl()?),
        );

        let ok = sequence
            .get(query_start_offset..query_end_excl_offset)?
            .iter()
            .all(u8::is_ascii_uppercase);

        if ok {
            sequence.make_ascii_uppercase();
            Some(ReferenceQueryResult {
                region: sequence_coordinates.clone(),
                sequence: sequence.into(),
                range_in_sequence: to_offset(query_coordinates.start())
                    ..to_offset(&query_coordinates.end_excl()?),
            })
        } else {
            None
        }
    }
}

impl TryFrom<&CliReferenceArg> for ReferenceReader {
    type Error = rust_htslib::errors::Error;

    fn try_from(value: &CliReferenceArg) -> Result<Self, Self::Error> {
        let path = PathBuf::from((&value.file).as_ref());
        let inner = faidx::Reader::from_path(&path)?;
        let name = path
            .file_name()
            .and_then(|str| str.to_str())
            .unwrap_or("<unknown>")
            .to_string();
        Ok(Self {
            name,
            inner: inner.into(),
            soft_mask: value.soft_mask,
        })
    }
}

#[derive(clap::Args, Clone, Debug)]
pub struct CliReferenceArg {
    #[command(flatten)]
    pub file: CliReferenceFileArg,

    /// When enabled, lowercase letters in the reference sequence will be excluded from any alignments.
    #[arg(long = "soft-mask")]
    pub soft_mask: bool,
}

#[derive(clap::Args, Clone, Debug)]
#[group(multiple = false, required = true)]
pub struct CliReferenceFileArg {
    /// The FASTA file containing the reference as the second positional argument.
    pub reference: Option<String>,

    /// The reference may also be provided as an argument with --reference at any position.
    #[arg(long = "reference", value_name = "FILE")]
    pub reference_arg: Option<String>,
}

impl AsRef<Path> for &CliReferenceFileArg {
    fn as_ref(&self) -> &Path {
        let path = match (&self.reference_arg, &self.reference) {
            (None, None) | (Some(_), Some(_)) => {
                #[expect(clippy::unreachable)]
                {
                    unreachable!("statically guaranteed by clap")
                }
            }
            (None, Some(path)) | (Some(path), None) => path,
        };
        path.as_ref()
    }
}

impl<S: AsRef<str>> From<S> for CliReferenceArg {
    fn from(value: S) -> Self {
        Self {
            file: CliReferenceFileArg {
                reference: Some(value.as_ref().to_string()),
                reference_arg: None,
            },
            soft_mask: false,
        }
    }
}

impl Default for CliReferenceArg {
    fn default() -> Self {
        Self {
            file: CliReferenceFileArg {
                reference: Some("reference.fa".to_string()),
                reference_arg: None,
            },
            soft_mask: false,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::common::contig::ContigName;

    /// Write a single-contig FASTA plus its `.fai` and open a reader on it.
    fn reader_with_chr1(dir: &Path) -> ReferenceReader {
        let path = dir.join("ref.fa");
        std::fs::write(&path, b">chr1\nACGTACGTAC\n").unwrap();
        faidx::build(&path).unwrap();
        ReferenceReader::try_from(&CliReferenceArg::from(path.to_str().unwrap())).unwrap()
    }

    fn region(contig: &[u8], start: usize, len: usize) -> GenomeRegion {
        GenomeRegion::new_bounded(GenomePosition::new_0(ContigName::new(contig), start), len)
    }

    /// A contig absent from the reference must yield an error, not abort the process.
    ///
    /// `faidx_seq_len` reports an unknown sequence as `-1`, which rust-htslib returns as
    /// `u64::MAX`; feeding that contig to `fetch_seq` gets a null pointer back that
    /// rust-htslib wraps in a `Vec` unchecked, killing the process with a non-unwinding
    /// panic. Reaching an assertion at all is the substance of these tests.
    #[tokio::test(flavor = "multi_thread")]
    async fn missing_contig_errors_get_seq_exact_unmasked() {
        let dir = tempfile::tempdir().unwrap();
        let reader = reader_with_chr1(dir.path());
        assert!(
            reader
                .get_seq_exact_unmasked(region(b"chr_absent", 0, 4))
                .await
                .is_err()
        );
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn missing_contig_errors_get_seq() {
        let dir = tempfile::tempdir().unwrap();
        let reader = reader_with_chr1(dir.path());
        assert!(
            reader
                .get_seq(region(b"chr_absent", 0, 4), 2, 2)
                .await
                .is_err()
        );
    }

    /// The guard must not reject contigs that are present.
    #[tokio::test(flavor = "multi_thread")]
    async fn present_contig_still_reads() {
        let dir = tempfile::tempdir().unwrap();
        let reader = reader_with_chr1(dir.path());
        let seq = reader
            .get_seq_exact_unmasked(region(b"chr1", 0, 4))
            .await
            .unwrap();
        assert_eq!(&*seq, b"ACGT");
    }
}