slow5lib 0.1.0

use std::collections::HashMap;
use std::io::{BufRead, Read, Write};
use std::path::Path;

use flate2::read::ZlibDecoder;

use crate::Result;
use crate::error::SlowError;
use crate::header::RecordCompression;

/// Binary index file format constants.
///
/// Wire layout of the `.blow5.idx` file header:
///
/// ```text
/// Offset  Size  Type      Description
///      0     9  [u8; 9]   Magic: b"SLOW5IDX\x01"
///      9     3  [u8; 3]   Version: major, minor, patch
///  12-63     -  zeros     Padding to offset 64
///     64+    -  records   Repeated: [read_id_len: u16 le][read_id: N bytes][offset: u64 le][size: u64 le]
///    end     8  [u8; 8]   EOF marker: b"XDI5WOLS"
/// ```
///
/// `offset` is the byte offset of the record's `record_size` field in the data file.
/// `size` = 8 + `record_size` (covers the 8-byte length prefix and the compressed data).
const INDEX_MAGIC: &[u8; 9] = b"SLOW5IDX\x01";
const INDEX_EOF: &[u8; 8] = b"XDI5WOLS";
const INDEX_HEADER_SIZE_OFFSET: usize = 64;

/// In-memory index over a `.slow5.idx` or `.blow5.idx` file.
///
/// Maps each `read_id` to the byte offset and length of its record in the
/// data file, enabling random access via [`Slow5IndexedReader::get()`](crate::Slow5IndexedReader::get).
/// The index is loaded or built automatically by `Slow5IndexedReader::open()`.
#[derive(Debug)]
pub struct Index {
    pub(crate) entries: HashMap<String, IndexEntry>,
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct IndexEntry {
    /// Byte offset of the `record_size` field in the data file.
    pub offset: u64,
    /// Total byte count: 8 (record_size field) + record_size (compressed data).
    pub size: u64,
}

impl Index {
    fn new() -> Self {
        Self {
            entries: HashMap::new(),
        }
    }

    /// Iterator over all read ids in the index.
    ///
    /// Iteration order is not defined (backed by a `HashMap`).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// for id in reader.read_ids() {
    ///     println!("{id}");
    /// }
    /// # Ok(()) }
    /// ```
    pub fn read_ids(&self) -> impl Iterator<Item = &str> {
        self.entries.keys().map(String::as_str)
    }

    pub(crate) fn get(&self, read_id: &str) -> Option<IndexEntry> {
        self.entries.get(read_id).copied()
    }

    /// Returns the number of reads in the index.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// println!("{} reads indexed", reader.read_ids().count());
    /// # Ok(()) }
    /// ```
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Returns `true` if the index contains no reads.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// assert!(!reader.read_ids().next().is_none());
    /// # Ok(()) }
    /// ```
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Build an index by scanning a BLOW5 file sequentially.
    ///
    /// `reader` must be positioned at the first record (after the file header).
    /// `start_offset` is the corresponding byte offset in the file.
    ///
    /// For compressed records, only the first few bytes are decompressed (enough
    /// to extract the read_id). The signal is never decoded during index building.
    pub(crate) fn build_from_reader<R: Read>(
        reader: &mut R,
        record_compression: RecordCompression,
        mut pos: u64,
    ) -> Result<Self> {
        let mut index = Self::new();
        let mut buf = Vec::new(); // reused across records to avoid per-record allocation

        loop {
            let rec_offset = pos;

            // Read first 5 bytes -- might be the start of the EOF marker
            let mut prefix = [0u8; 5];
            match reader.read_exact(&mut prefix) {
                Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
                    return Err(SlowError::InvalidFormat(
                        "unexpected EOF while building index".into(),
                    ));
                }
                Err(e) => return Err(SlowError::Io(e)),
                Ok(()) => {}
            }

            if &prefix == crate::blow5::EOF_MARKER {
                break;
            }

            let mut suffix = [0u8; 3];
            reader.read_exact(&mut suffix)?;

            let record_size = u64::from_le_bytes([
                prefix[0], prefix[1], prefix[2], prefix[3], prefix[4], suffix[0], suffix[1],
                suffix[2],
            ]);

            let total_size = 8 + record_size;
            pos += total_size;

            // Read the full compressed record (must advance the file position regardless).
            buf.resize(record_size as usize, 0);
            reader.read_exact(&mut buf)?;

            // Extract read_id with early-stop decompression -- the signal is never decoded.
            let read_id = read_id_from_record(&buf, record_compression)?;

            index.entries.insert(
                read_id,
                IndexEntry {
                    offset: rec_offset,
                    size: total_size,
                },
            );
        }

        Ok(index)
    }

    /// Build an index by scanning a SLOW5 text file sequentially.
    ///
    /// `reader` must be positioned at the first data line (after the file header).
    /// `start_offset` is the corresponding byte offset in the file.
    ///
    /// Each data line's byte range is stored as `IndexEntry { offset: line_start, size: line_len }`.
    /// The stored `size` includes the terminating `\n` exactly as `read_line` returns it.
    pub(crate) fn build_from_slow5_reader<R: BufRead>(
        reader: &mut R,
        mut pos: u64,
    ) -> Result<Self> {
        let mut index = Self::new();
        let mut line = String::new();

        loop {
            line.clear();
            let n = reader.read_line(&mut line).map_err(SlowError::Io)?;
            if n == 0 {
                break;
            }

            let line_start = pos;
            let line_len = n as u64;
            pos += line_len;

            let t = line.trim_end();
            if t.is_empty() || t.starts_with('#') || t.starts_with('@') {
                continue;
            }

            let read_id = t
                .split('\t')
                .next()
                .ok_or_else(|| SlowError::InvalidFormat("empty data line in SLOW5 file".into()))?
                .to_string();

            index.entries.insert(
                read_id,
                IndexEntry {
                    offset: line_start,
                    size: line_len,
                },
            );
        }

        Ok(index)
    }

    /// Read a `.blow5.idx` binary index file.
    ///
    /// Fails if the version embedded in the index does not match `expected_version`.
    pub(crate) fn read_from_path(path: &Path, expected_version: (u8, u8, u8)) -> Result<Self> {
        let mut file = std::fs::File::open(path)?;
        let mut buf = Vec::new();
        file.read_to_end(&mut buf)?;

        if buf.len() < INDEX_MAGIC.len() || &buf[..INDEX_MAGIC.len()] != INDEX_MAGIC.as_slice() {
            return Err(SlowError::Index("not a valid BLOW5 index file".into()));
        }

        let ver_start = INDEX_MAGIC.len();
        if buf.len() < ver_start + 3 {
            return Err(SlowError::Index("index file too short for version".into()));
        }
        let idx_version = (buf[ver_start], buf[ver_start + 1], buf[ver_start + 2]);
        if idx_version != expected_version {
            return Err(SlowError::Index(format!(
                "index version {}.{}.{} does not match file version {}.{}.{}; re-index required",
                idx_version.0,
                idx_version.1,
                idx_version.2,
                expected_version.0,
                expected_version.1,
                expected_version.2,
            )));
        }

        let mut pos = INDEX_HEADER_SIZE_OFFSET;
        let mut index = Self::new();

        loop {
            if pos >= buf.len() {
                return Err(SlowError::Index("index missing EOF marker".into()));
            }
            if buf.len() - pos >= INDEX_EOF.len()
                && &buf[pos..pos + INDEX_EOF.len()] == INDEX_EOF.as_slice()
            {
                break;
            }
            if pos + 2 > buf.len() {
                return Err(SlowError::Index(
                    "index truncated before read_id_len".into(),
                ));
            }
            let read_id_len = u16::from_le_bytes([buf[pos], buf[pos + 1]]) as usize;
            pos += 2;

            if pos + read_id_len > buf.len() {
                return Err(SlowError::Index("index truncated in read_id".into()));
            }
            let read_id = std::str::from_utf8(&buf[pos..pos + read_id_len])
                .map_err(|e| SlowError::Index(format!("read_id not valid UTF-8 in index: {e}")))?
                .to_string();
            pos += read_id_len;

            if pos + 16 > buf.len() {
                return Err(SlowError::Index("index truncated in offset/size".into()));
            }
            let offset = u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap());
            let size = u64::from_le_bytes(buf[pos + 8..pos + 16].try_into().unwrap());
            pos += 16;

            index.entries.insert(read_id, IndexEntry { offset, size });
        }

        Ok(index)
    }

    /// Write this index to a `.blow5.idx` binary file.
    pub(crate) fn write_to_path(&self, path: &Path, version: (u8, u8, u8)) -> Result<()> {
        let mut file = std::fs::File::create(path)?;

        file.write_all(INDEX_MAGIC)?;
        file.write_all(&[version.0, version.1, version.2])?;

        let padding_len = INDEX_HEADER_SIZE_OFFSET - INDEX_MAGIC.len() - 3;
        file.write_all(&vec![0u8; padding_len])?;

        for (read_id, entry) in &self.entries {
            let id_bytes = read_id.as_bytes();
            let id_len = id_bytes.len() as u16;
            file.write_all(&id_len.to_le_bytes())?;
            file.write_all(id_bytes)?;
            file.write_all(&entry.offset.to_le_bytes())?;
            file.write_all(&entry.size.to_le_bytes())?;
        }

        file.write_all(INDEX_EOF)?;

        Ok(())
    }
}

// ── Index build helpers ───────────────────────────────────────────────────────

/// Extract the read_id from a (possibly compressed) record buffer.
///
/// For compressed formats, decompression stops as soon as the read_id bytes
/// are available -- the signal is never decoded.
fn read_id_from_record(record: &[u8], compression: RecordCompression) -> Result<String> {
    match compression {
        RecordCompression::None => read_id_from_bytes(record),
        RecordCompression::Zstd => {
            let decoder =
                zstd::Decoder::new(record).map_err(|e| SlowError::Decompression(e.to_string()))?;
            read_id_from_stream(decoder)
        }
        RecordCompression::Zlib => read_id_from_stream(ZlibDecoder::new(record)),
    }
}

/// Read the read_id directly from an uncompressed byte slice.
fn read_id_from_bytes(buf: &[u8]) -> Result<String> {
    if buf.len() < 2 {
        return Err(SlowError::InvalidFormat(
            "record too short for read_id_len during index build".into(),
        ));
    }
    let n = u16::from_le_bytes([buf[0], buf[1]]) as usize;
    if buf.len() < 2 + n {
        return Err(SlowError::InvalidFormat(
            "record too short for read_id during index build".into(),
        ));
    }
    std::str::from_utf8(&buf[2..2 + n])
        .map_err(|e| SlowError::InvalidFormat(format!("read_id not valid UTF-8: {e}")))
        .map(str::to_string)
}

/// Read the read_id from a streaming decompressor, stopping immediately after.
///
/// Dropping the reader here terminates decompression -- the signal bytes are
/// never processed.
fn read_id_from_stream<R: Read>(mut r: R) -> Result<String> {
    let mut len_buf = [0u8; 2];
    r.read_exact(&mut len_buf)
        .map_err(|e| SlowError::InvalidFormat(format!("read_id_len read failed: {e}")))?;
    let n = u16::from_le_bytes(len_buf) as usize;

    let mut id_buf = vec![0u8; n];
    r.read_exact(&mut id_buf)
        .map_err(|e| SlowError::InvalidFormat(format!("read_id read failed: {e}")))?;

    std::str::from_utf8(&id_buf)
        .map_err(|e| SlowError::InvalidFormat(format!("read_id not valid UTF-8: {e}")))
        .map(str::to_string)
}