rasusa 5.0.1

use crate::format::OutputEncoding;
use crate::source::RecordSource;
use needletail::errors::ParseErrorKind::EmptyFile;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use thiserror::Error;

/// A collection of custom errors relating to the working with files for this package.
#[derive(Error, Debug)]
pub enum FastxError {
    /// Indicates that the specified input file could not be opened/read.
    #[error("Read error")]
    ReadError {
        source: needletail::errors::ParseError,
    },

    /// Indicates that a sequence record could not be parsed.
    #[error("Failed to parse record")]
    ParseError {
        source: needletail::errors::ParseError,
    },

    /// Indicates that the specified output file could not be created.
    #[error("Output file could not be created")]
    CreateError { source: std::io::Error },

    /// Indicates and error trying to create the compressor
    #[error(transparent)]
    CompressOutputError(#[from] niffler::Error),

    /// Indicates that some indices we expected to find in the input file weren't found.
    #[error("Some expected indices were not in the input file")]
    IndicesNotFound,

    /// Indicates that writing to the output file failed.
    #[error("Could not write to output file")]
    WriteError { source: anyhow::Error },

    /// Indicates an error when reading an unaligned alignment file (SAM/BAM/CRAM).
    #[error("Alignment read error: {source}")]
    AlignmentReadError { source: std::io::Error },

    /// Indicates that a mapped read was detected in the input alignment file.
    #[error("Error: Mapped read detected, please use `rasusa aln` for aligned data")]
    MappedReadDetected,
}

/// A `Struct` used for seamlessly dealing with either compressed or uncompressed fasta/fastq files.
#[derive(Debug, PartialEq)]
pub struct Fastx {
    /// The path for the file.
    path: PathBuf,
}

impl Fastx {
    /// Create a `Fastx` object from a `std::path::Path`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rasusa::fastx::Fastx;
    ///
    /// let path = std::path::Path::new("input.fa.gz");
    /// let fastx = Fastx::from_path(path);
    /// ```
    pub fn from_path(path: &Path) -> Self {
        Fastx {
            path: path.to_path_buf(),
        }
    }
}

impl Fastx {
    /// Opens the underlying (possibly compressed) file as a FASTA/Q record reader.
    ///
    /// Returns `Ok(None)` for a too-short/empty file, matching the pre-existing
    /// "no records" handling shared by [`RecordSource::read_lengths`] and
    /// [`RecordSource::count`].
    fn open_reader(&self) -> Result<Option<Box<dyn needletail::parser::FastxReader>>, FastxError> {
        let reader = match niffler::send::from_path(&self.path) {
            Ok((rdr, _)) => rdr,
            Err(niffler::error::Error::FileTooShort) => return Ok(None),
            Err(source) => return Err(FastxError::CompressOutputError(source)),
        };
        match needletail::parse_fastx_reader(reader) {
            Ok(rdr) => Ok(Some(rdr)),
            Err(e) if e.kind == EmptyFile => Ok(None),
            Err(source) => Err(FastxError::ReadError { source }),
        }
    }
}

impl RecordSource for Fastx {
    /// Returns a vector containing the lengths of all the reads in the file.
    ///
    /// # Errors
    /// If the file cannot be opened or there is an issue parsing any records then an
    /// `Err` containing a variant of [`FastxError`](#fastxerror) is returned.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rasusa::fastx::Fastx;
    /// use rasusa::source::RecordSource;
    /// use std::io::Write;
    /// let text = "@read1\nACGT\n+\n!!!!\n@read2\nG\n+\n!";
    /// let mut file = tempfile::Builder::new().suffix(".fq").tempfile().unwrap();
    /// file.write_all(text.as_bytes()).unwrap();
    /// let fastx = Fastx::from_path(file.path());
    /// let actual = fastx.read_lengths().unwrap();
    /// let expected: Vec<u32> = vec![4, 1];
    /// assert_eq!(actual, expected)
    /// ```
    fn read_lengths(&self) -> Result<Vec<u32>, FastxError> {
        let mut read_lengths: Vec<u32> = vec![];

        let mut reader = match self.open_reader()? {
            Some(rdr) => rdr,
            None => return Ok(read_lengths),
        };

        while let Some(record) = reader.next() {
            match record {
                Ok(rec) => read_lengths.push(rec.num_bases() as u32),
                Err(err) => return Err(FastxError::ParseError { source: err }),
            }
        }
        Ok(read_lengths)
    }

    /// Returns the number of records in the file, without materializing their lengths.
    ///
    /// # Example
    ///
    /// ```rust
    /// use rasusa::fastx::Fastx;
    /// use rasusa::source::RecordSource;
    /// use std::io::Write;
    /// let text = "@read1\nACGT\n+\n!!!!\n@read2\nG\n+\n!";
    /// let mut file = tempfile::Builder::new().suffix(".fq").tempfile().unwrap();
    /// file.write_all(text.as_bytes()).unwrap();
    /// let fastx = Fastx::from_path(file.path());
    /// assert_eq!(fastx.count().unwrap(), 2)
    /// ```
    fn count(&self) -> Result<usize, FastxError> {
        let mut count: usize = 0;

        let mut reader = match self.open_reader()? {
            Some(rdr) => rdr,
            None => return Ok(count),
        };

        while let Some(record) = reader.next() {
            match record {
                Ok(_) => count += 1,
                Err(err) => return Err(FastxError::ParseError { source: err }),
            }
        }
        Ok(count)
    }

    /// Writes reads, with indices contained within `reads_to_keep`, to the specified handle
    /// `write_to`.
    ///
    /// # Errors
    /// This function could raise an `Err` instance of [`FastxError`](#fastxerror) in the following
    /// circumstances:
    /// -   If the file (of `self`) cannot be opened.
    /// -   If writing to `write_to` fails.
    /// -   If, after iterating through all reads in the file, there is still elements left in
    ///     `reads_to_keep`. *Note: in this case, this function still writes all reads where indices
    ///     were found in the file.*
    fn filter_reads_into(
        &self,
        reads_to_keep: &[bool],
        nb_reads_keep: usize,
        write_to: &mut dyn Write,
        encoding: OutputEncoding,
    ) -> Result<usize, FastxError> {
        // Callers never ask a FASTA/Q source for alignment-format output (`Reads::run` rejects
        // that combination before any `RecordSource` is touched), so this is the only branch
        // this impl needs to handle.
        let is_fasta = match encoding {
            OutputEncoding::Fastx { fasta } => fasta,
            OutputEncoding::Alignment(_) => {
                unreachable!("fastx sources never receive alignment output encoding")
            }
        };

        let mut total_len = 0;

        let (reader, _) = niffler::send::from_path(&self.path)?;
        let mut reader = needletail::parse_fastx_reader(reader)
            .map_err(|source| FastxError::ReadError { source })?;
        let mut read_idx: usize = 0;
        let mut nb_reads_written = 0;

        while let Some(record) = reader.next() {
            match record {
                Err(source) => return Err(FastxError::ParseError { source }),
                Ok(rec) if read_idx < reads_to_keep.len() && reads_to_keep[read_idx] => {
                    total_len += rec.num_bases();
                    if is_fasta {
                        write_to
                            .write_all(b">")
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                        write_to
                            .write_all(rec.id())
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                        write_to
                            .write_all(b"\n")
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                        write_to
                            .write_all(&rec.seq())
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                        write_to
                            .write_all(b"\n")
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                    } else {
                        rec.write(write_to, None)
                            .map_err(|err| FastxError::WriteError {
                                source: anyhow::Error::from(err),
                            })?;
                    }
                    nb_reads_written += 1;
                    if nb_reads_keep == nb_reads_written {
                        break;
                    }
                }
                Ok(_) => (),
            }

            read_idx += 1;
        }

        if nb_reads_written == nb_reads_keep {
            Ok(total_len)
        } else {
            Err(FastxError::IndicesNotFound)
        }
    }
}

/// Create a file for writing.
///
/// # Errors
/// If the file cannot be created then an `Err` containing a variant of [`FastxError`](#fastxerror) is
/// returned.
pub fn create_output_writer(
    path: &Path,
    compression_lvl: Option<niffler::compression::Level>,
    compression_fmt: Option<niffler::compression::Format>,
) -> Result<Box<dyn Write>, FastxError> {
    let file = File::create(path).map_err(|source| FastxError::CreateError { source })?;
    let file_handle = Box::new(BufWriter::new(file));
    let fmt = compression_fmt.unwrap_or_else(|| crate::format::infer_compression_format(path));
    let compression_lvl =
        compression_lvl.unwrap_or_else(|| crate::format::default_compression_level(fmt));
    niffler::get_writer(file_handle, fmt, compression_lvl).map_err(FastxError::CompressOutputError)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::any::Any;
    use std::io::{Read, Write};
    use std::path::Path;
    use tempfile::Builder;

    #[test]
    fn fastx_from_fasta() {
        let path = Path::new("data/my.fa");

        let actual = Fastx::from_path(path);
        let expected = Fastx {
            path: path.to_path_buf(),
        };

        assert_eq!(actual, expected)
    }

    #[test]
    fn create_invalid_output_file_raises_error() {
        let path = Path::new("invalid/out/path.fq");

        let actual = create_output_writer(path, Some(niffler::Level::Eight), None)
            .err()
            .unwrap();
        let expected = FastxError::CreateError {
            source: std::io::Error::other(String::from("No such file or directory (os error 2)")),
        };

        assert_eq!(actual.type_id(), expected.type_id())
    }

    #[test]
    fn create_valid_output_file_and_can_write_to_it() {
        let file = Builder::new().suffix(".fastq").tempfile().unwrap();
        let mut writer =
            create_output_writer(file.path(), Some(niffler::Level::Eight), None).unwrap();

        let actual = writer.write(b"foo\nbar");

        assert!(actual.is_ok())
    }

    #[test]
    fn create_valid_compressed_output_file_and_can_write_to_it() {
        let file = Builder::new().suffix(".fastq.gz").tempfile().unwrap();
        let mut writer =
            create_output_writer(file.path(), Some(niffler::Level::Four), None).unwrap();

        let actual = writer.write(b"foo\nbar");

        assert!(actual.is_ok())
    }

    #[test]
    fn get_read_lengths_for_empty_fasta_returns_empty_vector() {
        let text = "";
        let mut file = Builder::new().suffix(".fa").tempfile().unwrap();
        file.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(file.path());

        let actual = fastx.read_lengths().unwrap();
        let expected: Vec<u32> = Vec::new();

        assert_eq!(actual, expected)
    }

    #[test]
    fn get_read_lengths_for_fasta() {
        let text = ">read1\nACGT\n>read2\nG";
        let mut file = Builder::new().suffix(".fa").tempfile().unwrap();
        file.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(file.path());

        let actual = fastx.read_lengths().unwrap();
        let expected: Vec<u32> = vec![4, 1];

        assert_eq!(actual, expected)
    }

    #[test]
    fn get_read_lengths_for_fastq() {
        let text = "@read1\nACGT\n+\n!!!!\n@read2\nG\n+\n!";
        let mut file = Builder::new().suffix(".fq").tempfile().unwrap();
        file.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(file.path());

        let actual = fastx.read_lengths().unwrap();
        let expected: Vec<u32> = vec![4, 1];

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_reads_empty_indices_no_output() {
        let text = "@read1\nACGT\n+\n!!!!";
        let mut input = Builder::new().suffix(".fastq").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![false];
        let output = Builder::new().suffix(".fastq").tempfile().unwrap();
        let mut out_fh = create_output_writer(output.path(), None, None).unwrap();
        let filter_result = fastx.filter_reads_into(
            &reads_to_keep,
            0,
            &mut out_fh,
            OutputEncoding::Fastx { fasta: false },
        );

        assert!(filter_result.is_ok());

        let mut actual = String::new();
        output.into_file().read_to_string(&mut actual).unwrap();
        let expected = String::new();

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fastq_reads_one_index_matches_only_read() {
        let text = "@read1\nACGT\n+\n!!!!\n";
        let mut input = Builder::new().suffix(".fastq").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![true];
        let output = Builder::new().suffix(".fastq").tempfile().unwrap();
        {
            let mut out_fh = create_output_writer(output.path(), None, None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                1,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: false },
            );
            assert!(filter_result.is_ok());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = text;

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fasta_reads_one_index_matches_only_read() {
        let text = ">read1\nACGT\n";
        let mut input = Builder::new().suffix(".fa").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![true];
        let output = Builder::new().suffix(".fa").tempfile().unwrap();
        {
            let mut out_fh = create_output_writer(output.path(), None, None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                1,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: true },
            );
            assert!(filter_result.is_ok());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = text;

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fastq_reads_one_index_matches_one_of_two_reads() {
        let text = "@read1\nACGT\n+\n!!!!\n@read2\nCCCC\n+\n$$$$\n";
        let mut input = Builder::new().suffix(".fastq").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![false, true];
        let output = Builder::new().suffix(".fastq").tempfile().unwrap();
        {
            let mut out_fh = create_output_writer(output.path(), None, None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                1,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: false },
            );
            assert!(filter_result.is_ok());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = "@read2\nCCCC\n+\n$$$$\n";

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fastq_reads_two_indices_matches_first_and_last_reads() {
        let text = "@read1\nACGT\n+\n!!!!\n@read2\nCCCC\n+\n$$$$\n@read3\nA\n+\n$\n";
        let mut input = Builder::new().suffix(".fastq").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![true, false, true];
        let output = Builder::new().suffix(".fastq").tempfile().unwrap();
        {
            let mut out_fh = create_output_writer(output.path(), None, None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                2,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: false },
            );
            assert!(filter_result.is_ok());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = "@read1\nACGT\n+\n!!!!\n@read3\nA\n+\n$\n";

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fasta_reads_one_index_out_of_range() {
        let text = ">read1 length=4\nACGT\n>read2\nCCCC\n";
        let mut input = Builder::new().suffix(".fa").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![true, false, true];
        let output = Builder::new().suffix(".fa").tempfile().unwrap();
        {
            let mut out_fh =
                create_output_writer(output.path(), Some(niffler::Level::Four), None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                2,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: true },
            );
            assert!(filter_result.is_err());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = ">read1 length=4\nACGT\n";

        assert_eq!(actual, expected)
    }

    #[test]
    fn filter_fastq_reads_one_index_out_of_range() {
        let text = "@read1 length=4\nACGT\n+\n!!!!\n@read2\nC\n+\n^\n";
        let mut input = Builder::new().suffix(".fq").tempfile().unwrap();
        input.write_all(text.as_bytes()).unwrap();
        let fastx = Fastx::from_path(input.path());
        let reads_to_keep: Vec<bool> = vec![true, false, true];
        let output = Builder::new().suffix(".fq").tempfile().unwrap();
        {
            let mut out_fh =
                create_output_writer(output.path(), Some(niffler::Level::Four), None).unwrap();
            let filter_result = fastx.filter_reads_into(
                &reads_to_keep,
                2,
                &mut out_fh,
                OutputEncoding::Fastx { fasta: false },
            );
            assert!(filter_result.is_err());
        }

        let actual = std::fs::read_to_string(output).unwrap();
        let expected = "@read1 length=4\nACGT\n+\n!!!!\n";

        assert_eq!(actual, expected)
    }
}