slow5lib 0.1.0

use std::fs::File;
use std::io::{BufReader, Seek, SeekFrom};
use std::path::Path;

const BUF_SIZE: usize = 512 * 1024;

use crate::Result;
use crate::error::SlowError;
use crate::header::{Header, RecordCompression};
use crate::index::Index;
use crate::record::Record;
use crate::slow5::PrimaryColIndices;

/// Format-specific state held by `Slow5IndexedReader`.
enum IndexedReaderState {
    Blow5,
    Slow5 { col_idx: PrimaryColIndices },
}

/// Random-access reader for indexed SLOW5 and BLOW5 files.
///
/// `Send + Sync`: each `get()` call uses `read_at` (pread) which is stateless
/// and safe to call concurrently from multiple threads. Works naturally with
/// rayon parallel iterators over read ids.
///
/// ```no_run
/// # use slow5lib::Slow5IndexedReader;
/// # fn example() -> slow5lib::Result<()> {
/// // BLOW5 indexed access
/// let reader = Slow5IndexedReader::open("reads.blow5")?;
/// let rec = reader.get("50abece6-9476-46be-a86f-d48cc3fddeb6")?;
/// println!("{} samples", rec.raw_signal.len());
///
/// // SLOW5 text indexed access (same API)
/// let reader = Slow5IndexedReader::open("reads.slow5")?;
/// let ids: Vec<&str> = reader.read_ids().collect();
/// # Ok(()) }
/// ```
pub struct Slow5IndexedReader {
    file: File,
    index: Index,
    header: Header,
    state: IndexedReaderState,
}

// Safety: File is Send; read_at (pread) does not mutate shared file-descriptor state.
// This is only correct on Unix where read_at is backed by pread(2).
// On Windows, seek_read is not concurrency-safe; the Sync impl would need revision.
unsafe impl Sync for Slow5IndexedReader {}

impl Slow5IndexedReader {
    /// Open a SLOW5 or BLOW5 file and load (or build) its sidecar index.
    ///
    /// Detects the format from the path extension (`.slow5` or `.blow5`).
    /// Looks for `<path>.idx` next to the data file. If not found, scans the
    /// file sequentially to build the index, then writes it to disk so future
    /// opens are fast. Index write failures are non-fatal (a warning is printed
    /// to stderr).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// // On first open, builds and caches the index sidecar.
    /// // On subsequent opens, loads the cached index from disk.
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// println!("{} reads available", reader.read_ids().count());
    ///
    /// let reader = Slow5IndexedReader::open("reads.slow5")?;
    /// println!("{} reads available", reader.read_ids().count());
    /// # 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 extension '.{other}': expected '.blow5' or '.slow5'"
            ))),
        }
    }

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

        let idx_path = sidecar_idx_path(path)?;

        let index = if idx_path.exists() {
            Index::read_from_path(&idx_path, header.version)?
        } else {
            file.seek(SeekFrom::Start(start_rec_offset))
                .map_err(SlowError::Io)?;
            let index = {
                let mut scan = BufReader::with_capacity(BUF_SIZE, &mut file);
                Index::build_from_reader(&mut scan, header.record_compression, start_rec_offset)?
            };
            if let Err(e) = index.write_to_path(&idx_path, header.version) {
                eprintln!("slow5lib: warning: failed to write index to {idx_path:?}: {e}");
            }
            index
        };

        Ok(Self {
            file,
            index,
            header,
            state: IndexedReaderState::Blow5,
        })
    }

    fn open_slow5(path: &Path) -> Result<Self> {
        let mut buf_reader = BufReader::with_capacity(BUF_SIZE, File::open(path)?);
        let (header, col_idx) = crate::slow5::parse_header(&mut buf_reader)?;
        let start_rec_offset = buf_reader.stream_position().map_err(SlowError::Io)?;
        let mut file = buf_reader.into_inner();

        let idx_path = sidecar_idx_path(path)?;

        let index = if idx_path.exists() {
            Index::read_from_path(&idx_path, header.version)?
        } else {
            file.seek(SeekFrom::Start(start_rec_offset))
                .map_err(SlowError::Io)?;
            let index = {
                let mut scan = BufReader::with_capacity(BUF_SIZE, &mut file);
                Index::build_from_slow5_reader(&mut scan, start_rec_offset)?
            };
            if let Err(e) = index.write_to_path(&idx_path, header.version) {
                eprintln!("slow5lib: warning: failed to write index to {idx_path:?}: {e}");
            }
            index
        };

        Ok(Self {
            file,
            index,
            header,
            state: IndexedReaderState::Slow5 { col_idx },
        })
    }

    /// Returns the parsed file header.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::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
    }

    /// Iterator over all read ids present in the index.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// let ids: Vec<&str> = reader.read_ids().collect();
    /// println!("{} reads", ids.len());
    /// # Ok(()) }
    /// ```
    pub fn read_ids(&self) -> impl Iterator<Item = &str> {
        self.index.read_ids()
    }

    /// Fetch and decompress a single record by read id.
    ///
    /// Uses `read_at` (pread) internally -- no file cursor is modified.
    /// Safe to call concurrently from multiple threads.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use slow5lib::Slow5IndexedReader;
    ///
    /// # fn example() -> slow5lib::Result<()> {
    /// let reader = Slow5IndexedReader::open("reads.blow5")?;
    /// let rec = reader.get("50abece6-9476-46be-a86f-d48cc3fddeb6")?;
    /// println!("{} samples", rec.raw_signal.len());
    /// # Ok(()) }
    /// ```
    pub fn get(&self, read_id: &str) -> Result<Record> {
        let entry = self
            .index
            .get(read_id)
            .ok_or_else(|| SlowError::NotFound(read_id.to_string()))?;

        match &self.state {
            IndexedReaderState::Blow5 => {
                // entry.offset -> start of [record_size: u64 le]
                // entry.size   -> 8 + record_size
                let mut buf = vec![0u8; entry.size as usize];
                read_exact_at(&self.file, &mut buf, entry.offset)?;

                let record_bytes = &buf[8..];
                let decompressed = match self.header.record_compression {
                    RecordCompression::None => record_bytes.to_vec(),
                    RecordCompression::Zstd => {
                        crate::compression::decompress_record_zstd(record_bytes)?
                    }
                    RecordCompression::Zlib => {
                        crate::compression::decompress_record_zlib(record_bytes)?
                    }
                };

                crate::reader::parse_raw_record(
                    &decompressed,
                    self.header.signal_compression,
                    &self.header.aux_meta,
                )?
                .decompress()
            }

            IndexedReaderState::Slow5 { col_idx } => {
                // entry.offset -> first byte of the text line
                // entry.size   -> byte count of the line including the trailing '\n'
                let mut buf = vec![0u8; entry.size as usize];
                read_exact_at(&self.file, &mut buf, entry.offset)?;

                let line = std::str::from_utf8(&buf).map_err(|e| {
                    SlowError::InvalidFormat(format!("SLOW5 line is not valid UTF-8: {e}"))
                })?;
                crate::slow5::parse_record_line(line.trim_end(), col_idx, &self.header.aux_meta)?
                    .decompress()
            }
        }
    }
}

/// Build the sidecar index path: `reads.blow5` -> `reads.blow5.idx`.
fn sidecar_idx_path(path: &Path) -> Result<std::path::PathBuf> {
    let name = path
        .file_name()
        .and_then(|n| n.to_str())
        .ok_or_else(|| SlowError::InvalidFormat("path has no file name".into()))?;
    Ok(path.with_file_name(format!("{name}.idx")))
}

/// Read exactly `buf.len()` bytes from `file` starting at `offset`, without
/// modifying the file's position cursor.
fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> Result<()> {
    #[cfg(unix)]
    {
        use std::os::unix::fs::FileExt;
        file.read_exact_at(buf, offset).map_err(SlowError::Io)
    }
    #[cfg(not(unix))]
    {
        let _ = (file, buf, offset);
        Err(SlowError::InvalidFormat(
            "indexed random access is not yet supported on this platform".into(),
        ))
    }
}