twitcher 0.6.8

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

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)?)
    }

    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 mut seq = tokio::task::block_in_place(|| {
            reader.fetch_seq(
                contig_name,
                query.start().position_0(),
                query.end_excl().unwrap().position_0() - 1,
            )
        })?;

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

    pub async fn get_seq(
        &self,
        mut 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}");
        }

        query = query.intersection(&this_contig).unwrap();

        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 mut reader_query =
                GenomeRegion::from_incl_excl(reader_query_start, reader_query_end)?;
            reader_query = reader_query.intersection(&this_contig).unwrap(); // might cut off at the end
            reader_query
        };

        let mut seq = tokio::task::block_in_place(|| {
            reader.fetch_seq(
                contig_name,
                reader_query.start().position_0(),
                reader_query.end_excl().unwrap().position_0() - 1,
            )
        })?;

        // 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().unwrap());
            Ok(Some(ReferenceQueryResult {
                region: reader_query,
                sequence: seq.into(),
                range_in_sequence: rs,
            }))
        }
    }

    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().unwrap()),
        );

        let ok = sequence[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().unwrap()),
            })
        } 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(_)) => unreachable!(),
            (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: &std::path::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");
    }
}