dryice-bio 0.3.0

//! Writer for the `dryice` format.

use std::{io::Write, marker::PhantomData};

use crate::{
    block::{
        BlockBuilder, BlockBuilderConfig,
        name::{NameCodec, OmittedNameCodec, RawNameCodec, SplitNameCodec},
        quality::{BinnedQualityCodec, OmittedQualityCodec, QualityCodec, RawQualityCodec},
        sequence::{
            OmittedSequenceCodec, RawAsciiCodec, SequenceCodec, TwoBitExactCodec, TwoBitLossyNCodec,
        },
    },
    config::{BlockLayoutOptions, BlockSizePolicy, DryIceWriterOptions},
    error::DryIceError,
    format,
    key::{Bytes8Key, Bytes16Key, Minimizer64, NoRecordKey, PrefixKmer64, RecordKey},
    record::{EMPTY_RECORD, SeqRecordLike},
};

/// Private marker type used to track a missing writer target in the builder.
pub struct MissingInner;

/// Builder for [`DryIceWriter`].
pub struct DryIceWriterBuilder<
    W = MissingInner,
    S = RawAsciiCodec,
    Q = RawQualityCodec,
    N = RawNameCodec,
    K = NoRecordKey,
> {
    inner: W,
    target_block_records: usize,
    _codec: PhantomData<S>,
    _quality: PhantomData<Q>,
    _name: PhantomData<N>,
    _key: PhantomData<K>,
}

impl DryIceWriterBuilder<MissingInner, RawAsciiCodec, RawQualityCodec, RawNameCodec, NoRecordKey> {
    fn new() -> Self {
        Self {
            inner: MissingInner,
            target_block_records: 8192,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }
}

impl<W, S, Q, N, K> DryIceWriterBuilder<W, S, Q, N, K> {
    /// Set the block size policy in records.
    #[must_use]
    pub fn target_block_records(mut self, n: usize) -> Self {
        self.target_block_records = n;
        self
    }
}

impl<S, Q, N, K> DryIceWriterBuilder<MissingInner, S, Q, N, K> {
    /// Set the writer's output target.
    #[must_use]
    pub fn inner<W>(self, inner: W) -> DryIceWriterBuilder<W, S, Q, N, K> {
        DryIceWriterBuilder {
            inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }
}

impl<W, Q, N, K> DryIceWriterBuilder<W, RawAsciiCodec, Q, N, K> {
    /// Configure the writer to use a user-defined sequence codec.
    #[must_use]
    pub fn sequence_codec<S: SequenceCodec>(self) -> DryIceWriterBuilder<W, S, Q, N, K> {
        DryIceWriterBuilder {
            inner: self.inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }

    /// Configure the writer to use the built-in 2-bit exact codec.
    #[must_use]
    pub fn two_bit_exact(self) -> DryIceWriterBuilder<W, TwoBitExactCodec, Q, N, K> {
        self.sequence_codec::<TwoBitExactCodec>()
    }

    /// Configure the writer to use the built-in lossy 2-bit codec
    /// that collapses all ambiguous bases to N.
    #[must_use]
    pub fn two_bit_lossy_n(self) -> DryIceWriterBuilder<W, TwoBitLossyNCodec, Q, N, K> {
        self.sequence_codec::<TwoBitLossyNCodec>()
    }

    /// Configure the writer to omit sequence data entirely.
    ///
    /// Records written through this builder must still satisfy the normal
    /// sequence/quality length invariant, which for omitted sequence payloads
    /// means writing empty sequence and quality fields.
    #[must_use]
    pub fn omit_sequence(self) -> DryIceWriterBuilder<W, OmittedSequenceCodec, Q, N, K> {
        self.sequence_codec::<OmittedSequenceCodec>()
    }
}

impl<W, K> DryIceWriterBuilder<W, RawAsciiCodec, RawQualityCodec, RawNameCodec, K> {
    /// Configure the writer to omit names, sequences, and qualities.
    ///
    /// This is convenience sugar over the three orthogonal omission controls:
    /// `omit_sequence()`, `omit_quality()`, and `omit_names()`.
    #[must_use]
    pub fn empty_payload(
        self,
    ) -> DryIceWriterBuilder<W, OmittedSequenceCodec, OmittedQualityCodec, OmittedNameCodec, K>
    {
        DryIceWriterBuilder {
            inner: self.inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }
}

impl<W, S, N, K> DryIceWriterBuilder<W, S, RawQualityCodec, N, K> {
    /// Configure the writer to use a user-defined quality codec.
    #[must_use]
    pub fn quality_codec<Q: QualityCodec>(self) -> DryIceWriterBuilder<W, S, Q, N, K> {
        DryIceWriterBuilder {
            inner: self.inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }

    /// Configure the writer to use the built-in binned quality codec.
    #[must_use]
    pub fn binned_quality(self) -> DryIceWriterBuilder<W, S, BinnedQualityCodec, N, K> {
        self.quality_codec::<BinnedQualityCodec>()
    }

    /// Configure the writer to omit quality scores entirely.
    #[must_use]
    pub fn omit_quality(self) -> DryIceWriterBuilder<W, S, OmittedQualityCodec, N, K> {
        self.quality_codec::<OmittedQualityCodec>()
    }
}

impl<W, S, Q, K> DryIceWriterBuilder<W, S, Q, RawNameCodec, K> {
    /// Configure the writer to use a user-defined name codec.
    #[must_use]
    pub fn name_codec<N: NameCodec>(self) -> DryIceWriterBuilder<W, S, Q, N, K> {
        DryIceWriterBuilder {
            inner: self.inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }

    /// Configure the writer to omit names entirely.
    #[must_use]
    pub fn omit_names(self) -> DryIceWriterBuilder<W, S, Q, OmittedNameCodec, K> {
        self.name_codec::<OmittedNameCodec>()
    }

    /// Configure the writer to split names on the first space.
    #[must_use]
    pub fn split_names(self) -> DryIceWriterBuilder<W, S, Q, SplitNameCodec, K> {
        self.name_codec::<SplitNameCodec>()
    }
}

impl<W, S, Q, N> DryIceWriterBuilder<W, S, Q, N, NoRecordKey> {
    /// Configure the writer to store a user-defined record-key type.
    #[must_use]
    pub fn record_key<K: RecordKey>(self) -> DryIceWriterBuilder<W, S, Q, N, K> {
        DryIceWriterBuilder {
            inner: self.inner,
            target_block_records: self.target_block_records,
            _codec: PhantomData,
            _quality: PhantomData,
            _name: PhantomData,
            _key: PhantomData,
        }
    }

    /// Configure the writer to store the built-in 8-byte key type.
    #[must_use]
    pub fn bytes8_key(self) -> DryIceWriterBuilder<W, S, Q, N, Bytes8Key> {
        self.record_key::<Bytes8Key>()
    }

    /// Configure the writer to store the built-in 16-byte key type.
    #[must_use]
    pub fn bytes16_key(self) -> DryIceWriterBuilder<W, S, Q, N, Bytes16Key> {
        self.record_key::<Bytes16Key>()
    }

    /// Configure the writer to store packed canonical prefix kmer keys.
    ///
    /// This is convenience sugar for `record_key::<PrefixKmer64<K>>()` and does
    /// not otherwise change payload layout or writing semantics.
    #[must_use]
    pub fn prefix_kmer_key<const K: u8>(self) -> DryIceWriterBuilder<W, S, Q, N, PrefixKmer64<K>> {
        self.record_key::<PrefixKmer64<K>>()
    }

    /// Configure the writer to store packed canonical minimizer keys.
    ///
    /// This is convenience sugar for `record_key::<Minimizer64<K, WN>>()` and
    /// does not otherwise change payload layout or writing semantics.
    #[must_use]
    pub fn minimizer_key<const K: u8, const WN: u8>(
        self,
    ) -> DryIceWriterBuilder<W, S, Q, N, Minimizer64<K, WN>> {
        self.record_key::<Minimizer64<K, WN>>()
    }

    /// Configure the writer to store prefix kmer keys using the library's default parameters.
    ///
    /// The current default is `PrefixKmer64<31>`.
    #[must_use]
    pub fn prefix_kmer_key_default(self) -> DryIceWriterBuilder<W, S, Q, N, PrefixKmer64<31>> {
        self.prefix_kmer_key::<31>()
    }

    /// Configure the writer to store minimizer keys using the library's default parameters.
    ///
    /// The current default is `Minimizer64<31, 15>`.
    #[must_use]
    pub fn minimizer_key_default(self) -> DryIceWriterBuilder<W, S, Q, N, Minimizer64<31, 15>> {
        self.minimizer_key::<31, 15>()
    }
}

impl<W> DryIceWriterBuilder<W, RawAsciiCodec, RawQualityCodec, RawNameCodec, NoRecordKey> {
    /// Configure the writer for key-only prefix-kmer output using the default parameters.
    ///
    /// This preset selects `PrefixKmer64<31>` and applies `empty_payload()`. It
    /// does not derive keys automatically; callers still write through the keyed
    /// writer APIs such as `write_key_only()` or `write_record_with_key()`.
    #[must_use]
    pub fn prefix_kmers(
        self,
    ) -> DryIceWriterBuilder<
        W,
        OmittedSequenceCodec,
        OmittedQualityCodec,
        OmittedNameCodec,
        PrefixKmer64<31>,
    > {
        self.prefix_kmer_key_default().empty_payload()
    }

    /// Configure the writer for prefix-kmer output while retaining sequences.
    ///
    /// This preset selects `PrefixKmer64<31>`, retains sequences, and omits
    /// names and qualities. Callers still supply or derive the key explicitly.
    #[must_use]
    pub fn prefix_kmers_with_sequences(
        self,
    ) -> DryIceWriterBuilder<
        W,
        RawAsciiCodec,
        OmittedQualityCodec,
        OmittedNameCodec,
        PrefixKmer64<31>,
    > {
        self.prefix_kmer_key_default().omit_quality().omit_names()
    }

    /// Configure the writer for prefix-kmer output while retaining names.
    ///
    /// This preset selects `PrefixKmer64<31>`, retains names, and omits
    /// sequences and qualities. Callers still supply or derive the key
    /// explicitly.
    #[must_use]
    pub fn prefix_kmers_with_names(
        self,
    ) -> DryIceWriterBuilder<
        W,
        OmittedSequenceCodec,
        OmittedQualityCodec,
        RawNameCodec,
        PrefixKmer64<31>,
    > {
        self.prefix_kmer_key_default()
            .omit_sequence()
            .omit_quality()
    }

    /// Configure the writer for key-only minimizer output using the default parameters.
    ///
    /// This preset selects `Minimizer64<31, 15>` and applies `empty_payload()`.
    /// It does not derive keys automatically; callers still write through the
    /// keyed writer APIs such as `write_key_only()` or `write_record_with_key()`.
    #[must_use]
    pub fn minimizers(
        self,
    ) -> DryIceWriterBuilder<
        W,
        OmittedSequenceCodec,
        OmittedQualityCodec,
        OmittedNameCodec,
        Minimizer64<31, 15>,
    > {
        self.minimizer_key_default().empty_payload()
    }

    /// Configure the writer for minimizer output while retaining sequences.
    ///
    /// This preset selects `Minimizer64<31, 15>`, retains sequences, and omits
    /// names and qualities. Callers still supply or derive the key explicitly.
    #[must_use]
    pub fn minimizers_with_sequences(
        self,
    ) -> DryIceWriterBuilder<
        W,
        RawAsciiCodec,
        OmittedQualityCodec,
        OmittedNameCodec,
        Minimizer64<31, 15>,
    > {
        self.minimizer_key_default().omit_quality().omit_names()
    }

    /// Configure the writer for minimizer output while retaining names.
    ///
    /// This preset selects `Minimizer64<31, 15>`, retains names, and omits
    /// sequences and qualities. Callers still supply or derive the key
    /// explicitly.
    #[must_use]
    pub fn minimizers_with_names(
        self,
    ) -> DryIceWriterBuilder<
        W,
        OmittedSequenceCodec,
        OmittedQualityCodec,
        RawNameCodec,
        Minimizer64<31, 15>,
    > {
        self.minimizer_key_default().omit_sequence().omit_quality()
    }
}

impl<W: Write, S: SequenceCodec, Q: QualityCodec, N: NameCodec>
    DryIceWriterBuilder<W, S, Q, N, NoRecordKey>
{
    /// Build an unkeyed writer.
    #[must_use]
    pub fn build(self) -> DryIceWriter<W, S, Q, N, NoRecordKey> {
        DryIceWriter::new_unkeyed(self.inner, self.target_block_records)
    }
}

impl<W: Write, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K: RecordKey>
    DryIceWriterBuilder<W, S, Q, N, K>
{
    /// Build a keyed writer.
    #[must_use]
    pub fn build(self) -> DryIceWriter<W, S, Q, N, K> {
        DryIceWriter::new_keyed(self.inner, self.target_block_records)
    }
}

#[cfg(feature = "async")]
impl<W, S: SequenceCodec, Q: QualityCodec, N: NameCodec>
    DryIceWriterBuilder<W, S, Q, N, NoRecordKey>
{
    /// Build an unkeyed async writer.
    #[must_use]
    pub fn build_async(self) -> crate::async_io::AsyncDryIceWriter<W, S, Q, N, NoRecordKey> {
        crate::async_io::AsyncDryIceWriter::new_unkeyed(self.inner, self.target_block_records)
    }
}

#[cfg(feature = "async")]
impl<W, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K: RecordKey>
    DryIceWriterBuilder<W, S, Q, N, K>
{
    /// Build a keyed async writer.
    #[must_use]
    pub fn build_async(self) -> crate::async_io::AsyncDryIceWriter<W, S, Q, N, K> {
        crate::async_io::AsyncDryIceWriter::new_keyed(self.inner, self.target_block_records)
    }
}

/// Writes sequencing records into the `dryice` block-oriented format.
pub struct DryIceWriter<
    W,
    S: SequenceCodec = RawAsciiCodec,
    Q: QualityCodec = RawQualityCodec,
    N: NameCodec = RawNameCodec,
    K = NoRecordKey,
> {
    inner: W,
    target_block_records: usize,
    block_builder: BlockBuilder<S, Q, N>,
    header_written: bool,
    _key: PhantomData<K>,
}

impl DryIceWriter<MissingInner, RawAsciiCodec, RawQualityCodec, RawNameCodec, NoRecordKey> {
    /// Start building a new writer.
    #[must_use]
    pub fn builder()
    -> DryIceWriterBuilder<MissingInner, RawAsciiCodec, RawQualityCodec, RawNameCodec, NoRecordKey>
    {
        DryIceWriterBuilder::new()
    }
}

impl<W, S: SequenceCodec, Q: QualityCodec, N: NameCodec> DryIceWriter<W, S, Q, N, NoRecordKey> {
    fn new_unkeyed(inner: W, target_block_records: usize) -> Self {
        let block_builder = BlockBuilder::new(&BlockBuilderConfig {
            record_key_width: None,
            record_key_tag: None,
            target_records: target_block_records,
        });

        Self {
            inner,
            target_block_records,
            block_builder,
            header_written: false,
            _key: PhantomData,
        }
    }
}

impl<W, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K: RecordKey> DryIceWriter<W, S, Q, N, K> {
    fn new_keyed(inner: W, target_block_records: usize) -> Self {
        let block_builder = BlockBuilder::new(&BlockBuilderConfig {
            record_key_width: Some(K::WIDTH),
            record_key_tag: Some(K::TYPE_TAG),
            target_records: target_block_records,
        });

        Self {
            inner,
            target_block_records,
            block_builder,
            header_written: false,
            _key: PhantomData,
        }
    }
}

impl<W, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K> DryIceWriter<W, S, Q, N, K> {
    fn ensure_header_written(&mut self) -> Result<(), DryIceError>
    where
        W: Write,
    {
        if !self.header_written {
            format::write_file_header(&mut self.inner)?;
            self.header_written = true;
        }
        Ok(())
    }
}

impl<W: Write, S: SequenceCodec, Q: QualityCodec, N: NameCodec>
    DryIceWriter<W, S, Q, N, NoRecordKey>
{
    /// Create an unkeyed writer from a pre-built options struct.
    ///
    /// # Errors
    ///
    /// Returns an error if the options request an unsupported block-size policy.
    pub fn from_options(inner: W, options: &DryIceWriterOptions) -> Result<Self, DryIceError> {
        let target_block_records = match options.layout.block_size {
            BlockSizePolicy::TargetRecords(n) => n,
            BlockSizePolicy::TargetBytes(_) => {
                return Err(DryIceError::InvalidWriterConfiguration(
                    "TargetBytes block size policy is not yet supported",
                ));
            },
        };

        Ok(Self::new_unkeyed(inner, target_block_records))
    }

    /// Assemble the current configuration into an options struct.
    #[must_use]
    pub fn options(&self) -> DryIceWriterOptions {
        DryIceWriterOptions {
            layout: BlockLayoutOptions {
                block_size: BlockSizePolicy::TargetRecords(self.target_block_records),
            },
        }
    }

    /// Write a single sequencing record to an unkeyed writer.
    ///
    /// # Errors
    ///
    /// Returns an error if the record fails validation, if the file header cannot
    /// be written, or if flushing the current block fails.
    pub fn write_record<R: SeqRecordLike>(&mut self, record: &R) -> Result<(), DryIceError> {
        self.ensure_header_written()?;
        self.block_builder.push_record(record)?;

        if self.block_builder.should_flush() {
            self.flush_block()?;
        }

        Ok(())
    }
}

impl<W: Write, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K: RecordKey>
    DryIceWriter<W, S, Q, N, K>
{
    /// Write a single sequencing record together with its accelerator key.
    ///
    /// # Errors
    ///
    /// Returns an error if the record fails validation, if the key cannot be
    /// encoded, if the file header cannot be written, or if flushing the current
    /// block fails.
    pub fn write_record_with_key<R: SeqRecordLike>(
        &mut self,
        record: &R,
        key: &K,
    ) -> Result<(), DryIceError> {
        self.ensure_header_written()?;
        self.block_builder.push_record_with_key(record, key)?;

        if self.block_builder.should_flush() {
            self.flush_block()?;
        }

        Ok(())
    }

    /// Write a key-only record with empty payload fields.
    ///
    /// # Errors
    ///
    /// Returns an error if the file header cannot be written, if the key cannot
    /// be encoded, or if flushing the current block fails.
    pub fn write_key_only(&mut self, key: &K) -> Result<(), DryIceError> {
        self.write_record_with_key(&EMPTY_RECORD, key)
    }
}

impl<W: Write, S: SequenceCodec, Q: QualityCodec, N: NameCodec, K> DryIceWriter<W, S, Q, N, K> {
    /// Flush any remaining buffered records and finalize the file.
    ///
    /// # Errors
    ///
    /// Returns an error if writing the file header or flushing the final block
    /// fails.
    pub fn finish(mut self) -> Result<W, DryIceError> {
        self.ensure_header_written()?;

        if !self.block_builder.is_empty() {
            self.flush_block()?;
        }

        Ok(self.inner)
    }

    fn flush_block(&mut self) -> Result<(), DryIceError> {
        self.block_builder.write_block(&mut self.inner)?;
        Ok(())
    }
}