slow5lib 0.1.0

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufReader, Read, Seek};
use std::path::Path;

const BUF_SIZE: usize = 512 * 1024;

use crate::Result;
use crate::aux::AuxMeta;
use crate::error::SlowError;
use crate::header::{Header, RecordCompression, SignalCompression};
use crate::record::{RawRecord, Record};

/// Sequential reader for SLOW5 (text) and BLOW5 (binary) files.
///
/// Not `Sync` -- owns a cursor over the file. For parallel decompression use
/// [`par_records()`](Slow5Reader::par_records) (requires the `rayon` feature).
///
/// ```no_run
/// # use slow5lib::Slow5Reader;
/// # fn example() -> slow5lib::Result<()> {
/// # let mut reader = Slow5Reader::open("reads.blow5")?;
/// for record in reader.records() { let _r = record?; }
/// # Ok(()) }
/// ```
pub struct Slow5Reader {
    header: Header,
    state: ReaderState,
}

enum ReaderState {
    Blow5 {
        inner: BufReader<File>,
        record_compression: RecordCompression,
        signal_compression: SignalCompression,
        #[allow(dead_code)]
        start_rec_offset: u64,
    },
    Slow5 {
        inner: BufReader<File>,
        col_idx: crate::slow5::PrimaryColIndices,
    },
}

impl Slow5Reader {
    /// Open a SLOW5 or BLOW5 file for sequential reading.
    ///
    /// The file format is detected from the path extension (`.blow5` or `.slow5`).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5Reader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let mut reader = Slow5Reader::open("reads.slow5")?;
    /// for record in reader.records() {
    ///     let r = record?;
    ///     println!("{}: {} samples", r.read_id, r.raw_signal.len());
    /// }
    /// # Ok(()) }
    /// ```
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref();

        let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");

        match ext {
            "blow5" => Self::open_blow5(path),
            "slow5" => Self::open_slow5(path),
            other => Err(SlowError::InvalidFormat(format!(
                "unrecognised file extension '.{other}': expected '.blow5' or '.slow5'"
            ))),
        }
    }

    fn open_blow5(path: &Path) -> Result<Self> {
        let file = File::open(path)?;
        let mut inner = BufReader::with_capacity(BUF_SIZE, file);
        let header = crate::blow5::parse_header(&mut inner)?;
        let start_rec_offset = inner.stream_position().map_err(SlowError::Io)?;
        let record_compression = header.record_compression;
        let signal_compression = header.signal_compression;

        Ok(Self {
            header,
            state: ReaderState::Blow5 {
                inner,
                record_compression,
                signal_compression,
                start_rec_offset,
            },
        })
    }

    fn open_slow5(path: &Path) -> Result<Self> {
        let file = File::open(path)?;
        let mut inner = BufReader::with_capacity(BUF_SIZE, file);
        let (header, col_idx) = crate::slow5::parse_header(&mut inner)?;
        Ok(Self {
            header,
            state: ReaderState::Slow5 { inner, col_idx },
        })
    }

    /// Returns the parsed file header.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5Reader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5Reader::open("reads.blow5")?;
    /// let h = reader.header();
    /// println!("version {}.{}.{}", h.version.0, h.version.1, h.version.2);
    /// # Ok(()) }
    /// ```
    pub fn header(&self) -> &Header {
        &self.header
    }

    /// Iterate over records without decompressing the signal.
    ///
    /// Each call to `next()` performs only I/O, reading compressed bytes off
    /// disk. Call [`RawRecord::decompress()`](crate::RawRecord::decompress) to pay the CPU cost.
    /// This split is the building block for custom batched parallel decompression.
    ///
    /// For SLOW5 text files, the signal is parsed from text on each call and
    /// stored as raw bytes; `decompress()` unpacks them with no additional
    /// decompression work.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5Reader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let mut reader = Slow5Reader::open("reads.blow5")?;
    /// for raw in reader.records_raw() {
    ///     let rec = raw?.decompress()?;
    ///     println!("{}", rec.read_id);
    /// }
    /// # Ok(()) }
    /// ```
    pub fn records_raw(&mut self) -> RecordsRaw<'_> {
        RecordsRaw {
            reader: self,
            exhausted: false,
        }
    }

    /// Iterate over fully decoded records.
    ///
    /// Equivalent to `records_raw().map(|r| r?.decompress())`. Single-threaded.
    /// For parallel decompression use [`par_records()`](Slow5Reader::par_records)
    /// (requires the `rayon` feature).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5Reader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let mut reader = Slow5Reader::open("reads.blow5")?;
    /// for record in reader.records() {
    ///     let r = record?;
    ///     let pa: Vec<f32> = r.raw_signal.iter()
    ///         .map(|&s| (s as f64 + r.offset) * r.range / r.digitisation)
    ///         .map(|v| v as f32)
    ///         .collect();
    ///     println!("{}: first pA = {:.2}", r.read_id, pa[0]);
    /// }
    /// # Ok(()) }
    /// ```
    pub fn records(&mut self) -> impl Iterator<Item = Result<Record>> + '_ {
        self.records_raw()
            .map(|r| r.and_then(RawRecord::decompress))
    }
}

#[cfg(feature = "rayon")]
impl Slow5Reader {
    /// Parallel-decompress all records using rayon's thread pool.
    ///
    /// Reads all compressed records sequentially in one I/O pass, then
    /// distributes decompression across the rayon pool with no lock contention.
    /// Memory usage is proportional to the compressed file size.
    ///
    /// For files too large to hold in memory, use [`records_raw()`](Slow5Reader::records_raw)
    /// and manage batching manually.
    ///
    /// Requires the `rayon` feature.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use slow5lib::Slow5Reader;
    /// use rayon::prelude::*;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let mut reader = Slow5Reader::open("reads.blow5")?;
    /// reader.par_records().for_each(|rec| {
    ///     let r = rec.unwrap();
    ///     println!("{}: {} samples", r.read_id, r.raw_signal.len());
    /// });
    /// # Ok(()) }
    /// ```
    pub fn par_records(&mut self) -> impl rayon::iter::ParallelIterator<Item = Result<Record>> {
        use rayon::iter::{IntoParallelIterator, ParallelIterator};
        let raw: Vec<Result<RawRecord>> = self.records_raw().collect();
        raw.into_par_iter()
            .map(|r| r.and_then(RawRecord::decompress))
    }
}

/// Iterator over raw (not yet decompressed) BLOW5 records.
///
/// Produced by `Slow5Reader::records_raw()`. Only I/O happens in `next()`;
/// signal decompression is deferred to `RawRecord::decompress()`.
pub struct RecordsRaw<'a> {
    reader: &'a mut Slow5Reader,
    exhausted: bool,
}

impl<'a> Iterator for RecordsRaw<'a> {
    type Item = Result<RawRecord>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.exhausted {
            return None;
        }
        let result = match &mut self.reader.state {
            ReaderState::Blow5 {
                inner,
                record_compression,
                signal_compression,
                ..
            } => read_next_raw_record(
                inner,
                *record_compression,
                *signal_compression,
                &self.reader.header.aux_meta,
            ),
            ReaderState::Slow5 { inner, col_idx } => {
                crate::slow5::read_next_record(inner, col_idx, &self.reader.header.aux_meta)
            }
        };
        match result {
            Ok(None) => {
                self.exhausted = true;
                None
            }
            Ok(Some(r)) => Some(Ok(r)),
            Err(e) => {
                self.exhausted = true;
                Some(Err(e))
            }
        }
    }
}

/// Read one raw record from `reader`.
///
/// Returns `Ok(None)` at the BLOW5 EOF marker, `Ok(Some(...))` for a valid record,
/// and `Err(...)` on I/O or format errors.
fn read_next_raw_record<R: Read>(
    reader: &mut R,
    record_compression: RecordCompression,
    signal_compression: SignalCompression,
    aux_meta: &AuxMeta,
) -> Result<Option<RawRecord>> {
    // Read the 8-byte record_size, but we must handle the 5-byte EOF marker.
    // Strategy: read 5 bytes first; if they match b"5WOLB" we're done.
    // Otherwise read 3 more bytes and assemble into u64 le.
    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 before record size or EOF marker".into(),
            ));
        }
        Err(e) => return Err(SlowError::Io(e)),
        Ok(()) => {}
    }

    if &prefix == crate::blow5::EOF_MARKER {
        return Ok(None);
    }

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

    let record_size_raw = u64::from_le_bytes([
        prefix[0], prefix[1], prefix[2], prefix[3], prefix[4], suffix[0], suffix[1], suffix[2],
    ]);
    // Cap before allocation -- prevents OOM panics on adversarial record_size fields.
    // 256 MiB is far larger than any real BLOW5 record (longest ONT reads compress to ~10 MiB).
    const MAX_RECORD_SIZE: u64 = 256 * 1024 * 1024;
    if record_size_raw > MAX_RECORD_SIZE {
        return Err(SlowError::InvalidFormat(format!(
            "record_size {record_size_raw} exceeds {MAX_RECORD_SIZE} byte limit"
        )));
    }
    let record_size = record_size_raw as usize;

    // Read the compressed record bytes
    let mut record_bytes = vec![0u8; record_size];
    reader.read_exact(&mut record_bytes)?;

    // Decompress the record bytes if needed
    let buf = match record_compression {
        RecordCompression::None => record_bytes,
        RecordCompression::Zstd => crate::compression::decompress_record_zstd(&record_bytes)?,
        RecordCompression::Zlib => crate::compression::decompress_record_zlib(&record_bytes)?,
    };

    parse_raw_record(&buf, signal_compression, aux_meta).map(Some)
}

/// Parse the decompressed record buffer into a `RawRecord`.
#[allow(unused_assignments)] // pos update in read_bytes! macro is flagged after the last call
pub(crate) fn parse_raw_record(
    buf: &[u8],
    signal_compression: SignalCompression,
    aux_meta: &AuxMeta,
) -> Result<RawRecord> {
    let mut pos: usize = 0;

    macro_rules! read_bytes {
        ($n:expr) => {{
            let end = pos.checked_add($n).ok_or_else(|| {
                SlowError::InvalidFormat(format!("record offset overflow at {pos} + {}", $n))
            })?;
            if end > buf.len() {
                return Err(SlowError::InvalidFormat(format!(
                    "record buffer too short at offset {pos}: need {}, have {}",
                    $n,
                    buf.len() - pos
                )));
            }
            let slice = &buf[pos..end];
            pos = end;
            slice
        }};
    }

    macro_rules! read_u16 {
        () => {{
            let b = read_bytes!(2);
            u16::from_le_bytes([b[0], b[1]])
        }};
    }

    macro_rules! read_u32 {
        () => {{
            let b = read_bytes!(4);
            u32::from_le_bytes([b[0], b[1], b[2], b[3]])
        }};
    }

    macro_rules! read_u64 {
        () => {{
            let b = read_bytes!(8);
            u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])
        }};
    }

    macro_rules! read_f64 {
        () => {{
            let b = read_bytes!(8);
            f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])
        }};
    }

    let read_id_len = read_u16!() as usize;
    let read_id_bytes = read_bytes!(read_id_len);
    let read_id = std::str::from_utf8(read_id_bytes)
        .map_err(|e| SlowError::InvalidFormat(format!("read_id is not valid UTF-8: {e}")))?
        .to_string();

    let read_group = read_u32!();
    let digitisation = read_f64!();
    let offset = read_f64!();
    let range = read_f64!();
    let sampling_rate = read_f64!();
    let len_raw_signal = read_u64!();

    // For uncompressed signal, len_raw_signal is the sample count; each sample is 2 bytes.
    // For SVB-ZD and ex-zd, len_raw_signal is the byte count of the compressed data.
    let signal_byte_count = match signal_compression {
        SignalCompression::None => (len_raw_signal as usize)
            .checked_mul(2)
            .ok_or_else(|| SlowError::InvalidFormat("len_raw_signal overflow".into()))?,
        SignalCompression::SvbZd | SignalCompression::ExZd => len_raw_signal as usize,
    };
    let signal_bytes = read_bytes!(signal_byte_count).to_vec();

    // Parse auxiliary fields from the remaining bytes
    let aux = if aux_meta.is_empty() {
        HashMap::new()
    } else {
        let mut aux_pos = pos;
        let mut map = HashMap::with_capacity(aux_meta.len());
        for (name, typ) in aux_meta.names.iter().zip(aux_meta.types.iter()) {
            let val = crate::aux::decode_binary(buf, &mut aux_pos, typ)?;
            map.insert(name.clone(), val);
        }
        map
    };

    Ok(RawRecord {
        read_id,
        read_group,
        digitisation,
        offset,
        range,
        sampling_rate,
        signal_bytes,
        signal_compression,
        aux,
    })
}