jellyfish-reader 0.1.0

use std::io::BufRead;

use crate::error::{Error, Result};
use crate::mer::MerDna;

/// Sequential reader for Jellyfish text/sorted format files.
///
/// Reads lines of the form `<kmer_string> <count>` and yields (MerDna, u64) pairs.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::BufReader;
/// use jellyfish_reader::{FileHeader, TextReader};
///
/// let file = File::open("output.txt").unwrap();
/// let mut reader = BufReader::new(file);
/// let header = FileHeader::read(&mut reader).unwrap();
///
/// let text_reader = TextReader::new(reader);
/// for result in text_reader {
///     let (mer, count) = result.unwrap();
///     println!("{}: {}", mer, count);
/// }
/// ```
pub struct TextReader<R: BufRead> {
    reader: R,
    line_buf: String,
}

impl<R: BufRead> TextReader<R> {
    /// Create a new text reader from a buffered reader positioned at the start of data.
    pub fn new(reader: R) -> Self {
        Self {
            reader,
            line_buf: String::new(),
        }
    }

    fn read_next(&mut self) -> Result<Option<(MerDna, u64)>> {
        self.line_buf.clear();
        let bytes_read = self.reader.read_line(&mut self.line_buf)?;
        if bytes_read == 0 {
            return Ok(None);
        }

        let line = self.line_buf.trim();
        if line.is_empty() {
            return Ok(None);
        }

        // Format: "<kmer> <count>" or "<kmer>\t<count>"
        let (mer_str, count_str) = line
            .split_once(|c: char| c.is_whitespace())
            .ok_or_else(|| Error::InvalidHeader(format!("invalid text format line: {line:?}")))?;

        let mer: MerDna = mer_str.parse()?;
        let count: u64 = count_str
            .trim()
            .parse()
            .map_err(|_| Error::InvalidHeader(format!("invalid count value: {count_str:?}")))?;

        Ok(Some((mer, count)))
    }
}

impl<R: BufRead> Iterator for TextReader<R> {
    type Item = Result<(MerDna, u64)>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.read_next() {
            Ok(Some(pair)) => Some(Ok(pair)),
            Ok(None) => None,
            Err(e) => Some(Err(e)),
        }
    }
}

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

    fn make_text_reader(content: &str) -> TextReader<std::io::BufReader<Cursor<Vec<u8>>>> {
        let cursor = Cursor::new(content.as_bytes().to_vec());
        TextReader::new(std::io::BufReader::new(cursor))
    }

    #[test]
    fn test_read_single_line() {
        let reader = make_text_reader("ACGT 42\n");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0.to_string(), "ACGT");
        assert_eq!(results[0].1, 42);
    }

    #[test]
    fn test_read_multiple_lines() {
        let reader = make_text_reader("AAAA 10\nACGT 42\nTTTT 100\n");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].0.to_string(), "AAAA");
        assert_eq!(results[0].1, 10);
        assert_eq!(results[1].0.to_string(), "ACGT");
        assert_eq!(results[1].1, 42);
        assert_eq!(results[2].0.to_string(), "TTTT");
        assert_eq!(results[2].1, 100);
    }

    #[test]
    fn test_read_tab_separated() {
        let reader = make_text_reader("ACGT\t42\n");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0.to_string(), "ACGT");
        assert_eq!(results[0].1, 42);
    }

    #[test]
    fn test_read_empty() {
        let reader = make_text_reader("");
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_read_no_trailing_newline() {
        let reader = make_text_reader("ACGT 42");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0.to_string(), "ACGT");
        assert_eq!(results[0].1, 42);
    }

    #[test]
    fn test_read_large_count() {
        let reader = make_text_reader("ACGT 18446744073709551615\n"); // u64::MAX
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results[0].1, u64::MAX);
    }

    #[test]
    fn test_read_longer_kmer() {
        let reader = make_text_reader("ACGTACGTACGTACGTACGTACGTA 99\n");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results[0].0.to_string(), "ACGTACGTACGTACGTACGTACGTA");
        assert_eq!(results[0].0.k(), 25);
        assert_eq!(results[0].1, 99);
    }

    #[test]
    fn test_invalid_count() {
        let reader = make_text_reader("ACGT notanumber\n");
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 1);
        assert!(results[0].is_err());
    }

    #[test]
    fn test_invalid_kmer() {
        let reader = make_text_reader("ACGN 42\n");
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 1);
        assert!(results[0].is_err());
    }

    #[test]
    fn test_malformed_line() {
        let reader = make_text_reader("justoneword\n");
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 1);
        assert!(results[0].is_err());
    }

    #[test]
    fn test_read_with_extra_whitespace() {
        let reader = make_text_reader("ACGT   42\n");
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results[0].1, 42);
    }
}