fastars 0.1.0

//! Split writer for dividing FASTQ output into multiple files.
//!
//! This module provides writers that can split output into multiple files
//! either by file count or by line count.

use std::path::{Path, PathBuf};

use anyhow::{Context, Result};

use super::{CompressionType, FastqWriter, OwnedRecord};
use crate::cli::{SplitConfig, SplitMode};

// ============================================================================
// Single-End Split Writer
// ============================================================================

/// Writer that splits single-end output into multiple files.
///
/// Files are named with a numeric suffix: base_0001.fastq.gz, base_0002.fastq.gz, etc.
pub struct SplitWriter {
    /// Base path for output files (without extension).
    base_path: PathBuf,
    /// Split configuration.
    config: SplitConfig,
    /// Compression type for output files.
    compression: CompressionType,
    /// Current file index (0-based).
    current_file_idx: usize,
    /// Current writer.
    current_writer: Option<FastqWriter>,
    /// Total records written to current file.
    records_in_current_file: usize,
    /// Total records written overall.
    total_records_written: usize,
    /// Records per file (for ByFile mode).
    records_per_file: Option<usize>,
}

impl SplitWriter {
    /// Create a new split writer.
    ///
    /// # Arguments
    /// - `base_path`: Base path for output files (without numeric suffix)
    /// - `config`: Split configuration
    /// - `total_records`: Total number of records (required for ByFile mode)
    ///
    /// For ByFile mode, `total_records` is used to calculate records per file.
    /// For ByLines mode, it's not used (pass None).
    pub fn new(
        base_path: impl AsRef<Path>,
        config: SplitConfig,
        total_records: Option<usize>,
    ) -> Result<Self> {
        let base_path = base_path.as_ref().to_path_buf();
        let compression = CompressionType::from_path(&base_path);

        // Calculate records per file for ByFile mode
        let records_per_file = match config.mode {
            SplitMode::ByFile(n_files) => {
                let total = total_records.ok_or_else(|| {
                    anyhow::anyhow!("total_records required for ByFile split mode")
                })?;
                Some((total + n_files - 1) / n_files) // Ceiling division
            }
            SplitMode::ByLines(_) => None,
        };

        Ok(Self {
            base_path,
            config,
            compression,
            current_file_idx: 0,
            current_writer: None,
            records_in_current_file: 0,
            total_records_written: 0,
            records_per_file,
        })
    }

    /// Write a single record, handling file switching as needed.
    pub fn write_record(&mut self, record: &OwnedRecord) -> Result<()> {
        // Check if we need to switch to a new file
        if self.should_switch_file() {
            self.switch_to_next_file()?;
        }

        // Ensure we have a writer
        if self.current_writer.is_none() {
            self.switch_to_next_file()?;
        }

        // Write the record
        if let Some(writer) = &mut self.current_writer {
            writer
                .write_record(record)
                .context("Failed to write record")?;
        }

        self.records_in_current_file += 1;
        self.total_records_written += 1;

        Ok(())
    }

    /// Write a batch of records.
    pub fn write_batch(&mut self, records: &[OwnedRecord]) -> Result<()> {
        for record in records {
            self.write_record(record)?;
        }
        Ok(())
    }

    /// Check if we should switch to a new file.
    fn should_switch_file(&self) -> bool {
        match self.config.mode {
            SplitMode::ByFile(_) => {
                if let Some(per_file) = self.records_per_file {
                    self.records_in_current_file >= per_file
                } else {
                    false
                }
            }
            SplitMode::ByLines(lines) => {
                let records_per_file = lines / 4;
                self.records_in_current_file >= records_per_file
            }
        }
    }

    /// Switch to the next output file.
    fn switch_to_next_file(&mut self) -> Result<()> {
        // Flush and close current writer
        if let Some(mut writer) = self.current_writer.take() {
            writer.flush()?;
        }

        // Generate the new file path
        let file_path = self.generate_file_path(self.current_file_idx);

        // Create a new writer
        let writer = FastqWriter::new(&file_path, self.compression)
            .with_context(|| format!("Failed to create split file: {}", file_path.display()))?;

        self.current_writer = Some(writer);
        self.current_file_idx += 1;
        self.records_in_current_file = 0;

        Ok(())
    }

    /// Generate the file path for a given file index.
    fn generate_file_path(&self, idx: usize) -> PathBuf {
        let suffix = format!("{:0width$}", idx + 1, width = self.config.prefix_digits);

        // Extract base name and parent directory
        let parent = self.base_path.parent().unwrap_or(Path::new(""));
        let stem = self
            .base_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("output");

        // Determine extension based on compression type
        let extension = match self.compression {
            CompressionType::Gzip | CompressionType::ParallelGzip => "fastq.gz",
            CompressionType::None => "fastq",
        };

        parent.join(format!("{}_{}.{}", stem, suffix, extension))
    }

    /// Finalize all files.
    pub fn finish(&mut self) -> Result<()> {
        if let Some(mut writer) = self.current_writer.take() {
            writer.flush()?;
        }
        Ok(())
    }

    /// Get the total number of records written.
    pub fn total_records_written(&self) -> usize {
        self.total_records_written
    }

    /// Get the number of output files created.
    pub fn num_files_created(&self) -> usize {
        self.current_file_idx
    }

    /// Get all created file paths.
    pub fn get_output_files(&self) -> Vec<PathBuf> {
        (0..self.current_file_idx)
            .map(|idx| self.generate_file_path(idx))
            .collect()
    }
}

impl Drop for SplitWriter {
    fn drop(&mut self) {
        let _ = self.finish();
    }
}

// ============================================================================
// Paired-End Split Writer
// ============================================================================

/// Writer that splits paired-end output into multiple file pairs.
///
/// Both R1 and R2 files are split synchronously with the same record counts.
/// Files are named: base_0001.R1.fastq.gz, base_0001.R2.fastq.gz, etc.
pub struct PairedSplitWriter {
    /// Base path for output files.
    base_path: PathBuf,
    /// Split configuration.
    config: SplitConfig,
    /// Compression type.
    compression: CompressionType,
    /// Current file index.
    current_file_idx: usize,
    /// Current R1 writer.
    current_r1_writer: Option<FastqWriter>,
    /// Current R2 writer.
    current_r2_writer: Option<FastqWriter>,
    /// Records in current file.
    records_in_current_file: usize,
    /// Total records written.
    total_records_written: usize,
    /// Records per file (for ByFile mode).
    records_per_file: Option<usize>,
}

impl PairedSplitWriter {
    /// Create a new paired split writer.
    pub fn new(
        base_path: impl AsRef<Path>,
        config: SplitConfig,
        total_records: Option<usize>,
    ) -> Result<Self> {
        let base_path = base_path.as_ref().to_path_buf();
        let compression = CompressionType::from_path(&base_path);

        let records_per_file = match config.mode {
            SplitMode::ByFile(n_files) => {
                let total = total_records.ok_or_else(|| {
                    anyhow::anyhow!("total_records required for ByFile split mode")
                })?;
                Some((total + n_files - 1) / n_files)
            }
            SplitMode::ByLines(_) => None,
        };

        Ok(Self {
            base_path,
            config,
            compression,
            current_file_idx: 0,
            current_r1_writer: None,
            current_r2_writer: None,
            records_in_current_file: 0,
            total_records_written: 0,
            records_per_file,
        })
    }

    /// Write a pair of records.
    pub fn write_pair(&mut self, r1: &OwnedRecord, r2: &OwnedRecord) -> Result<()> {
        // Check if we need to switch files
        if self.should_switch_file() {
            self.switch_to_next_file()?;
        }

        // Ensure we have writers
        if self.current_r1_writer.is_none() || self.current_r2_writer.is_none() {
            self.switch_to_next_file()?;
        }

        // Write both records
        if let (Some(r1_writer), Some(r2_writer)) =
            (&mut self.current_r1_writer, &mut self.current_r2_writer)
        {
            r1_writer
                .write_record(r1)
                .context("Failed to write R1 record")?;
            r2_writer
                .write_record(r2)
                .context("Failed to write R2 record")?;
        }

        self.records_in_current_file += 1;
        self.total_records_written += 1;

        Ok(())
    }

    /// Write a batch of paired records.
    pub fn write_batch(&mut self, pairs: &[(OwnedRecord, OwnedRecord)]) -> Result<()> {
        for (r1, r2) in pairs {
            self.write_pair(r1, r2)?;
        }
        Ok(())
    }

    /// Check if we should switch files.
    fn should_switch_file(&self) -> bool {
        match self.config.mode {
            SplitMode::ByFile(_) => {
                if let Some(per_file) = self.records_per_file {
                    self.records_in_current_file >= per_file
                } else {
                    false
                }
            }
            SplitMode::ByLines(lines) => {
                let records_per_file = lines / 4;
                self.records_in_current_file >= records_per_file
            }
        }
    }

    /// Switch to the next output file pair.
    fn switch_to_next_file(&mut self) -> Result<()> {
        // Flush and close current writers
        if let Some(mut writer) = self.current_r1_writer.take() {
            writer.flush()?;
        }
        if let Some(mut writer) = self.current_r2_writer.take() {
            writer.flush()?;
        }

        // Generate new file paths
        let (r1_path, r2_path) = self.generate_file_paths(self.current_file_idx);

        // Create new writers
        let r1_writer = FastqWriter::new(&r1_path, self.compression)
            .with_context(|| format!("Failed to create split file: {}", r1_path.display()))?;
        let r2_writer = FastqWriter::new(&r2_path, self.compression)
            .with_context(|| format!("Failed to create split file: {}", r2_path.display()))?;

        self.current_r1_writer = Some(r1_writer);
        self.current_r2_writer = Some(r2_writer);
        self.current_file_idx += 1;
        self.records_in_current_file = 0;

        Ok(())
    }

    /// Generate file paths for R1 and R2.
    fn generate_file_paths(&self, idx: usize) -> (PathBuf, PathBuf) {
        let suffix = format!("{:0width$}", idx + 1, width = self.config.prefix_digits);

        let parent = self.base_path.parent().unwrap_or(Path::new(""));
        let stem = self
            .base_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("output");

        let extension = match self.compression {
            CompressionType::Gzip | CompressionType::ParallelGzip => "fastq.gz",
            CompressionType::None => "fastq",
        };

        let r1_path = parent.join(format!("{}_{}.R1.{}", stem, suffix, extension));
        let r2_path = parent.join(format!("{}_{}.R2.{}", stem, suffix, extension));

        (r1_path, r2_path)
    }

    /// Finalize all files.
    pub fn finish(&mut self) -> Result<()> {
        if let Some(mut writer) = self.current_r1_writer.take() {
            writer.flush()?;
        }
        if let Some(mut writer) = self.current_r2_writer.take() {
            writer.flush()?;
        }
        Ok(())
    }

    /// Get the total number of records written.
    pub fn total_records_written(&self) -> usize {
        self.total_records_written
    }

    /// Get the number of output file pairs created.
    pub fn num_files_created(&self) -> usize {
        self.current_file_idx
    }

    /// Get all created file paths (R1 and R2).
    pub fn get_output_files(&self) -> Vec<PathBuf> {
        let mut files = Vec::new();
        for idx in 0..self.current_file_idx {
            let (r1, r2) = self.generate_file_paths(idx);
            files.push(r1);
            files.push(r2);
        }
        files
    }
}

impl Drop for PairedSplitWriter {
    fn drop(&mut self) {
        let _ = self.finish();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::SplitMode;
    use tempfile::tempdir;

    fn make_record(name: &str, seq: &str) -> OwnedRecord {
        let qual = vec![b'I'; seq.len()];
        OwnedRecord::new(name.as_bytes().to_vec(), seq.as_bytes().to_vec(), qual)
    }

    #[test]
    fn test_split_writer_by_file() {
        let dir = tempdir().unwrap();
        let base_path = dir.path().join("output");

        let config = SplitConfig {
            mode: SplitMode::ByFile(2),
            prefix_digits: 4,
        };

        let mut writer = SplitWriter::new(&base_path, config, Some(10)).unwrap();

        // Write 10 records
        for i in 0..10 {
            let record = make_record(&format!("read{}", i), "ACGTACGTACGT");
            writer.write_record(&record).unwrap();
        }

        writer.finish().unwrap();

        assert_eq!(writer.num_files_created(), 2);
        assert_eq!(writer.total_records_written(), 10);

        // Check that files exist (no .gz extension since base_path has no extension)
        assert!(dir.path().join("output_0001.fastq").exists());
        assert!(dir.path().join("output_0002.fastq").exists());
    }

    #[test]
    fn test_split_writer_by_lines() {
        let dir = tempdir().unwrap();
        let base_path = dir.path().join("output");

        let config = SplitConfig {
            mode: SplitMode::ByLines(8), // 2 records per file
            prefix_digits: 3,
        };

        let mut writer = SplitWriter::new(&base_path, config, None).unwrap();

        // Write 6 records
        for i in 0..6 {
            let record = make_record(&format!("read{}", i), "ACGTACGTACGT");
            writer.write_record(&record).unwrap();
        }

        writer.finish().unwrap();

        assert_eq!(writer.num_files_created(), 3);
        assert_eq!(writer.total_records_written(), 6);

        // Check file naming (no .gz extension since base_path has no extension)
        assert!(dir.path().join("output_001.fastq").exists());
        assert!(dir.path().join("output_002.fastq").exists());
        assert!(dir.path().join("output_003.fastq").exists());
    }

    #[test]
    fn test_paired_split_writer() {
        let dir = tempdir().unwrap();
        let base_path = dir.path().join("output");

        let config = SplitConfig {
            mode: SplitMode::ByFile(2),
            prefix_digits: 4,
        };

        let mut writer = PairedSplitWriter::new(&base_path, config, Some(6)).unwrap();

        // Write 6 pairs
        for i in 0..6 {
            let r1 = make_record(&format!("read{}/1", i), "ACGTACGTACGT");
            let r2 = make_record(&format!("read{}/2", i), "TGCATGCATGCA");
            writer.write_pair(&r1, &r2).unwrap();
        }

        writer.finish().unwrap();

        assert_eq!(writer.num_files_created(), 2);
        assert_eq!(writer.total_records_written(), 6);

        // Check that both R1 and R2 files exist (no .gz extension since base_path has no extension)
        assert!(dir.path().join("output_0001.R1.fastq").exists());
        assert!(dir.path().join("output_0001.R2.fastq").exists());
        assert!(dir.path().join("output_0002.R1.fastq").exists());
        assert!(dir.path().join("output_0002.R2.fastq").exists());
    }

    #[test]
    fn test_write_batch() {
        let dir = tempdir().unwrap();
        let base_path = dir.path().join("output");

        let config = SplitConfig {
            mode: SplitMode::ByLines(8),
            prefix_digits: 4,
        };

        let mut writer = SplitWriter::new(&base_path, config, None).unwrap();

        let records: Vec<_> = (0..4)
            .map(|i| make_record(&format!("read{}", i), "ACGTACGT"))
            .collect();

        writer.write_batch(&records).unwrap();
        writer.finish().unwrap();

        assert_eq!(writer.total_records_written(), 4);
    }
}