jellyfish-reader 0.1.0

use std::fs::File;
use std::io::BufReader;
use std::path::Path;

use memmap2::Mmap;

use crate::error::{Error, Result};
use crate::header::FileHeader;
use crate::matrix::RectangularBinaryMatrix;
use crate::mer::MerDna;

/// Random-access query interface for Jellyfish binary/sorted database files.
///
/// Uses memory-mapped I/O and matrix-guided binary search for efficient
/// k-mer lookups, matching the behavior of Jellyfish's C++ `binary_query` class.
///
/// # Examples
///
/// ```no_run
/// use jellyfish_reader::{QueryMerFile, MerDna};
///
/// let qf = QueryMerFile::open("output.jf").unwrap();
/// let mer: MerDna = "ACGTACGTACGTACGTACGTACGTA".parse().unwrap();
/// match qf.get(&mer) {
///     Some(count) => println!("{} appears {} times", mer, count),
///     None => println!("{} not found", mer),
/// }
/// ```
pub struct QueryMerFile {
    mmap: Mmap,
    data_offset: usize,
    key_len_bytes: usize,
    val_len_bytes: usize,
    record_len: usize,
    k: usize,
    matrix: RectangularBinaryMatrix,
    size_mask: u64,
    num_records: usize,
    header: FileHeader,
}

impl QueryMerFile {
    /// Open a Jellyfish database file for random-access queries.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = File::open(path.as_ref())?;

        // Parse header first to get metadata
        let mut buf_reader = BufReader::new(File::open(path.as_ref())?);
        let header = FileHeader::read(&mut buf_reader)?;

        if !header.is_binary() {
            return Err(Error::UnsupportedFormat(format!(
                "QueryMerFile requires binary/sorted format, got {:?}",
                header.format()
            )));
        }

        let key_len_bytes = header
            .key_bytes()
            .ok_or_else(|| Error::MissingField("key_len".to_string()))?;
        let val_len_bytes = header
            .data_val_len()
            .ok_or_else(|| Error::MissingField("counter_len or val_len".to_string()))?;
        let k = header
            .k()
            .ok_or_else(|| Error::MissingField("key_len".to_string()))?;
        let size = header
            .size()
            .ok_or_else(|| Error::MissingField("size".to_string()))?;

        let matrix = header.matrix(0)?;

        let record_len = key_len_bytes + val_len_bytes;
        let data_offset = header.offset();

        // Memory-map the entire file
        let mmap = unsafe { Mmap::map(&file)? };
        let file_data_len = mmap.len().saturating_sub(data_offset);
        let num_records = if record_len > 0 {
            file_data_len / record_len
        } else {
            0
        };

        Ok(Self {
            mmap,
            data_offset,
            key_len_bytes,
            val_len_bytes,
            record_len,
            k,
            matrix,
            size_mask: size - 1,
            num_records,
            header,
        })
    }

    /// Look up the count for a k-mer. Returns `None` if not found.
    pub fn get(&self, mer: &MerDna) -> Option<u64> {
        if mer.k() != self.k || self.num_records == 0 {
            return None;
        }

        // Compute expected hash position using matrix
        let hash_pos = self.matrix.times(mer.words()) & self.size_mask;

        // Try to find the k-mer using binary search around the expected position
        // Since records are sorted by hash, we can do a binary search
        self.binary_search_record(mer, hash_pos)
    }

    /// Query a k-mer given as a string. Convenience method.
    pub fn query(&self, kmer_str: &str) -> Result<Option<u64>> {
        let mer: MerDna = kmer_str.parse()?;
        Ok(self.get(&mer))
    }

    /// Get the file header.
    pub fn header(&self) -> &FileHeader {
        &self.header
    }

    /// K-mer length (number of bases).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Number of records in the file.
    pub fn num_records(&self) -> usize {
        self.num_records
    }

    /// Read the key (k-mer) at the given record index.
    fn read_key_at(&self, index: usize) -> Option<MerDna> {
        if index >= self.num_records {
            return None;
        }
        let offset = self.data_offset + index * self.record_len;
        let end = offset + self.key_len_bytes;
        if end > self.mmap.len() {
            return None;
        }
        Some(MerDna::from_bytes(&self.mmap[offset..end], self.k))
    }

    /// Read the value (count) at the given record index.
    fn read_val_at(&self, index: usize) -> Option<u64> {
        if index >= self.num_records {
            return None;
        }
        let offset = self.data_offset + index * self.record_len + self.key_len_bytes;
        let end = offset + self.val_len_bytes;
        if end > self.mmap.len() {
            return None;
        }
        let mut count = 0u64;
        for (i, &byte) in self.mmap[offset..end].iter().enumerate() {
            count |= (byte as u64) << (i * 8);
        }
        Some(count)
    }

    /// Binary search for a k-mer in the sorted records.
    fn binary_search_record(&self, mer: &MerDna, _hint_pos: u64) -> Option<u64> {
        if self.num_records == 0 {
            return None;
        }

        let mut lo = 0usize;
        let mut hi = self.num_records;

        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            let key = self.read_key_at(mid)?;
            match key.cmp(mer) {
                std::cmp::Ordering::Equal => return self.read_val_at(mid),
                std::cmp::Ordering::Less => lo = mid + 1,
                std::cmp::Ordering::Greater => hi = mid,
            }
        }

        None
    }
}

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

    /// Helper to create a minimal valid Jellyfish binary file.
    fn create_test_jf_file(k: usize, val_len: usize, records: &[(&str, u64)]) -> NamedTempFile {
        let key_len_bits = k * 2;
        let key_len_bytes = (key_len_bits + 7) / 8;

        // Sort records by k-mer binary representation
        let mut sorted_records: Vec<(MerDna, u64)> = records
            .iter()
            .map(|(s, c)| (s.parse::<MerDna>().unwrap(), *c))
            .collect();
        sorted_records.sort_by(|a, b| a.0.cmp(&b.0));

        let size = sorted_records.len().next_power_of_two().max(2);
        let json = serde_json::json!({
            "format": "binary/sorted",
            "key_len": key_len_bits,
            "val_len": val_len,
            "counter_len": val_len,
            "size": size,
            "canonical": false,
            "max_reprobe": 126
        });
        let json_str = serde_json::to_string(&json).unwrap();
        let header_len = json_str.len();

        let mut file = NamedTempFile::new().unwrap();

        // Write header length (9 digits)
        write!(file, "{:09}", header_len).unwrap();
        // Write JSON
        file.write_all(json_str.as_bytes()).unwrap();

        // Write sorted records
        for (mer, count) in &sorted_records {
            // Write key bytes
            let words = mer.words();
            let mut bytes_written = 0;
            for &word in words {
                for byte_idx in 0..8 {
                    if bytes_written >= key_len_bytes {
                        break;
                    }
                    file.write_all(&[(word >> (byte_idx * 8)) as u8]).unwrap();
                    bytes_written += 1;
                }
            }

            // Write value bytes (little-endian)
            for i in 0..val_len {
                file.write_all(&[(count >> (i * 8)) as u8]).unwrap();
            }
        }

        file.flush().unwrap();
        file
    }

    #[test]
    fn test_open_and_query_single() {
        let file = create_test_jf_file(4, 4, &[("ACGT", 42)]);
        let qf = QueryMerFile::open(file.path()).unwrap();
        assert_eq!(qf.k(), 4);
        assert_eq!(qf.num_records(), 1);

        let mer: MerDna = "ACGT".parse().unwrap();
        assert_eq!(qf.get(&mer), Some(42));
    }

    #[test]
    fn test_query_not_found() {
        let file = create_test_jf_file(4, 4, &[("ACGT", 42)]);
        let qf = QueryMerFile::open(file.path()).unwrap();

        let mer: MerDna = "TTTT".parse().unwrap();
        assert_eq!(qf.get(&mer), None);
    }

    #[test]
    fn test_query_multiple() {
        let file = create_test_jf_file(
            4,
            4,
            &[("AAAA", 10), ("ACGT", 42), ("CCCC", 7), ("TTTT", 100)],
        );
        let qf = QueryMerFile::open(file.path()).unwrap();

        assert_eq!(qf.get(&"AAAA".parse::<MerDna>().unwrap()), Some(10));
        assert_eq!(qf.get(&"ACGT".parse::<MerDna>().unwrap()), Some(42));
        assert_eq!(qf.get(&"CCCC".parse::<MerDna>().unwrap()), Some(7));
        assert_eq!(qf.get(&"TTTT".parse::<MerDna>().unwrap()), Some(100));
        assert_eq!(qf.get(&"GGGG".parse::<MerDna>().unwrap()), None);
    }

    #[test]
    fn test_query_string_convenience() {
        let file = create_test_jf_file(4, 4, &[("ACGT", 42)]);
        let qf = QueryMerFile::open(file.path()).unwrap();

        assert_eq!(qf.query("ACGT").unwrap(), Some(42));
        assert_eq!(qf.query("TTTT").unwrap(), None);
    }

    #[test]
    fn test_query_wrong_k() {
        let file = create_test_jf_file(4, 4, &[("ACGT", 42)]);
        let qf = QueryMerFile::open(file.path()).unwrap();

        // Wrong k-mer length should return None
        let mer: MerDna = "ACGTACGT".parse().unwrap(); // k=8 != 4
        assert_eq!(qf.get(&mer), None);
    }

    #[test]
    fn test_query_longer_kmer() {
        let seq = "ACGTACGTACGTACGTACGTACGTA"; // k=25
        let file = create_test_jf_file(25, 4, &[(seq, 99)]);
        let qf = QueryMerFile::open(file.path()).unwrap();

        assert_eq!(qf.get(&seq.parse::<MerDna>().unwrap()), Some(99));
    }

    #[test]
    fn test_header_access() {
        let file = create_test_jf_file(4, 4, &[("ACGT", 42)]);
        let qf = QueryMerFile::open(file.path()).unwrap();

        assert!(qf.header().is_binary());
        assert_eq!(qf.header().k(), Some(4));
    }

    #[test]
    fn test_empty_database() {
        let file = create_test_jf_file(4, 4, &[]);
        let qf = QueryMerFile::open(file.path()).unwrap();
        assert_eq!(qf.num_records(), 0);
        assert_eq!(qf.get(&"ACGT".parse::<MerDna>().unwrap()), None);
    }
}