jellyfish-reader 0.1.0

use std::io::Read;

use crate::error::{Error, Result};
use crate::matrix::RectangularBinaryMatrix;

/// The format string for binary/sorted Jellyfish files.
pub const FORMAT_BINARY_SORTED: &str = "binary/sorted";

/// The format string for text/sorted Jellyfish files.
pub const FORMAT_TEXT_SORTED: &str = "text/sorted";

/// The format string for bloom counter Jellyfish files.
pub const FORMAT_BLOOM_COUNTER: &str = "bloomcounter";

/// Number of ASCII digits used for the header length prefix.
const HEADER_LENGTH_DIGITS: usize = 9;

/// Parsed Jellyfish file header.
///
/// Contains all metadata from the JSON header of a Jellyfish output file.
/// The header includes k-mer parameters, hash table configuration, and
/// provenance information.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::BufReader;
/// use jellyfish_reader::FileHeader;
///
/// let file = File::open("output.jf").unwrap();
/// let mut reader = BufReader::new(file);
/// let header = FileHeader::read(&mut reader).unwrap();
///
/// println!("k = {}", header.k().unwrap());
/// println!("format: {}", header.format());
/// ```
#[derive(Debug, Clone)]
pub struct FileHeader {
    /// The raw JSON root object.
    raw: serde_json::Value,
    /// Total header size including length prefix (byte offset to data).
    offset: usize,
}

impl FileHeader {
    /// Read and parse a Jellyfish file header from the given reader.
    ///
    /// This reads the 9-digit length prefix, then the JSON header content.
    /// After this call, the reader is positioned at the start of the data section.
    pub fn read<R: Read>(reader: &mut R) -> Result<Self> {
        // Read the 9-digit decimal length prefix
        let mut len_buf = [0u8; HEADER_LENGTH_DIGITS];
        reader
            .read_exact(&mut len_buf)
            .map_err(|_| Error::InvalidHeader("could not read header length prefix".to_string()))?;

        let len_str = std::str::from_utf8(&len_buf)
            .map_err(|_| Error::InvalidHeader("header length is not valid ASCII".to_string()))?;

        let header_len: usize = len_str
            .trim()
            .parse()
            .map_err(|_| Error::InvalidHeader(format!("invalid header length: {len_str:?}")))?;

        // Read the JSON content
        let mut json_buf = vec![0u8; header_len];
        reader
            .read_exact(&mut json_buf)
            .map_err(|_| Error::InvalidHeader("could not read header JSON content".to_string()))?;

        // Strip trailing null bytes (alignment padding)
        let json_end = json_buf
            .iter()
            .position(|&b| b == 0)
            .unwrap_or(json_buf.len());
        let json_str = std::str::from_utf8(&json_buf[..json_end])
            .map_err(|_| Error::InvalidHeader("header JSON is not valid UTF-8".to_string()))?;

        let raw: serde_json::Value = serde_json::from_str(json_str)?;

        let offset = HEADER_LENGTH_DIGITS + header_len;

        Ok(Self { raw, offset })
    }

    /// Byte offset from the start of the file to the data section.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// The file format string (e.g., "binary/sorted", "text/sorted").
    pub fn format(&self) -> &str {
        self.raw
            .get("format")
            .and_then(|v| v.as_str())
            .unwrap_or("")
    }

    /// Whether this is a binary/sorted format file.
    pub fn is_binary(&self) -> bool {
        self.format() == FORMAT_BINARY_SORTED
    }

    /// Whether this is a text/sorted format file.
    pub fn is_text(&self) -> bool {
        self.format() == FORMAT_TEXT_SORTED
    }

    /// Hash table size (number of slots).
    pub fn size(&self) -> Option<u64> {
        self.raw.get("size").and_then(|v| v.as_u64())
    }

    /// K-mer length in bits (= 2 * k). Divide by 2 to get k.
    pub fn key_len(&self) -> Option<u64> {
        self.raw.get("key_len").and_then(|v| v.as_u64())
    }

    /// K-mer length (number of bases).
    pub fn k(&self) -> Option<usize> {
        self.key_len().map(|kl| (kl / 2) as usize)
    }

    /// Counter value length in bytes.
    pub fn val_len(&self) -> Option<u64> {
        self.raw.get("val_len").and_then(|v| v.as_u64())
    }

    /// Counter length in the data section (may differ from val_len).
    pub fn counter_len(&self) -> Option<u64> {
        self.raw.get("counter_len").and_then(|v| v.as_u64())
    }

    /// Maximum reprobing distance.
    pub fn max_reprobe(&self) -> Option<u64> {
        self.raw.get("max_reprobe").and_then(|v| v.as_u64())
    }

    /// Whether canonical k-mers were used.
    pub fn canonical(&self) -> bool {
        self.raw
            .get("canonical")
            .and_then(|v| v.as_bool())
            .unwrap_or(false)
    }

    /// Number of hash functions (for bloom counter format).
    pub fn nb_hashes(&self) -> Option<u64> {
        self.raw.get("nb_hashes").and_then(|v| v.as_u64())
    }

    /// False positive rate (for bloom counter format).
    pub fn fpr(&self) -> Option<f64> {
        self.raw.get("fpr").and_then(|v| v.as_f64())
    }

    /// The hostname of the machine that created this file.
    pub fn hostname(&self) -> Option<&str> {
        self.raw.get("hostname").and_then(|v| v.as_str())
    }

    /// The working directory when the file was created.
    pub fn pwd(&self) -> Option<&str> {
        self.raw.get("pwd").and_then(|v| v.as_str())
    }

    /// The timestamp when the file was created.
    pub fn time(&self) -> Option<&str> {
        self.raw.get("time").and_then(|v| v.as_str())
    }

    /// The path to the jellyfish executable that created this file.
    pub fn exe_path(&self) -> Option<&str> {
        self.raw.get("exe_path").and_then(|v| v.as_str())
    }

    /// The command line arguments used to create this file.
    pub fn cmdline(&self) -> Option<Vec<&str>> {
        self.raw.get("cmdline").and_then(|v| {
            v.as_array()
                .map(|arr| arr.iter().filter_map(|s| s.as_str()).collect())
        })
    }

    /// Parse the hash matrix at the given index from the header.
    pub fn matrix(&self, index: usize) -> Result<RectangularBinaryMatrix> {
        let key = format!("matrix{index}");
        match self.raw.get(&key) {
            Some(v) => RectangularBinaryMatrix::from_json(v),
            None => {
                // Try without index for the default matrix
                if index == 0 {
                    match self.raw.get("matrix") {
                        Some(v) => RectangularBinaryMatrix::from_json(v),
                        None => Ok(RectangularBinaryMatrix::identity(64)),
                    }
                } else {
                    Err(Error::MissingField(key))
                }
            }
        }
    }

    /// Get the reprobes array from the header.
    pub fn reprobes(&self) -> Option<Vec<u64>> {
        self.raw.get("reprobes").and_then(|v| {
            v.as_array().map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_u64().or_else(|| v.as_i64().map(|i| i as u64)))
                    .collect()
            })
        })
    }

    /// Access the raw JSON header value.
    pub fn raw_json(&self) -> &serde_json::Value {
        &self.raw
    }

    /// Compute key length in bytes for binary format.
    pub fn key_bytes(&self) -> Option<usize> {
        self.key_len().map(|bits| {
            let bits = bits as usize;
            bits.div_ceil(8)
        })
    }

    /// Compute the counter length in bytes for the data section.
    pub fn data_val_len(&self) -> Option<usize> {
        self.counter_len()
            .or_else(|| self.val_len())
            .map(|v| v as usize)
    }
}

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

    /// Helper to create a fake Jellyfish header from a JSON value.
    fn make_header_bytes(json: &serde_json::Value) -> Vec<u8> {
        let json_str = serde_json::to_string(json).unwrap();
        let len = json_str.len();
        let mut buf = Vec::new();
        buf.extend_from_slice(format!("{:09}", len).as_bytes());
        buf.extend_from_slice(json_str.as_bytes());
        buf
    }

    #[test]
    fn test_parse_minimal_header() {
        let json = serde_json::json!({
            "format": "binary/sorted",
            "key_len": 50,
            "val_len": 4,
            "size": 1024,
            "canonical": true
        });
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();

        assert_eq!(header.format(), "binary/sorted");
        assert!(header.is_binary());
        assert!(!header.is_text());
        assert_eq!(header.key_len(), Some(50));
        assert_eq!(header.k(), Some(25));
        assert_eq!(header.val_len(), Some(4));
        assert_eq!(header.size(), Some(1024));
        assert!(header.canonical());
    }

    #[test]
    fn test_parse_text_format_header() {
        let json = serde_json::json!({
            "format": "text/sorted",
            "key_len": 42,
            "counter_len": 8
        });
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();

        assert!(header.is_text());
        assert!(!header.is_binary());
        assert_eq!(header.k(), Some(21));
    }

    #[test]
    fn test_header_offset() {
        let json = serde_json::json!({"format": "binary/sorted"});
        let json_str = serde_json::to_string(&json).unwrap();
        let expected_offset = 9 + json_str.len();

        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();

        assert_eq!(header.offset(), expected_offset);
    }

    #[test]
    fn test_header_with_padding() {
        let json = serde_json::json!({"format": "binary/sorted", "key_len": 50});
        let json_str = serde_json::to_string(&json).unwrap();
        let padded_len = json_str.len() + 10; // add null padding
        let mut buf = Vec::new();
        buf.extend_from_slice(format!("{:09}", padded_len).as_bytes());
        buf.extend_from_slice(json_str.as_bytes());
        buf.extend_from_slice(&vec![0u8; 10]); // null padding

        let mut cursor = Cursor::new(buf);
        let header = FileHeader::read(&mut cursor).unwrap();
        assert_eq!(header.format(), "binary/sorted");
        assert_eq!(header.key_len(), Some(50));
    }

    #[test]
    fn test_header_metadata_fields() {
        let json = serde_json::json!({
            "format": "binary/sorted",
            "hostname": "testhost",
            "pwd": "/tmp",
            "time": "2024-01-01",
            "exe_path": "/usr/bin/jellyfish",
            "cmdline": ["jellyfish", "count", "-m", "25", "input.fa"],
            "key_len": 50,
            "val_len": 4,
            "counter_len": 4,
            "max_reprobe": 126,
            "canonical": false
        });
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();

        assert_eq!(header.hostname(), Some("testhost"));
        assert_eq!(header.pwd(), Some("/tmp"));
        assert_eq!(header.time(), Some("2024-01-01"));
        assert_eq!(header.exe_path(), Some("/usr/bin/jellyfish"));
        assert_eq!(
            header.cmdline(),
            Some(vec!["jellyfish", "count", "-m", "25", "input.fa"])
        );
        assert_eq!(header.max_reprobe(), Some(126));
        assert!(!header.canonical());
    }

    #[test]
    fn test_header_key_bytes() {
        let json = serde_json::json!({"format": "binary/sorted", "key_len": 50});
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();
        // 50 bits -> ceil(50/8) = 7 bytes
        assert_eq!(header.key_bytes(), Some(7));
    }

    #[test]
    fn test_header_key_bytes_aligned() {
        let json = serde_json::json!({"format": "binary/sorted", "key_len": 64});
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();
        // 64 bits -> 8 bytes exactly
        assert_eq!(header.key_bytes(), Some(8));
    }

    #[test]
    fn test_invalid_header_too_short() {
        let mut cursor = Cursor::new(b"123");
        assert!(FileHeader::read(&mut cursor).is_err());
    }

    #[test]
    fn test_invalid_header_bad_length() {
        let mut cursor = Cursor::new(b"not_a_num");
        assert!(FileHeader::read(&mut cursor).is_err());
    }

    #[test]
    fn test_invalid_header_bad_json() {
        let mut buf = Vec::new();
        let bad_json = b"not json!";
        buf.extend_from_slice(format!("{:09}", bad_json.len()).as_bytes());
        buf.extend_from_slice(bad_json);
        let mut cursor = Cursor::new(buf);
        assert!(FileHeader::read(&mut cursor).is_err());
    }

    #[test]
    fn test_missing_optional_fields() {
        let json = serde_json::json!({"format": "binary/sorted"});
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();

        assert_eq!(header.size(), None);
        assert_eq!(header.key_len(), None);
        assert_eq!(header.val_len(), None);
        assert_eq!(header.hostname(), None);
        assert_eq!(header.pwd(), None);
        assert_eq!(header.time(), None);
        assert_eq!(header.exe_path(), None);
        assert_eq!(header.cmdline(), None);
        assert!(!header.canonical()); // defaults to false
    }

    #[test]
    fn test_data_val_len_prefers_counter_len() {
        let json = serde_json::json!({
            "format": "binary/sorted",
            "val_len": 4,
            "counter_len": 8
        });
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();
        assert_eq!(header.data_val_len(), Some(8)); // counter_len takes priority
    }

    #[test]
    fn test_data_val_len_fallback_to_val_len() {
        let json = serde_json::json!({
            "format": "binary/sorted",
            "val_len": 4
        });
        let bytes = make_header_bytes(&json);
        let mut cursor = Cursor::new(bytes);
        let header = FileHeader::read(&mut cursor).unwrap();
        assert_eq!(header.data_val_len(), Some(4));
    }
}