rsomics-seq-grep 0.1.1

use std::io::Write;

use rsomics_common::{Result, RsomicsError};

/// Write a FASTA record wrapped to `line_width` (0 = no wrap), matching
/// `seqkit`'s `Record.FormatToWriter`: header line verbatim, then the
/// sequence split into fixed-width chunks each followed by `\n`.
pub fn write_fasta(name: &[u8], seq: &[u8], line_width: usize, out: &mut dyn Write) -> Result<()> {
    out.write_all(b">").map_err(RsomicsError::Io)?;
    out.write_all(name).map_err(RsomicsError::Io)?;
    out.write_all(b"\n").map_err(RsomicsError::Io)?;
    write_wrapped(seq, line_width, out)?;
    Ok(())
}

/// Write a FASTQ record. `seqkit grep` always forces `line_width` to 0 for
/// FASTQ input regardless of `-w` — the quality string wrapping upstream
/// supports elsewhere in `seqkit` is dead code on the grep path.
pub fn write_fastq(name: &[u8], seq: &[u8], qual: &[u8], out: &mut dyn Write) -> Result<()> {
    out.write_all(b"@").map_err(RsomicsError::Io)?;
    out.write_all(name).map_err(RsomicsError::Io)?;
    out.write_all(b"\n").map_err(RsomicsError::Io)?;
    out.write_all(seq).map_err(RsomicsError::Io)?;
    out.write_all(b"\n+\n").map_err(RsomicsError::Io)?;
    out.write_all(qual).map_err(RsomicsError::Io)?;
    out.write_all(b"\n").map_err(RsomicsError::Io)?;
    Ok(())
}

fn write_wrapped(seq: &[u8], line_width: usize, out: &mut dyn Write) -> Result<()> {
    if seq.is_empty() {
        out.write_all(b"\n").map_err(RsomicsError::Io)?;
    } else if line_width == 0 || seq.len() <= line_width {
        out.write_all(seq).map_err(RsomicsError::Io)?;
        out.write_all(b"\n").map_err(RsomicsError::Io)?;
    } else {
        for chunk in seq.chunks(line_width) {
            out.write_all(chunk).map_err(RsomicsError::Io)?;
            out.write_all(b"\n").map_err(RsomicsError::Io)?;
        }
    }
    Ok(())
}

/// Same as [`write_fasta`] but appends to a `Vec<u8>` — used on the
/// parallel by-seq path where each record is formatted on a rayon worker
/// before the single-threaded writer flushes buffers in input order.
pub fn format_fasta(name: &[u8], seq: &[u8], line_width: usize, buf: &mut Vec<u8>) {
    write_fasta(name, seq, line_width, buf).expect("writes to Vec<u8> are infallible");
}

pub fn format_fastq(name: &[u8], seq: &[u8], qual: &[u8], buf: &mut Vec<u8>) {
    write_fastq(name, seq, qual, buf).expect("writes to Vec<u8> are infallible");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fasta_wraps_at_width() {
        let mut buf = Vec::new();
        write_fasta(b"id desc", b"ACGTACGTAC", 4, &mut buf).unwrap();
        assert_eq!(buf, b">id desc\nACGT\nACGT\nAC\n");
    }

    #[test]
    fn fasta_no_wrap_when_zero() {
        let mut buf = Vec::new();
        write_fasta(b"id", b"ACGTACGTAC", 0, &mut buf).unwrap();
        assert_eq!(buf, b">id\nACGTACGTAC\n");
    }

    #[test]
    fn fasta_exact_multiple_is_single_line() {
        let mut buf = Vec::new();
        write_fasta(b"id", b"ACGTACGT", 4, &mut buf).unwrap();
        assert_eq!(buf, b">id\nACGT\nACGT\n");
    }

    #[test]
    fn fasta_empty_seq_is_blank_line() {
        let mut buf = Vec::new();
        write_fasta(b"id", b"", 60, &mut buf).unwrap();
        assert_eq!(buf, b">id\n\n");
    }

    #[test]
    fn fastq_ignores_line_width() {
        let mut buf = Vec::new();
        write_fastq(b"read1", b"ACGT", b"IIII", &mut buf).unwrap();
        assert_eq!(buf, b"@read1\nACGT\n+\nIIII\n");
    }
}