//! Chunk iterator: yields Rope slices that each end on a complete sequence record.
//!
//! Each `Rope` yielded by [`SeqChunkIter`] contains one or more blocks that
//! together form a self-contained block of complete sequence records.

use crate::DEFAULT_BLOCK_SIZE;
use crate::MimeTypeGuesser;
use crate::fasta;
use crate::fastq;
use crate::xopen::xopen;

use std::io::{self, Read};

use obikrope::Rope;

/// A splitter function: given the accumulated rope, returns the absolute byte
/// offset at which to cut, or `None` if no complete-record boundary was found.
pub type Splitter = fn(&Rope) -> Option<usize>;

/// Iterator that reads from `R` in blocks and yields `Rope` chunks,
/// each ending on a complete sequence record boundary.
pub struct SeqChunkIter<R> {
    mime_type: Option<&'static str>,
    source: R,
    rope: Rope,
    block_size: usize,
    splitter: Splitter,
    eof: bool,
}

impl<R: Read> SeqChunkIter<R> {
    /// Create a new iterator.
    pub fn new(
        source: R,
        block_size: usize,
        splitter: Splitter,
        mime_type: Option<&'static str>,
    ) -> Self {
        Self {
            mime_type,
            source,
            rope: Rope::new(mime_type),
            block_size,
            splitter,
            eof: false,
        }
    }

    fn read_block(&mut self) -> io::Result<Option<Vec<u8>>> {
        let mut buf = vec![0u8; self.block_size];
        let mut filled = 0;
        loop {
            match self.source.read(&mut buf[filled..]) {
                Ok(0) => break,
                Ok(n) => {
                    filled += n;
                    if filled == self.block_size {
                        break;
                    }
                }
                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => return Err(e),
            }
        }
        if filled == 0 {
            return Ok(None);
        }
        buf.truncate(filled);
        Ok(Some(buf))
    }
}

impl<R: Read> Iterator for SeqChunkIter<R> {
    type Item = io::Result<Rope>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.eof {
                if self.rope.is_empty() {
                    return None;
                }
                return Some(Ok(std::mem::replace(
                    &mut self.rope,
                    Rope::new(self.mime_type),
                )));
            }

            match self.read_block() {
                Err(e) => return Some(Err(e)),
                Ok(None) => {
                    self.eof = true;
                    continue;
                }
                Ok(Some(block)) => {
                    self.rope.push(block);
                }
            }

            if let Some(abs_offset) = (self.splitter)(&self.rope) {
                let tail = self
                    .rope
                    .split_off(abs_offset)
                    .expect("splitter returned valid offset");
                let chunk = std::mem::replace(&mut self.rope, tail);
                return Some(Ok(chunk));
            }
        }
    }
}

/// Open `path` and iterate over its FASTA records as chunks.
pub fn read_fasta_chunks(
    path: &str,
) -> io::Result<SeqChunkIter<MimeTypeGuesser<Box<dyn Read + Send>>>> {
    let input = match xopen(path) {
        Ok(mut i) => {
            if i.mime_type() == Some("text/fasta") {
                i
            } else {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "not a FASTA file",
                ));
            }
        }
        Err(e) => return Err(e),
    };
    Ok(fasta_chunks(input))
}

/// Open `path` and iterate over its FASTQ records as chunks.
pub fn read_fastq_chunks(
    path: &str,
) -> io::Result<SeqChunkIter<MimeTypeGuesser<Box<dyn Read + Send>>>> {
    let input = match xopen(path) {
        Ok(mut i) => {
            if i.mime_type() == Some("text/fastq") {
                i
            } else {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "not a FASTQ file",
                ));
            }
        }
        Err(e) => return Err(e),
    };
    Ok(fastq_chunks(input))
}

/// Open `path` and auto-detect whether it is FASTA or FASTQ, then iterate over its records.
///
/// Returns an error if the format cannot be identified as `text/fasta` or `text/fastq`.
pub fn read_sequence_chunks(
    path: &str,
) -> io::Result<SeqChunkIter<MimeTypeGuesser<Box<dyn Read + Send>>>> {
    read_sequence_chunks_sized(path, DEFAULT_BLOCK_SIZE)
}

/// Same as [`read_sequence_chunks`] but with an explicit I/O block size.
///
/// Larger values amortise per-partition open/close overhead across more superkmers.
pub fn read_sequence_chunks_sized(
    path: &str,
    block_size: usize,
) -> io::Result<SeqChunkIter<MimeTypeGuesser<Box<dyn Read + Send>>>> {
    let input = match xopen(path) {
        Ok(mut i) => match i.mime_type() {
            Some("text/fasta") => SeqChunkIter::new(i, block_size,
                fasta::end_of_last_fasta_entry, Some("text/fasta")),
            Some("text/fastq") => SeqChunkIter::new(i, block_size,
                fastq::end_of_last_fastq_entry, Some("text/fastq")),
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "not a FASTA or FASTQ file",
                ));
            }
        },
        Err(e) => return Err(e),
    };
    Ok(input)
}

/// Create a FASTA chunk iterator over `source`.
pub fn fasta_chunks<R: Read>(source: R) -> SeqChunkIter<R> {
    SeqChunkIter::new(
        source,
        DEFAULT_BLOCK_SIZE,
        fasta::end_of_last_fasta_entry,
        Some("text/fasta"),
    )
}

/// Create a FASTQ chunk iterator over `source`.
pub fn fastq_chunks<R: Read>(source: R) -> SeqChunkIter<R> {
    SeqChunkIter::new(
        source,
        DEFAULT_BLOCK_SIZE,
        fastq::end_of_last_fastq_entry,
        Some("text/fastq"),
    )
}

#[cfg(test)]
#[path = "tests/chunk.rs"]
mod tests;
