paraseq 0.4.14

use std::io;
use std::ops::RangeBounds;
use std::{borrow::Cow, thread};

use log::warn;

use crate::parallel::multi::{InterleavedMultiReader, MultiReader};
use crate::parallel::paired::{InterleavedPairedReader, PairedReader};
use crate::parallel::reader::{range_to_offset_limit, SingleReader};
use crate::parallel::single::{process_parallel_generic, process_parallel_generic_range};
use crate::ProcessError;
use crate::{fasta, fastq, Error, Record};

#[cfg(feature = "niffler")]
use crate::BoxedReader;
#[cfg(feature = "niffler")]
use std::path::Path;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Format {
    Fasta,
    Fastq,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CollectionType {
    Single,
    Paired,
    Interleaved,
    Multi { arity: usize },
    InterleavedMulti { arity: usize },
}

/// A reader over multiple `fastx::Reader`s.
pub struct Collection<R: io::Read> {
    inner: Vec<Reader<R>>,
    collection_type: CollectionType,
}
impl<R: io::Read> Collection<R> {
    pub fn new(inner: Vec<Reader<R>>, collection_type: CollectionType) -> crate::Result<Self> {
        let reader = Self {
            inner,
            collection_type,
        };
        reader.validate_arity()?;
        Ok(reader)
    }

    fn validate_arity(&self) -> crate::Result<()> {
        if self.inner.is_empty() {
            return Err(ProcessError::CollectionSizeMismatch { arity: 1, found: 0 });
        }
        match self.collection_type {
            CollectionType::Paired if !self.inner.len().is_multiple_of(2) => {
                return Err(ProcessError::CollectionSizeMismatch {
                    arity: 2,
                    found: self.inner.len(),
                });
            }
            CollectionType::Multi { arity } if !self.inner.len().is_multiple_of(arity) => {
                return Err(ProcessError::CollectionSizeMismatch {
                    arity,
                    found: self.inner.len(),
                });
            }
            _ => {}
        }
        Ok(())
    }

    /// Returns a reference to the inner vector of readers.
    ///
    /// This method is useful for inspecting the internal readers without modifying them.
    pub fn inner(&self) -> &Vec<Reader<R>> {
        &self.inner
    }

    /// Returns a mutable reference to the inner vector of readers.
    ///
    /// Note: This method is intended for advanced use cases where internal readers need to be modified directly.
    /// Use with caution, as improper modifications can lead to undefined behavior.
    pub fn inner_mut(&mut self) -> &mut Vec<Reader<R>> {
        &mut self.inner
    }

    /// Limit processing to the first `n` records per reader in the collection.
    ///
    /// When used with parallel processing, each reader will truncate batches to
    /// stay within the limit and return `false` once the limit is reached,
    /// stopping its worker threads cleanly.
    pub fn set_record_limit(&mut self, n: usize) {
        for reader in &mut self.inner {
            reader.set_record_limit(n);
        }
    }

    /// Checks if all readers have the same format.
    ///
    /// Will return `None` if the format is not unique and `Some(format)` if it is with the format.
    pub fn unique_format(&self) -> Option<Format> {
        let format = self.inner.first().map(|reader| reader.format());
        format?;
        let format = format.unwrap();
        if self.inner.iter().all(|reader| reader.format() == format) {
            Some(format)
        } else {
            None
        }
    }

    /// Returns the type of collection.
    pub fn collection_type(&self) -> CollectionType {
        self.collection_type
    }
}

#[cfg(feature = "niffler")]
impl Collection<BoxedReader> {
    pub fn from_paths<P: AsRef<Path>>(
        paths: &[P],
        collection_type: CollectionType,
    ) -> crate::Result<Self> {
        let mut inner = Vec::new();
        for path in paths {
            inner.push(Reader::from_path(path)?);
        }
        Self::new(inner, collection_type)
    }
}

impl<R: io::Read + Send> Collection<R> {
    // Generic handler for single-reader pattern
    fn handle_single_readers<T, F>(
        mut self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        scope_fn: F,
    ) -> crate::Result<()>
    where
        T: Clone + Send,
        F: Fn(Reader<R>, &mut T, usize) -> crate::Result<()> + Send + Sync,
    {
        let total_readers = self.inner.len();

        // Determine the maximum number of threads available (or provided)
        let total_threads = match total_threads {
            0 => num_cpus::get(),
            _ => num_cpus::get().min(total_threads),
        };

        // Calculate the number of threads per reader
        let threads_per_reader = match threads_per_reader {
            Some(num) => num.min(total_threads),
            None => (total_threads / total_readers).max(1),
        };

        // Find the batch size (i.e. number of readers per batch)
        let batch_size = total_threads / threads_per_reader;

        // Calculate the number of batches
        let num_batches = total_readers.div_ceil(batch_size);

        // eprintln!(
        //     "Processing {} readers; {} threads per reader; {} readers per batch; {} batches",
        //     total_readers, threads_per_reader, batch_size, num_batches
        // );

        thread::scope(|scope| -> crate::Result<()> {
            let scope_fn = &scope_fn;

            for _batch_idx in 0..num_batches {
                // Pull the readers for the batch
                let mut batch = Vec::new();
                let rbound = batch_size.min(self.inner.len());
                batch.extend(self.inner.drain(..rbound));

                // create threads for all readers in this batch
                let mut subhandles = Vec::new();
                for reader in batch {
                    let mut thread_proc = processor.clone();
                    subhandles.push(scope.spawn(move || -> crate::Result<()> {
                        scope_fn(reader, &mut thread_proc, threads_per_reader)?;
                        Ok(())
                    }));
                }

                // join all threads in this batch
                // eprintln!(
                //     "Joining threads in batch {_batch_idx}; # readers: {}",
                //     rbound,
                // );
                for handle in subhandles {
                    handle
                        .join()
                        .map_err(|_| crate::ProcessError::JoinError)??;
                }
            }

            Ok(())
        })
    }

    /// Generic handler for arity-based (grouped readers) pattern
    fn handle_grouped_readers<T, F>(
        mut self,
        processor: &mut T,
        total_threads: usize,
        threads_per_group: Option<usize>,
        arity: usize,
        scope_fn: F,
    ) -> crate::Result<()>
    where
        T: Clone + Send,
        F: Fn(Vec<Reader<R>>, &mut T, usize) -> crate::Result<()> + Send + Sync,
    {
        let total_groups = self.inner.len() / arity;

        // Determine the maximum number of threads available (or provided)
        let total_threads = match total_threads {
            0 => num_cpus::get(),
            _ => num_cpus::get().min(total_threads),
        };

        // Calculate the number of threads per group
        let threads_per_group = match threads_per_group {
            Some(num) => num.min(total_threads),
            None => {
                if total_threads >= arity {
                    (total_threads / total_groups).max(arity)
                } else {
                    (total_threads / total_groups).max(1)
                }
            }
        };

        // Find the batch size (i.e. number of groups per batch)
        let batch_size = total_threads / threads_per_group;

        // Calculate the number of batches
        let num_batches = total_groups.div_ceil(batch_size);

        // eprintln!("Total groups: {}, Total threads: {}, Threads per group: {}, Batch size: {}, Number of batches: {}", total_groups, total_threads, threads_per_group, batch_size, num_batches);

        thread::scope(|scope| -> crate::Result<()> {
            let scope_fn = &scope_fn;

            for _batch_idx in 0..num_batches {
                // Pull the groups for the batch
                let mut batch = Vec::new();
                let groups_in_batch = batch_size.min(total_groups - (_batch_idx * batch_size));

                for _ in 0..groups_in_batch {
                    let group: Vec<_> = self.inner.drain(..arity).collect();
                    batch.push(group);
                }

                // Create threads for all groups in this batch
                let mut subhandles = Vec::new();
                for group in batch {
                    let mut thread_proc = processor.clone();
                    subhandles.push(scope.spawn(move || -> crate::Result<()> {
                        scope_fn(group, &mut thread_proc, threads_per_group)?;
                        Ok(())
                    }));
                }

                // Join all threads in this batch
                // eprintln!(
                //     "Joining batch {}; number of groups: {}",
                //     _batch_idx, groups_in_batch
                // );
                for handle in subhandles {
                    handle
                        .join()
                        .map_err(|_| crate::ProcessError::JoinError)??;
                }
            }

            Ok(())
        })
    }

    /// Warn if collection type doesn't match expected processing mode
    fn warn_if_mismatch(&self, expected: CollectionType) {
        if self.collection_type == expected {
            return;
        }

        match (&self.collection_type, &expected) {
            // Exact arity matches - no warning needed
            (CollectionType::Multi { arity: a }, CollectionType::Multi { arity: b }) if a == b => {}
            (
                CollectionType::InterleavedMulti { arity: a },
                CollectionType::InterleavedMulti { arity: b },
            ) if a == b => {}

            // Multi arity=2 treated as Paired - no warning needed
            (CollectionType::Multi { arity: 2 }, CollectionType::Paired) => {}
            (CollectionType::InterleavedMulti { arity: 2 }, CollectionType::Paired) => {}
            (CollectionType::InterleavedMulti { arity: 2 }, CollectionType::Interleaved) => {}

            // Everything else gets a warning
            _ => {
                let from = match self.collection_type {
                    CollectionType::Single => "single reads".to_string(),
                    CollectionType::Paired => "paired reads".to_string(),
                    CollectionType::Interleaved => "interleaved reads".to_string(),
                    CollectionType::Multi { arity } => format!("multi reads (arity: {arity})"),
                    CollectionType::InterleavedMulti { arity } => {
                        format!("interleaved multi reads (arity: {arity})")
                    }
                };
                let to = match expected {
                    CollectionType::Single => "single reads".to_string(),
                    CollectionType::Paired => "paired reads".to_string(),
                    CollectionType::Interleaved => "interleaved reads".to_string(),
                    CollectionType::Multi { arity } => format!("multi-reads (arity={arity})"),
                    CollectionType::InterleavedMulti { arity } => {
                        format!("interleaved multi reads (arity: {arity})")
                    }
                };
                warn!("Processing {from} as {to}");
            }
        }
    }

    /// Get the effective arity for multi-read processing
    fn get_arity_for_multi(&self) -> usize {
        match self.collection_type {
            CollectionType::Single => {
                warn!("Processing single reads as multi-reads (arity=1)");
                1
            }
            CollectionType::Paired => {
                warn!("Processing paired reads as multi-reads (arity=2)");
                2
            }
            CollectionType::Interleaved => {
                warn!("Processing interleaved reads as multi-reads (arity=1)");
                1
            }
            CollectionType::Multi { arity } => arity,
            CollectionType::InterleavedMulti { arity } => {
                if arity != 2 {
                    warn!("Processing interleaved multi reads (arity: {arity}) as multi-reads (arity={arity})");
                }
                arity
            }
        }
    }

    /// Get the effective arity for interleaved multi-read processing
    fn get_arity_for_interleaved_multi(&self) -> usize {
        match self.collection_type {
            CollectionType::Single => {
                warn!("Processing single reads as interleaved multi reads (arity: 1)");
                1
            }
            CollectionType::Paired => {
                warn!("Processing paired reads as interleaved multi reads (arity: 2)");
                2
            }
            CollectionType::Interleaved => {
                warn!("Processing interleaved reads as interleaved multi reads (arity: 2)");
                2
            }
            CollectionType::Multi { arity } => {
                warn!("Processing multi reads (arity: {arity}) as interleaved multi reads");
                arity
            }
            CollectionType::InterleavedMulti { arity } => arity,
        }
    }

    pub fn process_parallel<T>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::ParallelProcessor<RefRecord<'a>>,
    {
        self.warn_if_mismatch(CollectionType::Single);
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            |reader, proc, threads| {
                process_parallel_generic(SingleReader::new(reader), proc, threads)
            },
        )
    }

    pub fn process_parallel_range<T, B>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        range: B,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::ParallelProcessor<RefRecord<'a>>,
        B: RangeBounds<usize>,
    {
        self.warn_if_mismatch(CollectionType::Single);
        let (start, limit) = range_to_offset_limit(range);
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            |reader, proc, threads| {
                process_parallel_generic_range(
                    SingleReader::new(reader),
                    proc,
                    threads,
                    start,
                    limit,
                )
            },
        )
    }

    pub fn process_parallel_paired<T>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::PairedParallelProcessor<RefRecord<'a>>,
    {
        self.warn_if_mismatch(CollectionType::Paired);
        self.handle_grouped_readers(
            processor,
            total_threads,
            threads_per_reader,
            2,
            |mut readers, proc, threads| {
                let r1 = readers.remove(0);
                let r2 = readers.remove(0);
                process_parallel_generic(PairedReader::new(r1, r2), proc, threads)
            },
        )
    }

    pub fn process_parallel_paired_range<T, B>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        range: B,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::PairedParallelProcessor<RefRecord<'a>>,
        B: RangeBounds<usize>,
    {
        self.warn_if_mismatch(CollectionType::Paired);
        let (start, limit) = range_to_offset_limit(range);
        self.handle_grouped_readers(
            processor,
            total_threads,
            threads_per_reader,
            2,
            |mut readers, proc, threads| {
                let r1 = readers.remove(0);
                let r2 = readers.remove(0);
                process_parallel_generic_range(
                    PairedReader::new(r1, r2),
                    proc,
                    threads,
                    start,
                    limit,
                )
            },
        )
    }

    pub fn process_parallel_interleaved<T>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::PairedParallelProcessor<RefRecord<'a>>,
    {
        self.warn_if_mismatch(CollectionType::Interleaved);
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            |reader, proc, threads| {
                process_parallel_generic(InterleavedPairedReader::new(reader), proc, threads)
            },
        )
    }

    pub fn process_parallel_interleaved_range<T, B>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        range: B,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::PairedParallelProcessor<RefRecord<'a>>,
        B: RangeBounds<usize>,
    {
        self.warn_if_mismatch(CollectionType::Interleaved);
        let (start, limit) = range_to_offset_limit(range);
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            |reader, proc, threads| {
                process_parallel_generic_range(
                    InterleavedPairedReader::new(reader),
                    proc,
                    threads,
                    start,
                    limit,
                )
            },
        )
    }

    pub fn process_parallel_multi<T>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::MultiParallelProcessor<RefRecord<'a>>,
        Self: Sized,
    {
        let arity = self.get_arity_for_multi();
        self.handle_grouped_readers(
            processor,
            total_threads,
            threads_per_reader,
            arity,
            |readers, proc, threads| {
                process_parallel_generic(MultiReader::new(readers), proc, threads)
            },
        )
    }

    pub fn process_parallel_multi_range<T, B>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        range: B,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::MultiParallelProcessor<RefRecord<'a>>,
        B: RangeBounds<usize>,
        Self: Sized,
    {
        let arity = self.get_arity_for_multi();
        let (start, limit) = range_to_offset_limit(range);
        self.handle_grouped_readers(
            processor,
            total_threads,
            threads_per_reader,
            arity,
            |readers, proc, threads| {
                process_parallel_generic_range(
                    MultiReader::new(readers),
                    proc,
                    threads,
                    start,
                    limit,
                )
            },
        )
    }

    pub fn process_parallel_multi_interleaved<T>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::MultiParallelProcessor<RefRecord<'a>>,
    {
        let arity = self.get_arity_for_interleaved_multi();
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            move |reader, proc, threads| {
                process_parallel_generic(InterleavedMultiReader::new(reader, arity), proc, threads)
            },
        )
    }

    pub fn process_parallel_multi_interleaved_range<T, B>(
        self,
        processor: &mut T,
        total_threads: usize,
        threads_per_reader: Option<usize>,
        range: B,
    ) -> crate::Result<()>
    where
        T: for<'a> crate::prelude::MultiParallelProcessor<RefRecord<'a>>,
        B: RangeBounds<usize>,
    {
        let arity = self.get_arity_for_interleaved_multi();
        let (start, limit) = range_to_offset_limit(range);
        self.handle_single_readers(
            processor,
            total_threads,
            threads_per_reader,
            move |reader, proc, threads| {
                process_parallel_generic_range(
                    InterleavedMultiReader::new(reader, arity),
                    proc,
                    threads,
                    start,
                    limit,
                )
            },
        )
    }
}

pub enum Reader<R: io::Read> {
    Fasta(fasta::Reader<R>),
    Fastq(fastq::Reader<R>),
}

#[cfg(feature = "niffler")]
impl Reader<BoxedReader> {
    pub fn from_path<P: AsRef<std::path::Path>>(path: P) -> Result<Self, Error> {
        let (reader, _format) = niffler::send::from_path(path)?;
        Self::new(reader)
    }

    pub fn from_stdin() -> Result<Self, Error> {
        let (reader, _format) = niffler::send::get_reader(Box::new(io::stdin()))?;
        Self::new(reader)
    }

    pub fn from_optional_path<P: AsRef<std::path::Path>>(path: Option<P>) -> Result<Self, Error> {
        match path {
            Some(path) => Self::from_path(path),
            None => Self::from_stdin(),
        }
    }

    pub fn from_path_with_batch_size<P: AsRef<std::path::Path>>(
        path: P,
        batch_size: usize,
    ) -> Result<Self, Error> {
        let (reader, _format) = niffler::send::from_path(path)?;
        Self::new_with_batch_size(reader, batch_size)
    }

    pub fn from_stdin_with_batch_size(batch_size: usize) -> Result<Self, Error> {
        let (reader, _format) = niffler::send::get_reader(Box::new(io::stdin()))?;
        Self::new_with_batch_size(reader, batch_size)
    }

    pub fn from_optional_path_with_batch_size<P: AsRef<std::path::Path>>(
        path: Option<P>,
        batch_size: usize,
    ) -> Result<Self, Error> {
        match path {
            Some(path) => Self::from_path_with_batch_size(path, batch_size),
            None => Self::from_stdin_with_batch_size(batch_size),
        }
    }
}

#[cfg(feature = "url")]
impl Reader<BoxedReader> {
    pub fn from_url(url: &str) -> Result<Self, Error> {
        let stream = reqwest::blocking::get(url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(stream))?;
        Self::new(reader)
    }

    pub fn from_url_with_batch_size(url: &str, batch_size: usize) -> Result<Self, Error> {
        let stream = reqwest::blocking::get(url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(stream))?;
        Self::new_with_batch_size(reader, batch_size)
    }
}

#[cfg(feature = "ssh")]
impl Reader<BoxedReader> {
    pub fn from_ssh(ssh_url: &str) -> Result<Self, Error> {
        let ssh_reader = crate::ssh::SshReader::new(ssh_url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(ssh_reader))?;
        Self::new(reader)
    }

    pub fn from_ssh_with_batch_size(ssh_url: &str, batch_size: usize) -> Result<Self, Error> {
        let ssh_reader = crate::ssh::SshReader::new(ssh_url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(ssh_reader))?;
        Self::new_with_batch_size(reader, batch_size)
    }
}

#[cfg(feature = "gcs")]
impl Reader<BoxedReader> {
    /// Create a GCS reader using Application Default Credentials
    pub fn from_gcs(gcs_url: &str) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::new(gcs_url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new(reader)
    }

    /// Create a GCS reader using custom gcloud arguments
    pub fn from_gcs_with_gcloud_args(gcs_url: &str, args: &[&str]) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::with_gcloud_args(gcs_url, args)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new(reader)
    }

    /// Create a GCS reader using a specific project ID
    pub fn from_gcs_with_project(gcs_url: &str, project_id: &str) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::with_project(gcs_url, project_id)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new(reader)
    }

    /// Create a GCS reader with custom batch size using Application Default Credentials
    pub fn from_gcs_with_batch_size(gcs_url: &str, batch_size: usize) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::new(gcs_url)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new_with_batch_size(reader, batch_size)
    }

    /// Create a GCS reader with custom batch size using custom gcloud arguments
    pub fn from_gcs_with_gcloud_args_and_batch_size(
        gcs_url: &str,
        gcloud_args: &[&str],
        batch_size: usize,
    ) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::with_gcloud_args(gcs_url, gcloud_args)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new_with_batch_size(reader, batch_size)
    }

    /// Create a GCS reader with custom batch size using a specific project ID
    pub fn from_gcs_with_project_and_batch_size(
        gcs_url: &str,
        project_id: &str,
        batch_size: usize,
    ) -> Result<Self, Error> {
        let gcs_reader = crate::gcs::GcsReader::with_project(gcs_url, project_id)?;
        let (reader, _format) = niffler::send::get_reader(Box::new(gcs_reader))?;
        Self::new_with_batch_size(reader, batch_size)
    }
}

impl<R: io::Read> Reader<R> {
    pub fn new(mut reader: R) -> Result<Self, Error> {
        let mut buffer = [0; 1];
        reader.read_exact(&mut buffer)?;
        match buffer {
            [b'@'] => {
                let mut rdr = fastq::Reader::new(reader);
                rdr.add_to_overflow(&buffer);
                Ok(Self::Fastq(rdr))
            }
            [b'>'] => {
                let mut rdr = fasta::Reader::new(reader);
                rdr.add_to_overflow(&buffer);
                Ok(Self::Fasta(rdr))
            }
            _ => Err(Error::InvalidStartCharacter(buffer[0].into())),
        }
    }

    pub fn new_with_batch_size(mut reader: R, batch_size: usize) -> Result<Self, Error> {
        let mut buffer = [0; 1];
        reader.read_exact(&mut buffer)?;
        match buffer {
            [b'@'] => {
                let mut rdr = fastq::Reader::with_batch_size(reader, batch_size)?;
                rdr.add_to_overflow(&buffer);
                Ok(Self::Fastq(rdr))
            }
            [b'>'] => {
                let mut rdr = fasta::Reader::with_batch_size(reader, batch_size)?;
                rdr.add_to_overflow(&buffer);
                Ok(Self::Fasta(rdr))
            }
            _ => Err(Error::InvalidStartCharacter(buffer[0].into())),
        }
    }

    /// Use the first record in the input to set the number of records per batch
    /// so that the expected length per batch is approximately `batch_size_in_bp`.
    pub fn update_batch_size_in_bp(&mut self, batch_size_in_bp: usize) -> Result<(), Error> {
        match self {
            Self::Fasta(inner) => inner.update_batch_size_in_bp(batch_size_in_bp),
            Self::Fastq(inner) => inner.update_batch_size_in_bp(batch_size_in_bp),
        }
    }

    pub fn format(&self) -> Format {
        match self {
            Self::Fasta(_) => Format::Fasta,
            Self::Fastq(_) => Format::Fastq,
        }
    }

    /// Limit processing to the first `n` records.
    ///
    /// When used with parallel processing, `fill()` will truncate batches to
    /// stay within the limit and return `false` once the limit is reached,
    /// stopping all worker threads cleanly.
    pub fn set_record_limit(&mut self, n: usize) {
        match self {
            Self::Fasta(inner) => inner.set_record_limit(n),
            Self::Fastq(inner) => inner.set_record_limit(n),
        }
    }

    pub fn new_record_set(&self) -> RecordSet {
        match self {
            Self::Fasta(inner) => RecordSet::Fasta(inner.new_record_set()),
            Self::Fastq(inner) => RecordSet::Fastq(inner.new_record_set()),
        }
    }

    pub fn new_record_set_with_size(&self, size: usize) -> RecordSet {
        match self {
            Self::Fasta(inner) => RecordSet::Fasta(inner.new_record_set_with_size(size)),
            Self::Fastq(inner) => RecordSet::Fastq(inner.new_record_set_with_size(size)),
        }
    }

    pub fn reload(&mut self, rset: &mut RecordSet) -> Result<(), Error> {
        match (self, rset) {
            (Self::Fasta(inner), RecordSet::Fasta(rset)) => {
                inner.reload(rset);
                Ok(())
            }
            (Self::Fastq(inner), RecordSet::Fastq(rset)) => {
                inner.reload(rset);
                Ok(())
            }
            _ => Err(Error::FormatMismatch),
        }
    }

    pub fn into_fasta_reader(self) -> Result<fasta::Reader<R>, Error> {
        match self {
            Self::Fasta(inner) => Ok(inner),
            _ => Err(Error::UnexpectedFormatRequest(
                "FASTQ".to_string(),
                "FASTA".to_string(),
            )),
        }
    }

    pub fn into_fastq_reader(self) -> Result<fastq::Reader<R>, Error> {
        match self {
            Self::Fastq(inner) => Ok(inner),
            _ => Err(Error::UnexpectedFormatRequest(
                "FASTA".to_string(),
                "FASTQ".to_string(),
            )),
        }
    }
}

pub enum RecordSet {
    Fasta(fasta::RecordSet),
    Fastq(fastq::RecordSet),
}
impl RecordSet {
    pub fn fill<R: io::Read>(&mut self, reader: &mut Reader<R>) -> Result<bool, Error> {
        match (self, reader) {
            (RecordSet::Fasta(records), Reader::Fasta(reader)) => records.fill(reader),
            (RecordSet::Fastq(records), Reader::Fastq(reader)) => records.fill(reader),
            _ => Err(Error::FormatMismatch),
        }
    }

    pub fn iter(&self) -> Box<dyn Iterator<Item = Result<RefRecord<'_>, Error>> + '_> {
        match self {
            RecordSet::Fasta(records) => Box::new(
                records
                    .iter()
                    .map(|record| -> Result<RefRecord<'_>, Error> {
                        Ok(RefRecord::Fasta(record?))
                    }),
            ),
            RecordSet::Fastq(records) => Box::new(
                records
                    .iter()
                    .map(|record| -> Result<RefRecord<'_>, Error> {
                        Ok(RefRecord::Fastq(record?))
                    }),
            ),
        }
    }
}

pub enum RefRecord<'a> {
    Fasta(fasta::RefRecord<'a>),
    Fastq(fastq::RefRecord<'a>),
}

impl Record for RefRecord<'_> {
    fn id(&self) -> &[u8] {
        match self {
            Self::Fasta(x) => x.id(),
            Self::Fastq(x) => x.id(),
        }
    }
    fn seq(&self) -> Cow<'_, [u8]> {
        match self {
            Self::Fasta(x) => x.seq(),
            Self::Fastq(x) => x.seq(),
        }
    }
    fn seq_raw(&self) -> &[u8] {
        match self {
            Self::Fasta(x) => x.seq_raw(),
            Self::Fastq(x) => x.seq_raw(),
        }
    }
    fn qual(&self) -> Option<&[u8]> {
        match self {
            Self::Fasta(x) => x.qual(),
            Self::Fastq(x) => x.qual(),
        }
    }
}

/// An internal trait implemented both by fasta and fastq types,
/// so we only have to write the reader and parallel IO implementations once.
pub trait GenericReader: Send {
    type RecordSet: Send + 'static;
    type Error;
    type RefRecord<'a>;

    fn new_record_set(&self) -> Self::RecordSet;
    fn fill(&mut self, record: &mut Self::RecordSet) -> std::result::Result<bool, Self::Error>;
    fn iter(
        record_set: &Self::RecordSet,
    ) -> impl ExactSizeIterator<Item = std::result::Result<Self::RefRecord<'_>, Self::Error>>;
    fn check_read_pair(
        _rec1: &Self::RefRecord<'_>,
        _rec2: &Self::RefRecord<'_>,
    ) -> std::result::Result<(), Self::Error> {
        Ok(())
    }
    fn set_threads(&mut self, _threads: usize) -> std::result::Result<(), Self::Error> {
        Ok(())
    }
}

impl<R> GenericReader for crate::fastx::Reader<R>
where
    R: io::Read + Send,
{
    type RecordSet = RecordSet;
    type Error = Error;
    type RefRecord<'a> = crate::fastx::RefRecord<'a>;

    fn new_record_set(&self) -> Self::RecordSet {
        match self {
            Reader::Fasta(inner) => RecordSet::Fasta(inner.new_record_set()),
            Reader::Fastq(inner) => RecordSet::Fastq(inner.new_record_set()),
        }
    }

    fn fill(&mut self, record: &mut Self::RecordSet) -> std::result::Result<bool, Self::Error> {
        record.fill(self)
    }

    fn iter(
        record_set: &Self::RecordSet,
    ) -> impl ExactSizeIterator<Item = std::result::Result<Self::RefRecord<'_>, Self::Error>> {
        match record_set {
            RecordSet::Fasta(record_set) => either::Either::Left(
                fasta::Reader::<R>::iter(record_set).map(|x| x.map(RefRecord::Fasta)),
            ),
            RecordSet::Fastq(record_set) => either::Either::Right(
                fastq::Reader::<R>::iter(record_set).map(|x| x.map(RefRecord::Fastq)),
            ),
        }
    }
}

#[cfg(feature = "niffler")]
#[cfg(test)]
mod testing {

    use crate::prelude::{ParallelProcessor, ParallelReader};
    use parking_lot::Mutex;
    use std::sync::Arc;

    use super::*;

    const FORMAT_EXTENSIONS: &[&str] = &[".fasta", ".fastq"];
    const COMPRESSION_EXTENSIONS: &[&str] = &["", ".gz", ".zst"];

    #[derive(Clone, Default)]
    struct Processor {
        local_count: usize,
        global_count: Arc<Mutex<usize>>,
    }
    impl Processor {
        pub fn n_records(&self) -> usize {
            *self.global_count.lock()
        }
    }
    impl<Rf: crate::Record> ParallelProcessor<Rf> for Processor {
        fn process_record(&mut self, _record: Rf) -> crate::parallel::Result<()> {
            self.local_count += 1;
            Ok(())
        }
        fn on_batch_complete(&mut self) -> crate::parallel::Result<()> {
            *self.global_count.lock() += self.local_count;
            self.local_count = 0;
            Ok(())
        }
    }

    #[test]
    fn test_fastx_reader_from_path() {
        let basename = "./data/sample";
        for format_ext in FORMAT_EXTENSIONS {
            for compression_ext in COMPRESSION_EXTENSIONS {
                let path = format!("{}{}{}", basename, format_ext, compression_ext);
                dbg!(&path);
                let reader = Reader::from_path(path).unwrap();
                let mut proc = Processor::default();
                reader.process_parallel(&mut proc, 1).unwrap();
                assert_eq!(proc.n_records(), 100);
            }
        }
    }

    #[test]
    fn test_fastx_reader_from_path_with_batch_size() {
        let basename = "./data/sample";
        for format_ext in FORMAT_EXTENSIONS {
            for compression_ext in COMPRESSION_EXTENSIONS {
                let path = format!("{}{}{}", basename, format_ext, compression_ext);
                dbg!(&path);
                let reader = Reader::from_path_with_batch_size(path, 10).unwrap();
                let mut proc = Processor::default();
                reader.process_parallel(&mut proc, 1).unwrap();
                assert_eq!(proc.n_records(), 100);
            }
        }
    }

    #[test]
    fn test_fastx_reload() {
        let path = "./data/sample.fastq";
        let mut reader = Reader::from_path(path).unwrap();
        let mut rset = reader.new_record_set_with_size(7);

        assert!(rset.fill(&mut reader).unwrap());
        let num_prefill = rset.iter().map(Result::unwrap).count();
        assert_eq!(num_prefill, 7);

        reader.reload(&mut rset).unwrap();

        let mut proc = Processor::default();
        reader.process_parallel(&mut proc, 1).unwrap();
        assert_eq!(proc.n_records(), 100);
    }

    #[derive(Clone, Default)]
    struct WriteProcessor {
        out_format: FormatKind,
        local_buf: Vec<u8>,
        global_buf: Arc<Mutex<Vec<u8>>>,
    }
    #[derive(Clone, Copy, Default)]
    enum FormatKind {
        #[default]
        Fasta,
        Fastq,
    }
    impl<Rf: crate::Record> ParallelProcessor<Rf> for WriteProcessor {
        fn process_record(&mut self, record: Rf) -> crate::parallel::Result<()> {
            match self.out_format {
                FormatKind::Fasta => record.write_fasta(&mut self.local_buf)?,
                FormatKind::Fastq => record.write_fastq(&mut self.local_buf)?,
            }
            Ok(())
        }
        fn on_batch_complete(&mut self) -> crate::parallel::Result<()> {
            self.global_buf.lock().extend_from_slice(&self.local_buf);
            self.local_buf.clear();
            Ok(())
        }
    }

    #[test]
    fn test_read_write_roundtrip() {
        let basename = "./data/sample";
        for format_ext in FORMAT_EXTENSIONS {
            for compression_ext in COMPRESSION_EXTENSIONS {
                for out_format in [FormatKind::Fasta, FormatKind::Fastq] {
                    let path = format!("{}{}{}", basename, format_ext, compression_ext);
                    dbg!(&path, out_format as u8);

                    let reader = Reader::from_path(&path).unwrap();
                    let mut writer = WriteProcessor {
                        out_format,
                        ..Default::default()
                    };
                    reader.process_parallel(&mut writer, 1).unwrap();
                    let written = writer.global_buf.lock().clone();

                    let reparsed = Reader::new(std::io::Cursor::new(written)).unwrap();
                    let mut proc = Processor::default();
                    reparsed.process_parallel(&mut proc, 1).unwrap();
                    assert_eq!(proc.n_records(), 100);
                }
            }
        }
    }

    #[derive(Clone, Default)]
    struct PairProcessor {
        local_count: usize,
        global_count: Arc<Mutex<usize>>,
    }
    impl PairProcessor {
        fn n_pairs(&self) -> usize {
            *self.global_count.lock()
        }
    }
    impl<Rf: crate::Record> crate::prelude::PairedParallelProcessor<Rf> for PairProcessor {
        fn process_record_pair(&mut self, _r1: Rf, _r2: Rf) -> crate::parallel::Result<()> {
            self.local_count += 1;
            Ok(())
        }
        fn on_batch_complete(&mut self) -> crate::parallel::Result<()> {
            *self.global_count.lock() += self.local_count;
            self.local_count = 0;
            Ok(())
        }
    }
    impl<Rf: crate::Record> crate::prelude::MultiParallelProcessor<Rf> for PairProcessor {
        fn process_multi_record(&mut self, _records: &[Rf]) -> crate::parallel::Result<()> {
            self.local_count += 1;
            Ok(())
        }
        fn on_batch_complete(&mut self) -> crate::parallel::Result<()> {
            *self.global_count.lock() += self.local_count;
            self.local_count = 0;
            Ok(())
        }
    }

    #[test]
    fn test_collection_single() {
        let collection =
            Collection::from_paths(&["./data/sample.fastq"], CollectionType::Single).unwrap();
        let mut proc = Processor::default();
        collection.process_parallel(&mut proc, 1, None).unwrap();
        assert_eq!(proc.n_records(), 100);
    }

    #[test]
    fn test_collection_paired() {
        let collection = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq"],
            CollectionType::Paired,
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_paired(&mut proc, 1, None)
            .unwrap();
        assert_eq!(proc.n_pairs(), 100);
    }

    #[test]
    fn test_collection_range() {
        let collection =
            Collection::from_paths(&["./data/sample.fastq"], CollectionType::Single).unwrap();
        let mut proc = Processor::default();
        collection
            .process_parallel_range(&mut proc, 1, None, 10..50)
            .unwrap();
        assert_eq!(proc.n_records(), 40);
    }

    #[test]
    fn test_collection_paired_range() {
        let collection = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq"],
            CollectionType::Paired,
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_paired_range(&mut proc, 1, None, 10..50)
            .unwrap();
        assert_eq!(proc.n_pairs(), 40);
    }

    #[test]
    fn test_collection_interleaved() {
        let collection =
            Collection::from_paths(&["./data/interleaved.fastq"], CollectionType::Interleaved)
                .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_interleaved(&mut proc, 1, None)
            .unwrap();
        assert_eq!(proc.n_pairs(), 100);
    }

    #[test]
    fn test_collection_interleaved_range() {
        let collection =
            Collection::from_paths(&["./data/interleaved.fastq"], CollectionType::Interleaved)
                .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_interleaved_range(&mut proc, 1, None, 10..50)
            .unwrap();
        assert_eq!(proc.n_pairs(), 40);
    }

    #[test]
    fn test_collection_multi() {
        let collection = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq", "./data/r3.fastq"],
            CollectionType::Multi { arity: 3 },
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_multi(&mut proc, 1, None)
            .unwrap();
        assert_eq!(proc.n_pairs(), 100);
    }

    #[test]
    fn test_collection_multi_range() {
        let collection = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq", "./data/r3.fastq"],
            CollectionType::Multi { arity: 3 },
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_multi_range(&mut proc, 1, None, 10..50)
            .unwrap();
        assert_eq!(proc.n_pairs(), 40);
    }

    #[test]
    fn test_collection_multi_interleaved() {
        let collection = Collection::from_paths(
            &["./data/r123.fastq"],
            CollectionType::InterleavedMulti { arity: 3 },
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_multi_interleaved(&mut proc, 1, None)
            .unwrap();
        assert_eq!(proc.n_pairs(), 100);
    }

    #[test]
    fn test_collection_multi_interleaved_range() {
        let collection = Collection::from_paths(
            &["./data/r123.fastq"],
            CollectionType::InterleavedMulti { arity: 3 },
        )
        .unwrap();
        let mut proc = PairProcessor::default();
        collection
            .process_parallel_multi_interleaved_range(&mut proc, 1, None, 10..50)
            .unwrap();
        assert_eq!(proc.n_pairs(), 40);
    }

    #[test]
    fn test_format() {
        let fasta_reader = Reader::from_path("./data/sample.fasta").unwrap();
        assert_eq!(fasta_reader.format(), Format::Fasta);

        let fastq_reader = Reader::from_path("./data/sample.fastq").unwrap();
        assert_eq!(fastq_reader.format(), Format::Fastq);
    }

    #[test]
    fn test_into_fasta_reader_and_into_fastq_reader() {
        let fasta_reader = Reader::from_path("./data/sample.fasta").unwrap();
        assert!(fasta_reader.into_fasta_reader().is_ok());

        let fasta_reader = Reader::from_path("./data/sample.fasta").unwrap();
        assert!(fasta_reader.into_fastq_reader().is_err());

        let fastq_reader = Reader::from_path("./data/sample.fastq").unwrap();
        assert!(fastq_reader.into_fastq_reader().is_ok());

        let fastq_reader = Reader::from_path("./data/sample.fastq").unwrap();
        assert!(fastq_reader.into_fasta_reader().is_err());
    }

    #[test]
    fn test_update_batch_size_in_bp() {
        for path in ["./data/sample.fasta", "./data/sample.fastq"] {
            let mut reader = Reader::from_path(path).unwrap();
            reader.update_batch_size_in_bp(1000).unwrap();

            let mut proc = Processor::default();
            reader.process_parallel(&mut proc, 1).unwrap();
            assert_eq!(proc.n_records(), 100);
        }
    }

    #[test]
    fn test_from_stdin() {
        if crate::test_util::is_stdin_child() {
            let reader = Reader::from_optional_path(None::<&str>).unwrap();
            let mut proc = Processor::default();
            reader.process_parallel(&mut proc, 1).unwrap();
            eprintln!("STDIN_COUNT={}", proc.n_records());
            return;
        }

        let stdin_data: String = (0..20)
            .map(|i| format!("@seq{i}\nACGT\n+\nIIII\n"))
            .collect();
        let output = crate::test_util::run_with_piped_stdin(
            "fastx::testing::test_from_stdin",
            stdin_data.as_bytes(),
        );
        assert!(output.status.success(), "child failed: {output:?}");
        let stderr = String::from_utf8_lossy(&output.stderr);
        assert!(stderr.contains("STDIN_COUNT=20"), "stderr: {stderr}");
    }

    #[test]
    fn test_ref_record_seq_raw() {
        let mut reader = Reader::from_path("./data/sample.fastq").unwrap();
        let mut rset = reader.new_record_set();
        assert!(rset.fill(&mut reader).unwrap());
        for record in rset.iter() {
            let record = record.unwrap();
            assert!(!record.seq_raw().is_empty());
        }
    }

    #[test]
    fn test_collection_empty_errors() {
        let err = Collection::from_paths(&[] as &[&str], CollectionType::Single)
            .err()
            .unwrap();
        assert!(err.to_string().contains("Collection size mismatch"));
    }

    #[test]
    fn test_collection_paired_arity_mismatch_errors() {
        let err = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq", "./data/r3.fastq"],
            CollectionType::Paired,
        )
        .err()
        .unwrap();
        assert!(err.to_string().contains("Collection size mismatch"));
    }

    #[test]
    fn test_collection_multi_arity_mismatch_errors() {
        let err = Collection::from_paths(
            &[
                "./data/r1.fastq",
                "./data/r2.fastq",
                "./data/r3.fastq",
                "./data/r4.fastq",
            ],
            CollectionType::Multi { arity: 3 },
        )
        .err()
        .unwrap();
        assert!(err.to_string().contains("Collection size mismatch"));
    }

    #[test]
    fn test_collection_unique_format() {
        let same_format = Collection::from_paths(
            &["./data/r1.fastq", "./data/r2.fastq"],
            CollectionType::Paired,
        )
        .unwrap();
        assert_eq!(same_format.unique_format(), Some(Format::Fastq));

        let mixed_format = Collection::from_paths(
            &["./data/sample.fasta", "./data/sample.fastq"],
            CollectionType::Paired,
        )
        .unwrap();
        assert_eq!(mixed_format.unique_format(), None);
    }

    #[test]
    fn test_collection_accessors() {
        let mut collection =
            Collection::from_paths(&["./data/sample.fastq"], CollectionType::Single).unwrap();
        assert_eq!(collection.collection_type(), CollectionType::Single);
        assert_eq!(collection.inner().len(), 1);

        collection.set_record_limit(10);
        collection.inner_mut().clear();
        assert!(collection.inner().is_empty());
    }
}