cqlite_core/storage/write_engine/

mod.rs

//! Write engine for SSTable generation and persistence (M5)
//!
//! This module provides the write path for CQLite, implementing WAL-backed
//! memtable flushing and K-way merge for producing valid Cassandra 5.0 SSTables.
//!
//! ## Architecture
//!
//! The WriteEngine is the public API that coordinates:
//! 1. WAL (Write-Ahead Log) - Durability
//! 2. Memtable - In-memory buffer
//! 3. SSTableWriter - On-disk persistence
//!
//! ## Write Flow
//!
//! 1. User calls `write(mutation)` or `execute(cql_statement)`
//! 2. WriteEngine appends to WAL (durability)
//! 3. WriteEngine inserts into Memtable
//! 4. When Memtable exceeds threshold → flush to SSTable
//! 5. After successful flush → truncate WAL
//!
//! ## Recovery
//!
//! On startup, the WriteEngine replays WAL entries into the memtable.

#[cfg(feature = "write-support")]
pub mod cql_to_mutation;
#[cfg(feature = "write-support")]
pub mod export;
#[cfg(feature = "write-support")]
pub mod memtable;
#[cfg(feature = "write-support")]
pub mod merge;
#[cfg(feature = "write-support")]
pub mod merge_policy;
#[cfg(feature = "write-support")]
pub mod mutation;
#[cfg(feature = "write-support")]
pub mod wal;

#[cfg(feature = "write-support")]
pub use export::{ExportOptions, ExportReport};
#[cfg(feature = "write-support")]
pub use memtable::Memtable;
#[cfg(feature = "write-support")]
pub use merge::KWayMerger;
#[cfg(feature = "write-support")]
pub use merge_policy::STCSPolicy;
#[cfg(feature = "write-support")]
pub use mutation::{
    CellOperation, ClusteringBound, ClusteringKey, DecoratedKey, Mutation, PartitionKey,
    PartitionTombstone, RangeTombstone, TableId,
};
#[cfg(feature = "write-support")]
pub use wal::WriteAheadLog;

use crate::error::{Error, Result};
use crate::schema::TableSchema;
use crate::storage::sstable::writer::SSTableInfo;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

/// Maintenance report from a maintenance_step() call (M5.2, Issue #384)
#[cfg(feature = "write-support")]
#[derive(Debug, Clone)]
pub struct MaintenanceReport {
    /// Time spent in this maintenance step
    pub time_spent: Duration,
    /// Completed merge output files (if any merge completed)
    pub completed_merges: Vec<PathBuf>,
    /// Number of rows merged in this step
    pub rows_merged: u64,
    /// Number of bytes written in this step
    pub bytes_written: u64,
    /// Whether there is pending compaction work
    pub pending_compaction: bool,
}

/// Cumulative statistics across all compaction operations (M5.2, Issue #474)
///
/// Tracks lifetime totals for monitoring compaction health and throughput.
/// Updated atomically at the end of each successful merge.
#[cfg(feature = "write-support")]
#[derive(Debug, Clone, Default)]
pub struct CompactionStats {
    /// Total number of completed compaction cycles
    pub compactions_completed: u64,
    /// Total number of input SSTables consumed
    pub sstables_merged_in: u64,
    /// Total number of output SSTables produced
    pub sstables_produced: u64,
    /// Total bytes read from input SSTables
    pub bytes_read: u64,
    /// Total bytes written to output SSTables
    pub bytes_written: u64,
    /// Total rows merged across all compactions
    pub rows_merged: u64,
    /// Total wall-clock time spent in compaction
    pub total_time: Duration,
}

/// Trait for merge policy implementations (M5.2, Issue #383)
///
/// A merge policy decides which SSTables should be compacted together.
/// This trait allows different compaction strategies (STCS, LCS, TWCS, etc.)
/// to be plugged into the WriteEngine.
#[cfg(feature = "write-support")]
pub trait MergePolicy: Send + std::fmt::Debug {
    /// Select SSTables for the next compaction
    ///
    /// # Arguments
    ///
    /// * `candidates` - Available SSTable paths in the data directory
    ///
    /// # Returns
    ///
    /// Paths to SSTables that should be merged, ordered newest to oldest.
    /// Returns empty Vec if no compaction is needed.
    fn select_merge(&self, candidates: &[PathBuf]) -> Result<Vec<PathBuf>>;
}

/// Active merge state for incremental compaction (M5.2, Issue #384)
#[cfg(feature = "write-support")]
#[derive(Debug)]
struct ActiveMerge {
    /// K-way merger performing the compaction
    merger: KWayMerger,
    /// Output SSTable writer (writes to `tmp_dir/keyspace/table/`)
    writer: crate::storage::sstable::writer::SSTableWriter,
    /// Input SSTable paths being merged (these remain intact until atomic rename succeeds)
    input_paths: Vec<PathBuf>,
    /// Root of the temporary directory tree used for this compaction output.
    ///
    /// The SSTableWriter appends `keyspace/table/` to this path, so component
    /// files land at `tmp_dir/keyspace/table/nb-{gen}-big-*.{ext}`.
    ///
    /// After `writer.finish()` the files are atomically renamed to the final
    /// SSTable directory. Only then are the inputs deleted.
    ///
    /// Invariant: if the process crashes before the renames complete, `tmp_dir`
    /// may contain partial output but the input SSTables remain intact.
    tmp_dir: PathBuf,
    /// Final SSTable directory (`data_dir/keyspace/table/`)
    ///
    /// Stored here so `finalize_merge_async` doesn't have to recompute it.
    sstable_dir: PathBuf,
    /// Number of rows merged so far (updated per partition)
    rows_merged: u64,
    /// Total bytes read from input SSTables (approximate: sum of Data.db file sizes)
    bytes_read: u64,
    /// When this merge started
    started_at: Instant,
}

/// WAL durability mode for the write engine.
///
/// Controls whether `write` and `write_async` append to and fsync the
/// write-ahead log on every call.  The default (`SyncEachWrite`) matches the
/// pre-existing behavior and is the **only safe choice for production
/// workloads** — a process crash between a successful `write` and a later
/// `flush` will lose mutations written with `Disabled`.
///
/// ## When to use `Disabled`
///
/// - **Bulk-load / import pipelines** where the source data is replayable and
///   you are willing to re-run the load on failure.
/// - **Benchmarking** where you want to isolate CPU-bound write throughput from
///   fsync latency.  The companion `write/ingest_wal_off` Criterion bench uses
///   this variant (see `cqlite-core/benches/write.rs`).
///
/// In both cases, call [`WriteEngine::flush`] (and, optionally,
/// [`WriteEngine::close`]) when the load is finished so the data is durably
/// persisted to SSTables.
///
/// ## WAL replay on restart
///
/// When `Disabled`, no WAL entries are written.  Reopening the engine on the
/// same `wal_dir` after a crash will replay **zero** mutations, even if
/// `flush` was never called.  If you need crash-safe recovery, use
/// `SyncEachWrite`.
///
/// # Example
///
/// ```rust,ignore
/// use cqlite_core::storage::write_engine::{Durability, WriteEngineConfig};
///
/// // Production (default)
/// let config = WriteEngineConfig::new(data, wal, schema);
///
/// // Bulk-load / benchmarking
/// let config = WriteEngineConfig::new(data, wal, schema)
///     .with_durability(Durability::Disabled);
/// ```
#[cfg(feature = "write-support")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Durability {
    /// Append to the WAL and call `fsync` on every `write` / `write_async`
    /// call.  A successful return guarantees the mutation is durable on disk.
    ///
    /// This is the **default** and the safe choice for all production
    /// workloads.
    #[default]
    SyncEachWrite,

    /// Skip WAL append **and** fsync on every `write` / `write_async` call.
    /// Mutations are buffered in the memtable only.  Data is durable only
    /// after a successful [`WriteEngine::flush`].
    ///
    /// **Use only for bulk-load pipelines and benchmarks where durability can
    /// be traded for throughput.**
    Disabled,
}

/// Write engine configuration
#[cfg(feature = "write-support")]
#[derive(Debug, Clone)]
pub struct WriteEngineConfig {
    /// Directory for SSTable data files
    pub data_dir: PathBuf,
    /// Directory for WAL files
    pub wal_dir: PathBuf,
    /// Memtable flush threshold in bytes (default: 64MB)
    pub memtable_flush_threshold: usize,
    /// Memtable hard limit in bytes (default: 256MB)
    /// When this limit is reached, writes will fail with an error
    pub memtable_hard_limit: usize,
    /// Table schema for column metadata
    pub schema: TableSchema,
    /// WAL durability mode (default: [`Durability::SyncEachWrite`])
    pub durability: Durability,
}

#[cfg(feature = "write-support")]
impl WriteEngineConfig {
    /// Default flush threshold (64 MB)
    pub const DEFAULT_FLUSH_THRESHOLD: usize = 64 * 1024 * 1024;
    /// Default hard limit (256 MB)
    pub const DEFAULT_HARD_LIMIT: usize = 256 * 1024 * 1024;

    /// Create a new configuration with default flush threshold
    pub fn new(data_dir: PathBuf, wal_dir: PathBuf, schema: TableSchema) -> Self {
        Self {
            data_dir,
            wal_dir,
            memtable_flush_threshold: Self::DEFAULT_FLUSH_THRESHOLD,
            memtable_hard_limit: Self::DEFAULT_HARD_LIMIT,
            schema,
            durability: Durability::default(),
        }
    }

    /// Set a custom flush threshold
    pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
        self.memtable_flush_threshold = threshold;
        self
    }

    /// Set a custom hard limit
    pub fn with_hard_limit(mut self, limit: usize) -> Self {
        self.memtable_hard_limit = limit;
        self
    }

    /// Set the WAL durability mode.
    ///
    /// Mirrors `with_flush_threshold` in style. See [`Durability`] for the
    /// trade-offs between [`Durability::SyncEachWrite`] (default, production)
    /// and [`Durability::Disabled`] (bulk-load / benchmarking).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use cqlite_core::storage::write_engine::{Durability, WriteEngineConfig};
    ///
    /// let config = WriteEngineConfig::new(data, wal, schema)
    ///     .with_durability(Durability::Disabled);
    /// ```
    pub fn with_durability(mut self, durability: Durability) -> Self {
        self.durability = durability;
        self
    }
}

/// Write engine coordinator
///
/// Orchestrates WAL, memtable, and SSTable flushing for write operations.
/// This is the primary public API for all write operations in CQLite.
///
/// ## Thread Safety
///
/// WriteEngine follows a single-writer model. It is NOT thread-safe and
/// should be used from a single thread or protected by external locking.
/// The `closed` flag uses atomic operations for safe concurrent access checking.
///
/// ## Example
///
/// ```rust,ignore
/// use cqlite_core::storage::write_engine::{WriteEngine, WriteEngineConfig, Mutation};
/// use std::path::PathBuf;
///
/// // Create configuration
/// let config = WriteEngineConfig::new(
///     PathBuf::from("data"),
///     PathBuf::from("wal"),
///     schema
/// );
///
/// // Create engine
/// let mut engine = WriteEngine::new(config)?;
///
/// // Write a mutation
/// engine.write(mutation)?;
///
/// // Execute CQL statement
/// engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')")?;
///
/// // Flush to SSTable
/// engine.flush()?;
///
/// // Close cleanly
/// engine.close()?;
/// ```
#[cfg(feature = "write-support")]
#[derive(Debug)]
pub struct WriteEngine {
    /// Configuration
    config: WriteEngineConfig,
    /// Write-ahead log for durability
    wal: WriteAheadLog,
    /// In-memory write buffer
    memtable: Memtable,
    /// SSTable generation counter (increments on each flush)
    generation: u64,
    /// Whether the engine has been closed (atomic for thread safety)
    closed: AtomicBool,
    /// Active merge state for incremental compaction (M5.2)
    active_merge: Option<ActiveMerge>,
    /// Merge policy for compaction decisions (M5.2)
    merge_policy: Option<Box<dyn MergePolicy>>,
    /// Cumulative compaction statistics (M5.2, Issue #474)
    cumulative_stats: CompactionStats,
}

/// Reject any mutation that contains a counter cell write.
///
/// Counter columns require server-side distributed increment semantics and
/// cannot be expressed as a last-write-wins mutation.  Both the sync
/// `write()` and the async `write_async()` paths call this guard immediately
/// after the closed-check.
#[cfg(feature = "write-support")]
fn reject_counter_cells(mutation: &Mutation) -> Result<()> {
    for op in &mutation.operations {
        match op {
            CellOperation::Write { value, .. } | CellOperation::WriteWithTtl { value, .. } => {
                if matches!(value, crate::types::Value::Counter(_)) {
                    return Err(Error::invalid_operation(
                        "counter writes are not supported via the standard mutation path; \
                         counter columns require server-side distributed increment semantics",
                    ));
                }
            }
            _ => {}
        }
    }
    Ok(())
}

#[cfg(feature = "write-support")]
impl WriteEngine {
    /// Create a new write engine
    ///
    /// This initializes the WAL and memtable. If a WAL exists in the
    /// wal_dir, it will be replayed to recover in-flight writes.
    ///
    /// # Arguments
    ///
    /// * `config` - Write engine configuration
    ///
    /// # Returns
    ///
    /// A new WriteEngine ready to accept writes.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - WAL directory doesn't exist
    /// - Data directory doesn't exist
    /// - WAL replay fails
    pub fn new(config: WriteEngineConfig) -> Result<Self> {
        // Ensure directories exist
        std::fs::create_dir_all(&config.data_dir).map_err(|e| {
            Error::Storage(format!(
                "Failed to create data directory {:?}: {}",
                config.data_dir, e
            ))
        })?;

        std::fs::create_dir_all(&config.wal_dir).map_err(|e| {
            Error::Storage(format!(
                "Failed to create WAL directory {:?}: {}",
                config.wal_dir, e
            ))
        })?;

        // Startup sweep: remove orphaned compaction artifacts left by a previous crash.
        //
        // Two kinds of orphans can be left if the process crashes mid-rename in
        // `finalize_merge_async`:
        //
        //   (a) A `.compaction-tmp-{gen}/` directory under `data_dir` with partial
        //       component files.
        //
        //   (b) A partial set of renamed components in `data_dir/{keyspace}/{table}/`
        //       — specifically one or more `nb-{gen}-big-*.db` files without a
        //       matching `TOC.txt`. Because `scan_data_files` discovers SSTables by
        //       `nb-*-big-Data.db` glob, an orphaned Data.db without TOC.txt will be
        //       picked up by the merge policy and fed to `KWayMerger`, which may
        //       produce garbled output.
        //
        // Both sweeps are best-effort: individual failures are logged as warnings but
        // do not abort engine startup.
        Self::sweep_orphaned_compaction_tmp(&config.data_dir);
        Self::sweep_orphaned_partial_sstables(
            &config.data_dir,
            &config.schema.keyspace,
            &config.schema.table,
        );

        // Initialize WAL
        let wal_path = config.wal_dir.join(WriteAheadLog::WAL_FILENAME);
        let wal = if wal_path.exists() {
            // Recover from existing WAL
            WriteAheadLog::open_existing(&wal_path)?
        } else {
            // Create new WAL
            WriteAheadLog::create(&config.wal_dir)?
        };

        // Replay WAL into memtable
        let mut memtable = Memtable::new();
        let mutations = wal.replay()?;

        if !mutations.is_empty() {
            log::info!("Replaying {} mutations from WAL", mutations.len());

            for mutation in mutations {
                // Compute decorated key
                let decorated_key = mutation.decorated_key(&config.schema)?;

                // Insert into memtable
                memtable.insert_with_key(decorated_key, mutation)?;
            }

            log::info!(
                "WAL replay complete: {} rows in memtable, {} bytes",
                memtable.row_count(),
                memtable.size_bytes()
            );
        }

        // Determine next generation number by scanning data directory
        let generation = Self::determine_next_generation(&config.data_dir)?;

        Ok(Self {
            config,
            wal,
            memtable,
            generation,
            closed: AtomicBool::new(false),
            active_merge: None,
            merge_policy: None,
            cumulative_stats: CompactionStats::default(),
        })
    }

    /// Write a mutation to the write engine
    ///
    /// This appends the mutation to the WAL for durability, then inserts it
    /// into the memtable. If the memtable exceeds the flush threshold,
    /// an automatic flush is triggered.
    ///
    /// **Note**: Automatic flush is disabled when called from an async context.
    /// Use `write_async()` for async contexts with automatic flush support.
    ///
    /// # Arguments
    ///
    /// * `mutation` - The mutation to write
    ///
    /// # Returns
    ///
    /// Ok(()) on success, or an error if the write fails.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Engine has been closed
    /// - WAL append fails
    /// - Memtable insert fails
    /// - Automatic flush fails (sync context only)
    pub fn write(&mut self, mutation: Mutation) -> Result<()> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(Error::InvalidInput(
                "WriteEngine has been closed".to_string(),
            ));
        }

        reject_counter_cells(&mutation)?;

        // Check hard limit before accepting write
        if self.memtable.size_bytes() >= self.config.memtable_hard_limit {
            return Err(Error::Storage(format!(
                "Memtable at hard limit ({} bytes >= {} bytes). Flush required before accepting more writes.",
                self.memtable.size_bytes(),
                self.config.memtable_hard_limit
            )));
        }

        // 1. Append to WAL (durability) — skipped when Durability::Disabled
        if self.config.durability == Durability::SyncEachWrite {
            self.wal.append(&mutation)?;
            self.wal.sync()?;
        }

        // 2. Compute decorated key from partition key
        let decorated_key = mutation.decorated_key(&self.config.schema)?;

        // 3. Insert into memtable
        self.memtable.insert_with_key(decorated_key, mutation)?;

        // 4. Check if memtable should be flushed (only in non-async context)
        if self
            .memtable
            .should_flush(self.config.memtable_flush_threshold)
        {
            log::warn!(
                "Memtable size {} exceeds threshold {} - call flush() manually in async context",
                self.memtable.size_bytes(),
                self.config.memtable_flush_threshold
            );

            // Try to flush synchronously only if we're not in an async context
            if tokio::runtime::Handle::try_current().is_err() {
                log::info!("Triggering automatic flush");
                self.flush_internal()?;
            }
        }

        Ok(())
    }

    /// Write a mutation with async automatic flush support
    ///
    /// This is the async version of `write()` that supports automatic flushing
    /// in async contexts. Use this method when calling from async code.
    ///
    /// # Arguments
    ///
    /// * `mutation` - The mutation to write
    ///
    /// # Returns
    ///
    /// Ok(()) on success, or an error if the write fails.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Engine has been closed
    /// - WAL append fails
    /// - Memtable insert fails
    /// - Automatic flush fails
    pub async fn write_async(&mut self, mutation: Mutation) -> Result<()> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(Error::InvalidInput(
                "WriteEngine has been closed".to_string(),
            ));
        }

        reject_counter_cells(&mutation)?;

        // Check hard limit before accepting write
        if self.memtable.size_bytes() >= self.config.memtable_hard_limit {
            return Err(Error::Storage(format!(
                "Memtable at hard limit ({} bytes >= {} bytes). Flush required before accepting more writes.",
                self.memtable.size_bytes(),
                self.config.memtable_hard_limit
            )));
        }

        // 1. Append to WAL (durability) — skipped when Durability::Disabled
        if self.config.durability == Durability::SyncEachWrite {
            self.wal.append(&mutation)?;
            self.wal.sync()?;
        }

        // 2. Compute decorated key from partition key
        let decorated_key = mutation.decorated_key(&self.config.schema)?;

        // 3. Insert into memtable
        self.memtable.insert_with_key(decorated_key, mutation)?;

        // 4. Check if memtable should be flushed
        if self
            .memtable
            .should_flush(self.config.memtable_flush_threshold)
        {
            log::info!(
                "Memtable size {} exceeds threshold {}, triggering flush",
                self.memtable.size_bytes(),
                self.config.memtable_flush_threshold
            );
            self.flush_internal_async().await?;
        }

        Ok(())
    }

    /// Execute a CQL statement (INSERT, UPDATE, DELETE)
    ///
    /// This parses the CQL statement and converts it to a mutation,
    /// then writes it using the `write()` method.
    ///
    /// # Arguments
    ///
    /// * `statement` - CQL statement string
    ///
    /// # Returns
    ///
    /// Ok(()) on success, or an error if parsing or writing fails.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - CQL parsing fails
    /// - Statement is not a mutation (INSERT/UPDATE/DELETE)
    /// - Mutation conversion fails
    /// - Write fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')")?;
    /// engine.execute("UPDATE users SET name = 'Bob' WHERE id = 1")?;
    /// engine.execute("DELETE FROM users WHERE id = 1")?;
    /// ```
    pub fn execute(&mut self, statement: &str) -> Result<()> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(Error::InvalidInput(
                "WriteEngine has been closed".to_string(),
            ));
        }

        let trimmed = statement.trim();

        // BATCH statements produce multiple mutations
        if trimmed.len() >= 5 && trimmed.as_bytes()[..5].eq_ignore_ascii_case(b"BEGIN") {
            let mutations =
                cql_to_mutation::convert_cql_to_mutations(trimmed, &self.config.schema)?;
            for mutation in mutations {
                self.write(mutation)?;
            }
            Ok(())
        } else {
            let mutation = self.parse_cql_to_mutation(statement)?;
            self.write(mutation)
        }
    }

    /// Force a flush of the memtable to SSTable
    ///
    /// This writes all data in the memtable to a new SSTable generation,
    /// then truncates the WAL. The memtable is cleared after a successful flush.
    ///
    /// # Returns
    ///
    /// Returns `Some(SSTableInfo)` if data was flushed, or `None` if the
    /// memtable was empty.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Engine has been closed
    /// - SSTable write fails
    /// - WAL truncate fails
    pub async fn flush(&mut self) -> Result<Option<SSTableInfo>> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(Error::InvalidInput(
                "WriteEngine has been closed".to_string(),
            ));
        }

        self.flush_internal_async().await
    }

    /// Internal synchronous flush helper.
    ///
    /// Bridges to the async flush via [`merge::block_on_async`], which is safe to
    /// call whether or not a Tokio runtime is already running on this thread
    /// (Issue #587).
    fn flush_internal(&mut self) -> Result<()> {
        merge::block_on_async(self.flush_internal_async())?;
        Ok(())
    }

    /// Internal async flush implementation
    async fn flush_internal_async(&mut self) -> Result<Option<SSTableInfo>> {
        // Check if memtable is empty
        if self.memtable.is_empty() {
            return Ok(None);
        }

        log::info!(
            "Flushing memtable: {} partitions, {} rows, {} bytes",
            self.memtable.iter().count(),
            self.memtable.row_count(),
            self.memtable.size_bytes()
        );

        // Create SSTable writer with hint for Bloom filter sizing
        let partition_count_hint = self.memtable.iter().count();
        let mut writer = crate::storage::sstable::writer::SSTableWriter::with_expected_partitions(
            self.config.data_dir.clone(),
            self.generation,
            &self.config.schema,
            partition_count_hint,
        )?;

        // Write all partitions from memtable (already in token order)
        for (decorated_key, mutations) in self.memtable.iter() {
            writer.write_partition(decorated_key.clone(), mutations.to_vec())?;
        }

        // Finalize SSTable
        let info = writer.finish().await?;

        log::info!(
            "SSTable flush complete: generation {}, {} partitions, {} bytes",
            self.generation,
            info.partition_count,
            info.data_size
        );

        // Truncate WAL (data now persisted to SSTable)
        // If truncate fails, log warning but don't fail - data is already in SSTable
        if let Err(e) = self.wal.truncate() {
            log::warn!(
                "Failed to truncate WAL after successful SSTable flush: {}. \
                Data is safe in SSTable, but WAL cleanup failed.",
                e
            );
            // Don't return error - SSTable write succeeded, which is the important part
        }

        // Clear memtable
        self.memtable.clear();

        // Increment generation for next flush
        self.generation += 1;

        Ok(Some(info))
    }

    /// Close the write engine
    ///
    /// This flushes any remaining data in the memtable to SSTable,
    /// syncs the WAL, then marks the engine as closed. After calling close(),
    /// the engine cannot be used for further writes.
    ///
    /// This method is idempotent - calling it multiple times is safe.
    ///
    /// # Returns
    ///
    /// Ok(()) on success.
    ///
    /// # Errors
    ///
    /// Returns an error if the final flush fails. If the WAL truncate fails
    /// after a successful SSTable write, a warning is logged but no error
    /// is returned (the data is already persisted).
    pub async fn close(&mut self) -> Result<()> {
        // Check if already closed (idempotent)
        if self.closed.swap(true, Ordering::SeqCst) {
            return Ok(());
        }

        log::info!("Closing WriteEngine");

        // Flush any remaining data
        if !self.memtable.is_empty() {
            log::info!("Flushing memtable before close");

            // Attempt to flush to SSTable
            match self.flush_internal_async().await {
                Ok(_) => {
                    log::info!("Memtable flushed successfully");
                }
                Err(e) => {
                    // If flush fails, log error and return it
                    log::error!("Failed to flush memtable during close: {}", e);
                    // Reset closed flag since we failed to close cleanly
                    self.closed.store(false, Ordering::SeqCst);
                    return Err(e);
                }
            }
        }

        // Sync WAL before closing
        if let Err(e) = self.wal.sync() {
            log::warn!("Failed to sync WAL during close: {}", e);
            // Don't fail close if sync fails - data is already persisted to SSTable
        }

        log::info!("WriteEngine closed");

        Ok(())
    }

    /// Get the current memtable size in bytes
    pub fn memtable_size(&self) -> usize {
        self.memtable.size_bytes()
    }

    /// Get the current memtable row count
    pub fn memtable_row_count(&self) -> usize {
        self.memtable.row_count()
    }

    /// Get the current WAL size in bytes
    pub fn wal_size(&self) -> u64 {
        self.wal.size()
    }

    /// Get the current generation number
    pub fn generation(&self) -> u64 {
        self.generation
    }

    /// Parse a CQL statement to a Mutation
    ///
    /// Supports INSERT, UPDATE, and DELETE statements.
    fn parse_cql_to_mutation(&self, statement: &str) -> Result<Mutation> {
        cql_to_mutation::convert_cql_to_mutation(statement, &self.config.schema)
    }

    /// Determine the next SSTable generation number
    ///
    /// Scans the data directory for existing SSTable files and returns
    /// the next generation number.
    fn determine_next_generation(data_dir: &Path) -> Result<u64> {
        let mut max_generation = 0u64;

        if !data_dir.exists() {
            return Ok(1);
        }

        // Recursively scan for SSTable files (writer places them in keyspace/table/ subdirs)
        Self::scan_generations(
            data_dir,
            &mut max_generation,
            crate::storage::sstable::MAX_SSTABLE_SCAN_DEPTH,
        )?;

        Ok(max_generation + 1)
    }

    /// Recursively scan directory for SSTable generation numbers
    fn scan_generations(dir: &Path, max_generation: &mut u64, depth: usize) -> Result<()> {
        for entry in std::fs::read_dir(dir)
            .map_err(|e| Error::Storage(format!("Failed to read data directory: {}", e)))?
        {
            let entry = entry
                .map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;

            let filename = entry.file_name();
            let filename_str = filename.to_string_lossy();

            // Parse generation from filename: nb-{generation}-big-{Component}.db
            if filename_str.starts_with("nb-") && filename_str.contains("-big-") {
                if let Some(gen_str) = filename_str
                    .strip_prefix("nb-")
                    .and_then(|s| s.split('-').next())
                {
                    if let Ok(gen) = gen_str.parse::<u64>() {
                        *max_generation = (*max_generation).max(gen);
                    }
                }
            } else if depth > 0 {
                let path = entry.path();
                if path.is_dir() {
                    Self::scan_generations(&path, max_generation, depth - 1)?;
                }
            }
        }
        Ok(())
    }

    /// Set the merge policy for background compaction (M5.2, Issue #383)
    ///
    /// # Arguments
    ///
    /// * `policy` - Merge policy implementation (e.g., STCS, LCS, TWCS)
    pub fn set_merge_policy(&mut self, policy: Box<dyn MergePolicy>) -> Result<()> {
        self.merge_policy = Some(policy);
        Ok(())
    }

    /// Return cumulative compaction statistics (M5.2, Issue #474)
    ///
    /// Returns a snapshot of the lifetime totals accumulated across all compaction
    /// cycles that have completed since the `WriteEngine` was created. The snapshot
    /// is cheaply cloneable and safe to inspect from any thread (no lock required,
    /// because `WriteEngine` itself is not `Sync`).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let stats = engine.maintenance_stats();
    /// println!(
    ///     "Completed {} compactions, merged {} rows, wrote {} bytes",
    ///     stats.compactions_completed,
    ///     stats.rows_merged,
    ///     stats.bytes_written,
    /// );
    /// ```
    pub fn maintenance_stats(&self) -> CompactionStats {
        self.cumulative_stats.clone()
    }

    /// Perform incremental maintenance work (M5.2, Issue #384)
    ///
    /// This method performs background compaction work within a time budget.
    /// It can be called repeatedly from a background thread or task scheduler
    /// to make incremental progress on compaction.
    ///
    /// ## Runtime contexts
    ///
    /// This is a synchronous method, but its internal async-to-sync bridge is
    /// runtime-aware (see [`merge::block_on_async`]), so it is safe to call from
    /// **either** a plain synchronous context **or** from within an active Tokio
    /// runtime — including `#[tokio::main]`/`#[tokio::test]` worker threads and
    /// `async fn` callers. Prior to Issue #587 calling it from inside a runtime
    /// panicked with "Cannot start a runtime from within a runtime" once a merge
    /// had input SSTables to read. The sync signature is preserved so the CLI and
    /// Python bindings can keep calling it directly. (The Node binding wraps it in
    /// `spawn_blocking`, which remains correct.)
    ///
    /// ## Behavior
    ///
    /// 1. If no active merge exists, consult the merge policy for work
    /// 2. If merge work is available, start a new merge
    /// 3. Process the active merge until budget is exhausted
    /// 4. Return progress report
    ///
    /// ## Invariants
    ///
    /// - Budget is honored within 10% tolerance
    /// - At least one partition is processed per call (minimum progress guarantee)
    /// - Merge state is preserved across calls for resumption
    ///
    /// ## Budget Enforcement
    ///
    /// The budget is honored within approximately 10% tolerance. This tolerance
    /// exists to avoid interrupting partition processing mid-stream, which would
    /// require complex state management to resume. The tolerance ensures forward
    /// progress on each call while remaining responsive to time constraints.
    ///
    /// # Arguments
    ///
    /// * `budget` - Maximum time to spend in this call
    ///
    /// # Returns
    ///
    /// A report containing progress metrics and whether more work is pending.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Engine has been closed
    /// - Merge policy returns an error
    /// - SSTable reading or writing fails
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use std::time::Duration;
    ///
    /// // Background compaction loop
    /// loop {
    ///     let report = engine.maintenance_step(Duration::from_millis(100))?;
    ///
    ///     if !report.pending_compaction {
    ///         // No more work, sleep or exit
    ///         break;
    ///     }
    ///
    ///     // Log progress
    ///     println!("Merged {} rows in {:?}", report.rows_merged, report.time_spent);
    /// }
    /// ```
    pub fn maintenance_step(&mut self, budget: Duration) -> Result<MaintenanceReport> {
        if self.closed.load(Ordering::SeqCst) {
            return Err(Error::InvalidInput(
                "WriteEngine has been closed".to_string(),
            ));
        }

        let start = Instant::now();
        let mut report = MaintenanceReport {
            time_spent: Duration::from_secs(0),
            completed_merges: Vec::new(),
            rows_merged: 0,
            bytes_written: 0,
            pending_compaction: false,
        };

        // If no merge policy is set, no maintenance work to do
        let merge_policy = match &self.merge_policy {
            Some(policy) => policy,
            None => {
                report.time_spent = start.elapsed();
                return Ok(report);
            }
        };

        // If no active merge exists, check if we should start one
        if self.active_merge.is_none() {
            let candidates = self.scan_sstable_candidates()?;
            let selected = merge_policy.select_merge(&candidates)?;

            if !selected.is_empty() {
                // Start a new merge
                self.start_merge(selected)?;
            } else {
                // No work selected by policy
                report.time_spent = start.elapsed();
                report.pending_compaction = false;
                return Ok(report);
            }
        }

        // Process active merge within budget
        let budget_tolerance = budget.mul_f32(1.1); // 10% tolerance
        let mut partitions_processed = 0;

        while let Some(merge) = &mut self.active_merge {
            // Check budget (but always process at least one partition)
            if partitions_processed > 0 && start.elapsed() >= budget_tolerance {
                break;
            }

            // Process one partition from the merge
            let step = merge.merger.step()?;

            match step {
                merge::MergeStep::Partition { key, rows } => {
                    partitions_processed += 1;
                    let row_count = rows.len() as u64;

                    // Convert MergeEntry rows to Mutation format
                    // (collect into a vec first to release the borrow on merge)
                    let entries_vec: Vec<_> = rows.into_iter().collect();

                    // Now we can call self methods without conflict
                    let mutations = entries_vec
                        .into_iter()
                        .map(|entry| self.merge_entry_to_mutation(entry))
                        .collect::<Result<Vec<_>>>()?;

                    // Write partition to output SSTable
                    // Re-borrow active_merge to write
                    if let Some(merge) = &mut self.active_merge {
                        merge.writer.write_partition(key, mutations)?;
                        merge.rows_merged += row_count;
                    }

                    // Update stats
                    report.rows_merged += row_count;
                }
                merge::MergeStep::Complete => {
                    // Merge is complete - finalize and clean up
                    // Use blocking call to handle async finalization
                    self.finalize_merge_blocking(&mut report)?;
                    break;
                }
            }
        }

        // Check if more work is pending
        report.pending_compaction = self.active_merge.is_some();
        report.time_spent = start.elapsed();

        Ok(report)
    }

    /// Scan data directory for SSTable candidates (M5.2 helper)
    /// Startup sweep (a): remove any `.compaction-tmp-*` directories left under
    /// `data_dir` by a previous crash mid-rename.  Best-effort — individual
    /// failures are logged but do not abort engine startup.
    fn sweep_orphaned_compaction_tmp(data_dir: &Path) {
        let read_dir = match std::fs::read_dir(data_dir) {
            Ok(rd) => rd,
            Err(e) => {
                log::debug!(
                    "sweep_orphaned_compaction_tmp: cannot read {:?}: {}",
                    data_dir,
                    e
                );
                return;
            }
        };

        for entry in read_dir.flatten() {
            let path = entry.path();
            let name = entry.file_name();
            let name_str = name.to_string_lossy();
            if name_str.starts_with(".compaction-tmp-") && path.is_dir() {
                log::warn!("removing orphaned compaction tmp directory: {:?}", path);
                if let Err(e) = std::fs::remove_dir_all(&path) {
                    log::warn!(
                        "failed to remove orphaned compaction tmp directory {:?}: {}",
                        path,
                        e
                    );
                }
            }
        }
    }

    /// Startup sweep (b): remove any `nb-{gen}-big-Data.db` (and its siblings)
    /// under `data_dir/keyspace/table/` that lack a matching `TOC.txt`.
    /// Such files are left when a crash occurs after some component renames but
    /// before `TOC.txt` is renamed (the publication barrier).  Best-effort.
    fn sweep_orphaned_partial_sstables(data_dir: &Path, keyspace: &str, table: &str) {
        let sstable_dir = data_dir.join(keyspace).join(table);

        let read_dir = match std::fs::read_dir(&sstable_dir) {
            Ok(rd) => rd,
            Err(_) => {
                // Directory doesn't exist yet — nothing to sweep.
                return;
            }
        };

        for entry in read_dir.flatten() {
            let path = entry.path();
            let name = entry.file_name();
            let name_str = name.to_string_lossy();

            // Look for Data.db files produced by the writer (nb-{gen}-big-Data.db)
            if !name_str.starts_with("nb-")
                || !name_str.ends_with("-big-Data.db")
                || !path.is_file()
            {
                continue;
            }

            // Extract the base prefix (e.g. "nb-5-big") to find the TOC sibling
            let base = match name_str.strip_suffix("-Data.db") {
                Some(b) => b.to_owned(),
                None => continue,
            };

            // Extract the generation number for the log message
            let gen_str = base
                .strip_prefix("nb-")
                .and_then(|s| s.strip_suffix("-big"))
                .unwrap_or(&base);

            let toc_path = sstable_dir.join(format!("{}-TOC.txt", base));
            if !toc_path.exists() {
                log::warn!(
                    "removing orphaned partial SSTable components for generation {}: missing TOC.txt",
                    gen_str
                );
                if let Err(e) = Self::delete_sstable_files_static(&path) {
                    log::warn!(
                        "failed to remove orphaned partial SSTable for generation {}: {}",
                        gen_str,
                        e
                    );
                }
            }
        }
    }

    fn scan_sstable_candidates(&self) -> Result<Vec<PathBuf>> {
        let mut candidates = Vec::new();

        if !self.config.data_dir.exists() {
            return Ok(candidates);
        }

        Self::scan_data_files(
            &self.config.data_dir,
            &mut candidates,
            crate::storage::sstable::MAX_SSTABLE_SCAN_DEPTH,
        )?;
        Ok(candidates)
    }

    /// Recursively scan for Data.db files
    fn scan_data_files(dir: &Path, candidates: &mut Vec<PathBuf>, depth: usize) -> Result<()> {
        for entry in std::fs::read_dir(dir)
            .map_err(|e| Error::Storage(format!("Failed to read data directory: {}", e)))?
        {
            let entry = entry
                .map_err(|e| Error::Storage(format!("Failed to read directory entry: {}", e)))?;

            let path = entry.path();
            let filename = path.file_name().unwrap_or_default().to_string_lossy();

            // Only consider Data.db files
            if filename.starts_with("nb-") && filename.ends_with("-big-Data.db") {
                // Honor the TOC.txt publication barrier (Issue #591). A Data.db
                // without a sibling TOC.txt is NOT a published SSTable: it is
                // either a crash-interrupted partial rename or a deferred-delete
                // orphan whose TOC was removed first while its data file stayed
                // pinned by an open/mapped reader (Windows). Feeding such a file
                // to the merger would re-compact an unpublished input and could
                // produce garbled output, so it is skipped here just as the
                // read path discovers SSTables by TOC.txt. The startup orphan
                // sweep reclaims the leftover components.
                let base = filename.trim_end_matches("-Data.db");
                let toc_path = path.with_file_name(format!("{base}-TOC.txt"));
                if toc_path.exists() {
                    candidates.push(path);
                } else {
                    log::debug!(
                        "scan_data_files: skipping unpublished SSTable (no TOC.txt): {:?}",
                        path
                    );
                }
            } else if depth > 0 && path.is_dir() {
                Self::scan_data_files(&path, candidates, depth - 1)?;
            }
        }
        Ok(())
    }

    /// Start a new merge operation (M5.2 helper, Issue #474)
    ///
    /// ## Atomicity design
    ///
    /// The SSTableWriter is pointed at a temporary directory (`tmp_dir`) that lives
    /// alongside the final SSTable directory. After `writer.finish()` succeeds, each
    /// output component is atomically renamed into the final directory. Input files are
    /// deleted only after all renames complete. This guarantees:
    ///
    /// - A crash before renames: tmp files are incomplete; inputs intact.
    /// - A crash mid-rename: at worst a partial output component exists in the final
    ///   directory, but the old inputs are still there and the TOC.txt (publication
    ///   barrier) has not been renamed yet, so the partial output is never discovered
    ///   by readers scanning for `TOC.txt`.
    /// - A crash after all renames but before input deletion: a harmless duplicate
    ///   exists until next compaction.
    fn start_merge(&mut self, input_paths: Vec<PathBuf>) -> Result<()> {
        log::info!(
            "Starting compaction merge of {} SSTables",
            input_paths.len()
        );

        // Measure total bytes read (sum of Data.db file sizes as an approximation)
        let bytes_read: u64 = input_paths
            .iter()
            .map(|p| std::fs::metadata(p).map(|m| m.len()).unwrap_or(0))
            .sum();

        let output_generation = self.generation;

        // Final SSTable directory: data_dir/keyspace/table/
        let sstable_dir = self
            .config
            .data_dir
            .join(&self.config.schema.keyspace)
            .join(&self.config.schema.table);

        // Temporary root: data_dir/.compaction-tmp-{gen}/
        //
        // Placing the tmp root inside `data_dir` (not inside `sstable_dir`) keeps
        // the path simple and guarantees the rename is within the same filesystem.
        // The SSTableWriter appends `keyspace/table/` internally, so component files
        // land at: data_dir/.compaction-tmp-{gen}/keyspace/table/nb-{gen}-big-*.db
        let tmp_dir = self
            .config
            .data_dir
            .join(format!(".compaction-tmp-{}", output_generation));

        // Create the tmp root (SSTableWriter::finish will create the subdirs)
        std::fs::create_dir_all(&tmp_dir).map_err(|e| {
            Error::Storage(format!(
                "Failed to create compaction tmp directory {:?}: {}",
                tmp_dir, e
            ))
        })?;

        // Create K-way merger
        let merger = KWayMerger::new(input_paths.clone(), &self.config.schema)?;

        // Point the SSTableWriter at the tmp root; it will write to
        // tmp_dir/keyspace/table/nb-{gen}-big-*.db
        let writer = crate::storage::sstable::writer::SSTableWriter::new(
            tmp_dir.clone(),
            output_generation,
            &self.config.schema,
        )?;

        // Increment generation for next operation
        self.generation += 1;

        self.active_merge = Some(ActiveMerge {
            merger,
            writer,
            input_paths,
            tmp_dir,
            sstable_dir,
            rows_merged: 0,
            bytes_read,
            started_at: Instant::now(),
        });

        Ok(())
    }

    /// Finalize the active merge - blocking version (M5.2 helper).
    ///
    /// Bridges to the async finalizer via [`merge::block_on_async`], which is
    /// safe to call from within an active Tokio runtime (e.g. the CLI's
    /// `#[tokio::main]` worker threads). A nested `Handle::block_on` here would
    /// otherwise panic with "Cannot start a runtime from within a runtime"
    /// (Issue #587).
    fn finalize_merge_blocking(&mut self, report: &mut MaintenanceReport) -> Result<()> {
        merge::block_on_async(self.finalize_merge_async(report))
    }

    /// Finalize the active merge - async version (M5.2 helper, Issue #474)
    ///
    /// ## Atomicity protocol
    ///
    /// 1. `writer.finish()` flushes all component files to the tmp directory.
    /// 2. Each component file is renamed from `tmp_dir/` to `sstable_dir/`.
    ///    The TOC.txt rename is performed **last** (it is the publication barrier:
    ///    readers discover SSTables by scanning for `TOC.txt`).
    /// 3. Only after all renames succeed are the input SSTable files deleted.
    /// 4. The now-empty tmp directory is removed.
    ///
    /// If any step 2 rename fails, the partially-renamed output components are
    /// cleaned up and an error is returned. The input SSTables remain intact.
    async fn finalize_merge_async(&mut self, report: &mut MaintenanceReport) -> Result<()> {
        let merge = match self.active_merge.take() {
            Some(m) => m,
            None => return Ok(()),
        };

        let elapsed = merge.started_at.elapsed();
        log::info!(
            "Finalizing compaction merge: {} rows, {:?} elapsed",
            merge.rows_merged,
            elapsed
        );

        // Step 1: Finish writing all components to the tmp directory.
        // If this fails the tmp directory may contain partial files, but inputs are safe.
        let tmp_info = match merge.writer.finish().await {
            Ok(info) => info,
            Err(e) => {
                // Clean up tmp directory on failure (best effort)
                let _ = std::fs::remove_dir_all(&merge.tmp_dir);
                return Err(Error::Storage(format!(
                    "Compaction merge write failed (inputs intact): {}",
                    e
                )));
            }
        };

        log::info!(
            "Compaction tmp output: {} bytes, {} partitions",
            tmp_info.data_size,
            tmp_info.partition_count
        );

        // Step 2: Atomically rename each component from the tmp sub-directory to
        // the final SSTable directory. Because both directories are under the same
        // `data_dir`, the rename is within the same filesystem (POSIX atomic).
        // We rename TOC.txt last because it is the publication barrier.
        let sstable_dir = &merge.sstable_dir;

        // Ensure the final SSTable directory exists (it normally does, but handle
        // the edge case where it was created by start_merge only in the tmp path).
        std::fs::create_dir_all(sstable_dir).map_err(|e| {
            Error::Storage(format!(
                "Failed to create SSTable directory {:?}: {}",
                sstable_dir, e
            ))
        })?;

        // Helper: map a tmp component path to its final destination.
        // tmp_info paths look like: data_dir/.compaction-tmp-N/keyspace/table/nb-N-big-X.db
        // Final destination:        data_dir/keyspace/table/nb-N-big-X.db
        let make_rename = |src: &PathBuf| -> Result<(PathBuf, PathBuf)> {
            let filename = src
                .file_name()
                .ok_or_else(|| Error::Storage("Component path has no filename".to_string()))?;
            let dst = sstable_dir.join(filename);
            Ok((src.clone(), dst))
        };

        // Build list of (src, dst) renames. TOC.txt goes last (publication barrier).
        let mut renames: Vec<(PathBuf, PathBuf)> = Vec::new();

        // Non-TOC components first
        for src in &[
            &tmp_info.data_path,
            &tmp_info.index_path,
            &tmp_info.filter_path,
            &tmp_info.summary_path,
            &tmp_info.stats_path,
            &tmp_info.digest_path,
        ] {
            renames.push(make_rename(src)?);
        }
        if let Some(ref ci_path) = tmp_info.compression_info_path {
            renames.push(make_rename(ci_path)?);
        }
        // TOC.txt is last (publication barrier)
        renames.push(make_rename(&tmp_info.toc_path)?);

        // Perform the renames. On failure, remove any already-renamed files so
        // we don't leave a half-published SSTable, then return the error.
        let mut renamed: Vec<PathBuf> = Vec::with_capacity(renames.len());
        let mut rename_error: Option<Error> = None;

        for (src, dst) in &renames {
            match std::fs::rename(src, dst) {
                Ok(()) => {
                    log::debug!(
                        "Renamed {:?} → {:?}",
                        src.file_name().unwrap_or_default(),
                        dst.file_name().unwrap_or_default()
                    );
                    renamed.push(dst.clone());
                }
                Err(e) => {
                    rename_error = Some(Error::Storage(format!(
                        "Atomic rename of {:?} to {:?} failed (rolling back, inputs intact): {}",
                        src, dst, e
                    )));
                    break;
                }
            }
        }

        if let Some(err) = rename_error {
            // Roll back already-renamed files (best effort)
            for dst in &renamed {
                let _ = std::fs::remove_file(dst);
            }
            // Clean up tmp directory
            let _ = std::fs::remove_dir_all(&merge.tmp_dir);
            return Err(err);
        }

        // Step 3: All renames succeeded. The new SSTable is now visible.
        // Delete input SSTable files. If deletion fails we log a warning but do
        // NOT return an error — the merge output is correct.
        //
        // Issue #591 (write-while-mapped / Windows policy): the inputs were read
        // through buffered I/O and fully drained into memory by `KWayMerger::new`
        // before this point, so the merger holds no mapping over them — deleting
        // them cannot fault with SIGBUS. `delete_sstable_files` removes each
        // input's TOC.txt first (unpublishing it) and is best-effort on the data
        // components, so a component still pinned by a concurrent mapped reader on
        // Windows becomes an invisible orphan reclaimed on the next startup rather
        // than a hard failure or a source of duplicate rows.
        for input_path in &merge.input_paths {
            if let Err(e) = self.delete_sstable_files(input_path) {
                log::warn!(
                    "Failed to delete compaction input {:?}: {} \
                     (merge output is valid; inputs will be re-evaluated next cycle)",
                    input_path,
                    e
                );
            }
        }

        // Step 4: Remove the now-empty tmp directory (best effort).
        if let Err(e) = std::fs::remove_dir_all(&merge.tmp_dir) {
            log::debug!(
                "Failed to remove compaction tmp directory {:?}: {}",
                merge.tmp_dir,
                e
            );
        }

        // The final Data.db path is in sstable_dir (renamed from tmp)
        let final_data_path = sstable_dir.join(
            tmp_info
                .data_path
                .file_name()
                .ok_or_else(|| Error::Storage("Data.db path has no filename".to_string()))?,
        );

        // Compute total bytes written across ALL output SSTable components (not just Data.db).
        // We stat the final paths (post-rename) so we measure what was actually persisted.
        let total_bytes_written: u64 = [
            &tmp_info.data_path,
            &tmp_info.index_path,
            &tmp_info.filter_path,
            &tmp_info.summary_path,
            &tmp_info.stats_path,
            &tmp_info.digest_path,
        ]
        .iter()
        .map(|p| {
            let filename = p.file_name().unwrap_or_default();
            let final_path = sstable_dir.join(filename);
            std::fs::metadata(&final_path).map(|m| m.len()).unwrap_or(0)
        })
        .sum::<u64>()
            + tmp_info
                .compression_info_path
                .as_ref()
                .and_then(|p| {
                    let filename = p.file_name()?;
                    std::fs::metadata(sstable_dir.join(filename))
                        .ok()
                        .map(|m| m.len())
                })
                .unwrap_or(0);

        // Update per-step report
        report.completed_merges.push(final_data_path);
        report.bytes_written += total_bytes_written;

        // Update cumulative lifetime stats
        self.cumulative_stats.compactions_completed += 1;
        self.cumulative_stats.sstables_merged_in += merge.input_paths.len() as u64;
        self.cumulative_stats.sstables_produced += 1;
        self.cumulative_stats.bytes_read += merge.bytes_read;
        self.cumulative_stats.bytes_written += total_bytes_written;
        self.cumulative_stats.rows_merged += merge.rows_merged;
        self.cumulative_stats.total_time += elapsed;

        log::info!(
            "Compaction complete: merged {} inputs → 1 output ({} bytes total across all components, {} rows, {:?})",
            merge.input_paths.len(),
            total_bytes_written,
            merge.rows_merged,
            elapsed
        );

        Ok(())
    }

    /// Delete all component files for an SSTable (M5.2 helper)
    fn delete_sstable_files(&self, data_path: &Path) -> Result<()> {
        Self::delete_sstable_files_static(data_path)
    }

    /// Static helper that deletes all component files for an SSTable given the
    /// Data.db path.  Called from both `delete_sstable_files` and the startup
    /// orphan sweep, which runs before `self` is fully constructed.
    ///
    /// ## Deferred-delete / Windows policy (Issue #591)
    ///
    /// `TOC.txt` is removed **first**. TOC.txt is the publication barrier — both
    /// the read path (`SSTableManager`) and the compaction candidate scan
    /// (`scan_data_files`, since #591) treat a Data.db without a sibling TOC.txt
    /// as unpublished. Removing TOC.txt first therefore *unpublishes* the SSTable
    /// atomically, before any data component is touched, so it can never be
    /// observed (no duplicate rows, never re-fed to the merger) even if the
    /// remaining components cannot be removed yet.
    ///
    /// The remaining components are then deleted **best-effort**: a failure on
    /// any one of them (most plausibly a Windows sharing violation when a
    /// concurrent reader still has the file open or memory-mapped) is logged but
    /// does NOT abort the rest or fail the operation. Such a leftover is a
    /// harmless orphan — invisible because its TOC.txt is gone — and is reclaimed
    /// by [`Self::sweep_orphaned_partial_sstables`] on the next engine startup,
    /// by which time the reader's handle has been released. This is the
    /// "deferred delete" half of the policy; Unix removes the inode immediately
    /// while any mapping keeps the bytes alive until it is dropped.
    fn delete_sstable_files_static(data_path: &Path) -> Result<()> {
        // Extract base path: nb-{gen}-big
        let filename = data_path
            .file_name()
            .and_then(|s| s.to_str())
            .ok_or_else(|| Error::Storage("Invalid SSTable path".to_string()))?;

        let base = filename
            .strip_suffix("-Data.db")
            .ok_or_else(|| Error::Storage("Invalid Data.db filename".to_string()))?;

        let parent_dir = data_path.parent().ok_or_else(|| {
            Error::Storage(format!(
                "Data.db path has no parent directory: {:?}",
                data_path
            ))
        })?;

        // TOC.txt FIRST — the publication barrier (Issue #591). Once it is gone
        // the SSTable is unpublished regardless of whether the data components
        // can be removed. Remaining components follow, best-effort.
        let components = [
            "TOC.txt",
            "Data.db",
            "Index.db",
            "Summary.db",
            "Statistics.db",
            "CompressionInfo.db",
            "Filter.db",
            "Digest.crc32",
        ];

        let mut failures: Vec<String> = Vec::new();
        for component in &components {
            let component_path = parent_dir.join(format!("{}-{}", base, component));
            if component_path.exists() {
                match std::fs::remove_file(&component_path) {
                    Ok(()) => log::debug!("Deleted compaction input: {:?}", component_path),
                    Err(e) => {
                        // Best-effort: do not abort. A leftover data component
                        // whose TOC.txt is already gone is an invisible orphan
                        // reclaimed by the startup sweep (Issue #591).
                        log::warn!(
                            "Deferred delete of {:?}: {} (component left as orphan; \
                             unpublished via TOC.txt removal, reclaimed on next startup)",
                            component_path,
                            e
                        );
                        failures.push(format!("{:?}: {}", component_path, e));
                    }
                }
            }
        }

        if failures.is_empty() {
            Ok(())
        } else {
            // Surface a non-fatal error so callers can log it. The SSTable is
            // already unpublished (TOC.txt removed first), so callers treat this
            // as a deferred reclamation, not a correctness failure.
            Err(Error::Storage(format!(
                "Deferred delete left {} orphaned component(s) (unpublished, reclaimed on \
                 next startup): {}",
                failures.len(),
                failures.join("; ")
            )))
        }
    }

    /// Convert MergeEntry to Mutation (M5.2 helper)
    ///
    /// Delegates to `KWayMerger::merge_entry_to_mutation` to avoid duplication.
    fn merge_entry_to_mutation(
        &self,
        entry: merge::MergeEntry,
    ) -> Result<crate::storage::write_engine::mutation::Mutation> {
        merge::KWayMerger::merge_entry_to_mutation(entry, &self.config.schema)
    }
}

#[cfg(all(test, feature = "write-support"))]
mod tests {
    use super::*;
    use crate::schema::{Column, KeyColumn};
    use crate::storage::write_engine::mutation::{CellOperation, PartitionKey, TableId};
    use crate::types::Value;
    use std::collections::HashMap;
    use tempfile::TempDir;

    fn create_test_schema() -> TableSchema {
        TableSchema {
            keyspace: "test_ks".to_string(),
            table: "test_table".to_string(),
            partition_keys: vec![KeyColumn {
                name: "id".to_string(),
                data_type: "int".to_string(),
                position: 0,
            }],
            clustering_keys: vec![],
            columns: vec![
                Column {
                    name: "id".to_string(),
                    data_type: "int".to_string(),
                    nullable: false,
                    default: None,
                    is_static: false,
                },
                Column {
                    name: "name".to_string(),
                    data_type: "text".to_string(),
                    nullable: true,
                    default: None,
                    is_static: false,
                },
            ],
            comments: HashMap::new(),
        }
    }

    fn create_test_mutation(id: i32, name: &str, timestamp: i64) -> Mutation {
        let table_id = TableId::new("test_ks", "test_table");
        let pk = PartitionKey::single("id", Value::Integer(id));
        let ops = vec![CellOperation::Write {
            column: "name".to_string(),
            value: Value::Text(name.to_string()),
        }];

        Mutation::new(table_id, pk, None, ops, timestamp, None)
    }

    #[test]
    fn test_set_merge_policy() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Should succeed now (was previously returning error)
        let policy = Box::new(crate::storage::write_engine::STCSPolicy::default());
        engine.set_merge_policy(policy).unwrap();

        // With policy set but no SSTables, should return quickly with no work
        let report = engine
            .maintenance_step(std::time::Duration::from_millis(100))
            .unwrap();
        assert!(!report.pending_compaction);
        assert_eq!(report.rows_merged, 0);
    }

    #[test]
    fn test_write_engine_config() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        assert_eq!(
            config.memtable_flush_threshold,
            WriteEngineConfig::DEFAULT_FLUSH_THRESHOLD
        );
        assert_eq!(
            config.memtable_hard_limit,
            WriteEngineConfig::DEFAULT_HARD_LIMIT
        );

        let config = config.with_flush_threshold(128 * 1024 * 1024);
        assert_eq!(config.memtable_flush_threshold, 128 * 1024 * 1024);

        let config = config.with_hard_limit(512 * 1024 * 1024);
        assert_eq!(config.memtable_hard_limit, 512 * 1024 * 1024);
    }

    #[test]
    fn test_write_engine_new() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        assert_eq!(engine.generation(), 1);
        assert_eq!(engine.memtable_size(), 0);
        assert_eq!(engine.memtable_row_count(), 0);
        assert!(!engine.closed.load(std::sync::atomic::Ordering::Relaxed));
    }

    #[test]
    fn test_write_engine_write_single_mutation() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        let mutation = create_test_mutation(1, "Alice", 1000000);
        engine.write(mutation).unwrap();

        assert_eq!(engine.memtable_row_count(), 1);
        assert!(engine.memtable_size() > 0);
        assert!(engine.wal_size() > 0);
    }

    #[test]
    fn test_write_engine_write_multiple_mutations() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Write 10 mutations
        for i in 0..10 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            engine.write(mutation).unwrap();
        }

        assert_eq!(engine.memtable_row_count(), 10);
        assert!(engine.memtable_size() > 0);
    }

    #[tokio::test]
    async fn test_write_engine_flush_empty() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Flush empty memtable
        let result = engine.flush().await.unwrap();
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_write_engine_flush_with_data() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Write mutations
        for i in 0..5 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            engine.write(mutation).unwrap();
        }

        let initial_generation = engine.generation();

        // Flush
        let info = engine.flush().await.unwrap();
        assert!(info.is_some());

        let info = info.unwrap();
        assert_eq!(info.partition_count, 5);
        assert!(info.data_size > 0);
        assert!(info.data_path.exists());

        // Memtable should be empty after flush
        assert_eq!(engine.memtable_row_count(), 0);
        assert_eq!(engine.memtable_size(), 0);

        // WAL should be truncated
        assert_eq!(engine.wal_size(), 0);

        // Generation should increment
        assert_eq!(engine.generation(), initial_generation + 1);
    }

    #[test]
    fn test_write_engine_automatic_flush() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        // Set very low flush threshold (1KB)
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_flush_threshold(1024);

        let mut engine = WriteEngine::new(config).unwrap();

        // Write enough mutations to trigger automatic flush
        for i in 0..100 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            engine.write(mutation).unwrap();
        }

        // Should have automatically flushed
        // Memtable may have some data if writes continued after flush
        // But generation should have incremented
        assert!(engine.generation() > 1 || engine.memtable_size() < 10000);
    }

    #[tokio::test]
    async fn test_write_engine_close_with_data() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Write mutations
        for i in 0..5 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            engine.write(mutation).unwrap();
        }

        // Close should flush
        engine.close().await.unwrap();

        // Verify SSTable was created
        let data_dir = temp_dir.path().join("data");
        let entries: Vec<_> = std::fs::read_dir(&data_dir).unwrap().collect();
        assert!(!entries.is_empty(), "SSTable files should exist");
    }

    #[tokio::test]
    async fn test_write_engine_close_empty() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close empty engine
        engine.close().await.unwrap();
    }

    #[test]
    fn test_write_engine_write_after_close() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close
        tokio::runtime::Runtime::new()
            .unwrap()
            .block_on(engine.close())
            .unwrap();

        // Create new engine with same config (simulating restart)
        let schema2 = create_test_schema();
        let config2 = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema2,
        );

        let mut engine2 = WriteEngine::new(config2).unwrap();

        // Write should fail on closed engine (if we still had reference)
        // But new engine should work
        let mutation = create_test_mutation(1, "Alice", 1000000);
        engine2.write(mutation).unwrap();
        assert_eq!(engine2.memtable_row_count(), 1);
    }

    #[test]
    fn test_write_engine_wal_recovery() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema.clone(),
        );

        // Write mutations and close without flushing
        {
            let mut engine = WriteEngine::new(config.clone()).unwrap();

            for i in 0..5 {
                let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
                engine.write(mutation).unwrap();
            }

            // Don't flush - just drop engine (simulating crash)
        }

        // Create new engine - should recover from WAL
        let config2 = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config2).unwrap();

        // Should have recovered 5 mutations
        assert_eq!(engine.memtable_row_count(), 5);
        assert!(engine.memtable_size() > 0);
    }

    #[test]
    fn test_write_engine_generation_tracking() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema.clone(),
        );

        // First engine
        {
            let mut engine = WriteEngine::new(config.clone()).unwrap();
            assert_eq!(engine.generation(), 1);

            // Write and flush
            let mutation = create_test_mutation(1, "Alice", 1000000);
            engine.write(mutation).unwrap();

            tokio::runtime::Runtime::new()
                .unwrap()
                .block_on(engine.flush())
                .unwrap();

            assert_eq!(engine.generation(), 2);
        }

        // Second engine - should detect existing generation
        let config2 = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config2).unwrap();
        assert_eq!(engine.generation(), 2);
    }

    #[test]
    fn test_write_engine_execute_table_mismatch() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Schema defines test_table, but statement targets users → table mismatch
        let result = engine.execute("INSERT INTO users (id, name) VALUES (1, 'Alice')");
        let err_msg = result.unwrap_err().to_string();
        assert!(
            err_msg.contains("targets table 'users'")
                && err_msg.contains("schema is for 'test_table'"),
            "Expected table mismatch error, got: {}",
            err_msg
        );
    }

    #[test]
    fn test_write_engine_execute_insert_success() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        assert_eq!(engine.memtable_row_count(), 0);

        // INSERT matching the test schema: test_ks.test_table(id int PK, name text)
        let result = engine.execute("INSERT INTO test_table (id, name) VALUES (1, 'Alice')");
        assert!(
            result.is_ok(),
            "execute() failed: {:?}",
            result.unwrap_err()
        );

        assert_eq!(engine.memtable_row_count(), 1);
    }

    #[test]
    fn test_determine_next_generation_empty_dir() {
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
        assert_eq!(generation, 1);
    }

    #[test]
    fn test_determine_next_generation_with_sstables() {
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        // Create dummy SSTable files
        std::fs::write(data_dir.join("nb-1-big-Data.db"), b"").unwrap();
        std::fs::write(data_dir.join("nb-2-big-Data.db"), b"").unwrap();
        std::fs::write(data_dir.join("nb-5-big-Data.db"), b"").unwrap();

        let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
        assert_eq!(generation, 6);
    }

    #[tokio::test]
    async fn test_write_engine_close_idempotent() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close once
        engine.close().await.unwrap();
        assert!(engine.closed.load(Ordering::SeqCst));

        // Close again - should be idempotent
        engine.close().await.unwrap();
        assert!(engine.closed.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_write_engine_close_syncs_wal() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Write a mutation
        let mutation = create_test_mutation(1, "Alice", 1000000);
        engine.write(mutation).unwrap();

        // Close should sync WAL before completing
        engine.close().await.unwrap();

        // Verify WAL was truncated (because data was flushed to SSTable)
        assert_eq!(engine.wal_size(), 0);
    }

    #[test]
    fn test_write_engine_closed_flag_atomic() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        // Verify closed flag is atomic
        assert!(!engine.closed.load(Ordering::SeqCst));

        // Store true
        engine.closed.store(true, Ordering::SeqCst);
        assert!(engine.closed.load(Ordering::SeqCst));

        // Swap back to false
        let prev = engine.closed.swap(false, Ordering::SeqCst);
        assert!(prev);
        assert!(!engine.closed.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_write_engine_write_after_close_fails() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close the engine
        engine.close().await.unwrap();

        // Try to write - should fail
        let mutation = create_test_mutation(1, "Alice", 1000000);
        let result = engine.write(mutation);

        assert!(result.is_err());
        match result {
            Err(Error::InvalidInput(_)) => {}
            _ => panic!("Expected InvalidInput error"),
        }
    }

    #[tokio::test]
    async fn test_write_engine_flush_after_close_fails() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close the engine
        engine.close().await.unwrap();

        // Try to flush - should fail
        let result = engine.flush().await;

        assert!(result.is_err());
        match result {
            Err(Error::InvalidInput(_)) => {}
            _ => panic!("Expected InvalidInput error"),
        }
    }

    #[test]
    fn test_write_engine_hard_limit_enforcement() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        // Set very low hard limit (2KB) with flush threshold higher (to prevent auto-flush)
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_flush_threshold(10 * 1024) // 10KB flush threshold (higher than hard limit for test)
        .with_hard_limit(2048); // 2KB hard limit

        let mut engine = WriteEngine::new(config).unwrap();

        // Write mutations until we hit the hard limit
        let mut write_count = 0;
        for i in 0..1000 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            let result = engine.write(mutation);

            match result {
                Ok(()) => {
                    write_count += 1;
                }
                Err(Error::Storage(msg)) => {
                    assert!(msg.contains("hard limit"));
                    break;
                }
                Err(e) => panic!("Expected Storage error, got: {:?}", e),
            }
        }

        // Should have stopped before 1000 writes due to hard limit
        assert!(
            write_count < 1000,
            "Should have hit hard limit before 1000 writes"
        );
        assert!(
            write_count > 0,
            "Should have accepted at least some writes before hitting limit"
        );
    }

    #[tokio::test]
    async fn test_write_engine_hard_limit_enforcement_async() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        // Set very low hard limit (2KB) with flush threshold higher (to prevent auto-flush)
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_flush_threshold(10 * 1024) // 10KB flush threshold (higher than hard limit for test)
        .with_hard_limit(2048); // 2KB hard limit

        let mut engine = WriteEngine::new(config).unwrap();

        // Write mutations until we hit the hard limit
        let mut write_count = 0;
        for i in 0..1000 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            let result = engine.write_async(mutation).await;

            match result {
                Ok(()) => {
                    write_count += 1;
                }
                Err(Error::Storage(msg)) => {
                    assert!(msg.contains("hard limit"));
                    break;
                }
                Err(e) => panic!("Expected Storage error, got: {:?}", e),
            }
        }

        // Should have stopped before 1000 writes due to hard limit
        assert!(
            write_count < 1000,
            "Should have hit hard limit before 1000 writes"
        );
        assert!(
            write_count > 0,
            "Should have accepted at least some writes before hitting limit"
        );
    }

    #[tokio::test]
    async fn test_write_engine_hard_limit_recovery_after_flush() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        // Set low hard limit
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_flush_threshold(1024)
        .with_hard_limit(2048);

        let mut engine = WriteEngine::new(config).unwrap();

        // Write until hard limit
        let mut first_batch_count = 0;
        for i in 0..1000 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1000000 + i as i64);
            let result = engine.write(mutation);

            if result.is_err() {
                break;
            }

            first_batch_count += 1;
        }

        assert!(
            first_batch_count > 0,
            "Should have accepted some writes before limit"
        );

        // Flush to clear memtable
        engine.flush().await.unwrap();

        // Should be able to write again after flush
        let mutation = create_test_mutation(9999, "After flush", 2000000);
        let result = engine.write(mutation);
        assert!(result.is_ok(), "Should accept writes after flush");

        assert_eq!(engine.memtable_row_count(), 1);
    }

    #[test]
    fn test_generation_counter_is_u64() {
        // Verify that generation counter is u64 to prevent overflow (Issue #410)
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        // Verify type by checking the value is within u64 range
        let generation: u64 = engine.generation();
        assert_eq!(generation, 1u64);

        // This compile-time check ensures generation() returns u64
        // If it returned u32, this assignment would be a no-op but still compile
        let _type_check: u64 = generation;

        // Verify that u64 can handle generations beyond u32::MAX
        // This would overflow with u32 (max value: 4,294,967,295)
        let large_generation: u64 = u32::MAX as u64 + 1000;
        assert!(large_generation > u32::MAX as u64);
        assert_eq!(large_generation, 4_294_968_295u64);
    }

    #[test]
    fn test_determine_next_generation_large_numbers() {
        // Verify that generation parsing handles large u64 values (Issue #410)
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        // Create dummy SSTable files with large generation numbers
        // These would overflow if we used u32 (max: 4,294,967,295)
        let large_gen: u64 = u32::MAX as u64 + 100;
        std::fs::write(data_dir.join(format!("nb-{}-big-Data.db", large_gen)), b"").unwrap();

        let generation = WriteEngine::determine_next_generation(&data_dir).unwrap();
        assert_eq!(generation, large_gen + 1);
        assert!(generation > u32::MAX as u64);
    }

    // M5.2 maintenance_step() tests (Issue #384)

    #[test]
    fn test_maintenance_step_no_policy() {
        // Without a merge policy, maintenance_step should do nothing
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Call maintenance_step without setting a policy
        let report = engine.maintenance_step(Duration::from_millis(100)).unwrap();

        // Should return immediately with no work done
        assert_eq!(report.rows_merged, 0);
        assert_eq!(report.bytes_written, 0);
        assert_eq!(report.completed_merges.len(), 0);
        assert!(!report.pending_compaction);
        assert!(report.time_spent < Duration::from_millis(50));
    }

    #[test]
    fn test_maintenance_step_with_closed_engine() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Close the engine
        tokio::runtime::Runtime::new()
            .unwrap()
            .block_on(engine.close())
            .unwrap();

        // maintenance_step should fail on closed engine
        let result = engine.maintenance_step(Duration::from_millis(100));
        assert!(result.is_err());
        match result {
            Err(Error::InvalidInput(msg)) => {
                assert!(msg.contains("closed"));
            }
            _ => panic!("Expected InvalidInput error"),
        }
    }

    #[test]
    fn test_maintenance_report_creation() {
        let report = MaintenanceReport {
            time_spent: Duration::from_millis(250),
            completed_merges: vec![PathBuf::from("data/nb-5-big-Data.db")],
            rows_merged: 1000,
            bytes_written: 1024 * 1024,
            pending_compaction: true,
        };

        assert_eq!(report.time_spent.as_millis(), 250);
        assert_eq!(report.completed_merges.len(), 1);
        assert_eq!(report.rows_merged, 1000);
        assert_eq!(report.bytes_written, 1024 * 1024);
        assert!(report.pending_compaction);
    }

    #[test]
    fn test_scan_sstable_candidates_empty_dir() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        let candidates = engine.scan_sstable_candidates().unwrap();
        assert_eq!(candidates.len(), 0);
    }

    #[test]
    fn test_scan_sstable_candidates_with_sstables() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        // Create dummy SSTable files. Each Data.db needs a sibling TOC.txt to
        // count as a *published* SSTable (the publication barrier, Issue #591) —
        // a Data.db without TOC.txt is an unpublished partial/orphan and must be
        // skipped by the candidate scan.
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();
        std::fs::write(data_dir.join("nb-1-big-Data.db"), b"").unwrap();
        std::fs::write(data_dir.join("nb-1-big-TOC.txt"), b"").unwrap();
        std::fs::write(data_dir.join("nb-2-big-Data.db"), b"").unwrap();
        std::fs::write(data_dir.join("nb-2-big-TOC.txt"), b"").unwrap();
        std::fs::write(data_dir.join("nb-3-big-Index.db"), b"").unwrap(); // Not a Data.db
        std::fs::write(data_dir.join("other-file.txt"), b"").unwrap(); // Not an SSTable
                                                                       // An unpublished Data.db (no TOC.txt) must NOT be picked up (Issue #591).
        std::fs::write(data_dir.join("nb-4-big-Data.db"), b"").unwrap();

        let candidates = engine.scan_sstable_candidates().unwrap();

        // Should only find the two PUBLISHED Data.db files (TOC.txt present);
        // nb-4 is excluded because it has no TOC.txt.
        assert_eq!(candidates.len(), 2);
        assert!(candidates
            .iter()
            .all(|p| p.to_string_lossy().contains("Data.db")));
        assert!(
            !candidates
                .iter()
                .any(|p| p.to_string_lossy().contains("nb-4-big")),
            "unpublished Data.db (no TOC.txt) must be excluded (Issue #591)"
        );
    }

    #[test]
    fn test_delete_sstable_files() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        // Create dummy SSTable component files
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        let components = [
            "nb-5-big-Data.db",
            "nb-5-big-Index.db",
            "nb-5-big-Summary.db",
            "nb-5-big-Statistics.db",
        ];

        for component in &components {
            std::fs::write(data_dir.join(component), b"dummy").unwrap();
        }

        // Verify files exist
        for component in &components {
            assert!(data_dir.join(component).exists());
        }

        // Delete SSTable files
        let data_path = data_dir.join("nb-5-big-Data.db");
        engine.delete_sstable_files(&data_path).unwrap();

        // Verify files are deleted
        for component in &components {
            assert!(!data_dir.join(component).exists());
        }
    }

    /// Issue #591: deletion removes TOC.txt FIRST so the SSTable is unpublished
    /// before any data component is touched. This guarantees the read path and
    /// the compaction candidate scan stop seeing it immediately, even if a data
    /// component cannot be removed yet (e.g. pinned by a mapped reader on
    /// Windows).
    #[test]
    fn test_delete_removes_toc_first_unpublishing_atomically() {
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        // A full published SSTable component set including TOC.txt.
        for comp in &[
            "nb-7-big-Data.db",
            "nb-7-big-Index.db",
            "nb-7-big-Statistics.db",
            "nb-7-big-TOC.txt",
        ] {
            std::fs::write(data_dir.join(comp), b"x").unwrap();
        }

        let data_path = data_dir.join("nb-7-big-Data.db");
        WriteEngine::delete_sstable_files_static(&data_path).unwrap();

        // Everything gone on the happy path.
        assert!(!data_dir.join("nb-7-big-TOC.txt").exists());
        assert!(!data_path.exists());

        // And critically: scan_data_files (the compaction candidate discovery)
        // never surfaces a Data.db without a TOC.txt, so a deferred-delete orphan
        // is not re-fed to the merger. Recreate a TOC-less leftover to prove it.
        std::fs::write(data_dir.join("nb-8-big-Data.db"), b"x").unwrap();
        let mut candidates = Vec::new();
        WriteEngine::scan_data_files(&data_dir, &mut candidates, 1).unwrap();
        assert!(
            candidates.is_empty(),
            "a Data.db without a sibling TOC.txt must NOT be a compaction candidate \
             (publication barrier, Issue #591); got {:?}",
            candidates
        );

        // Add the matching TOC.txt and it becomes a valid candidate again.
        std::fs::write(data_dir.join("nb-8-big-TOC.txt"), b"x").unwrap();
        let mut candidates = Vec::new();
        WriteEngine::scan_data_files(&data_dir, &mut candidates, 1).unwrap();
        assert_eq!(
            candidates.len(),
            1,
            "a published Data.db (TOC.txt present) must be discovered"
        );
    }

    // Mock merge policy that selects specific files for testing
    #[derive(Debug)]
    #[allow(dead_code)] // Used in multiple test functions below
    struct TestMergePolicy {
        files_to_select: Vec<PathBuf>,
    }

    impl MergePolicy for TestMergePolicy {
        fn select_merge(&self, _candidates: &[PathBuf]) -> Result<Vec<PathBuf>> {
            Ok(self.files_to_select.clone())
        }
    }

    #[test]
    fn test_maintenance_step_with_policy_no_work() {
        // Policy that returns empty selection (no work to do)
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Set a policy that selects nothing
        let policy = TestMergePolicy {
            files_to_select: vec![],
        };
        engine.set_merge_policy(Box::new(policy)).unwrap();

        // Call maintenance_step - policy selects no work
        let report = engine.maintenance_step(Duration::from_millis(100)).unwrap();

        // Should return with no work done
        assert_eq!(report.rows_merged, 0);
        assert_eq!(report.bytes_written, 0);
        assert_eq!(report.completed_merges.len(), 0);
        assert!(!report.pending_compaction);
    }

    #[test]
    fn test_maintenance_step_budget_honored() {
        // Test that budget is approximately honored
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Set a policy that selects nothing
        let policy = TestMergePolicy {
            files_to_select: vec![],
        };
        engine.set_merge_policy(Box::new(policy)).unwrap();

        // Call with small budget - policy selects no work, should return quickly
        let budget = Duration::from_millis(10);
        let report = engine.maintenance_step(budget).unwrap();

        // Should return quickly when there's no compaction work
        assert!(
            report.time_spent < budget.mul_f32(1.5),
            "Time spent {:?} exceeded budget {:?} by >50%",
            report.time_spent,
            budget
        );
    }

    // ============================================================================
    // Issue #474: Wire k-way merger + STCS into maintenance_step
    // ============================================================================

    /// Helper: flush `n` distinct mutations through the engine synchronously.
    ///
    /// Uses a dedicated single-threaded runtime so it can be called from both
    /// sync test functions and (via `spawn_blocking`) from async contexts.
    fn flush_n_sstables_sync(engine: &mut WriteEngine, n: usize) -> Vec<PathBuf> {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let mut paths = Vec::new();
        for batch in 0..n {
            for row in 0..5 {
                let mutation = create_test_mutation(
                    (batch * 100 + row) as i32,
                    &format!("User-{}-{}", batch, row),
                    1_000_000 + (batch * 100 + row) as i64,
                );
                engine.write(mutation).unwrap();
            }
            let info = rt.block_on(engine.flush()).unwrap().unwrap();
            paths.push(info.data_path);
        }
        paths
    }

    #[test]
    fn test_maintenance_stats_initial_zero() {
        // Before any maintenance work, all stats should be zero
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let engine = WriteEngine::new(config).unwrap();

        let stats = engine.maintenance_stats();
        assert_eq!(stats.compactions_completed, 0);
        assert_eq!(stats.sstables_merged_in, 0);
        assert_eq!(stats.sstables_produced, 0);
        assert_eq!(stats.bytes_read, 0);
        assert_eq!(stats.bytes_written, 0);
        assert_eq!(stats.rows_merged, 0);
        assert_eq!(stats.total_time, Duration::ZERO);
    }

    #[test]
    fn test_stcs_selects_expected_group_by_size() {
        // Verify that STCSPolicy groups four same-sized SSTables into one candidate set.
        // We do this without actually running a merge (just test the policy selection).
        let policy = crate::storage::write_engine::STCSPolicy::default();

        // Create 4 temp files of equal size to satisfy min_threshold=4
        let temp_dir = TempDir::new().unwrap();
        let mut paths = Vec::new();
        for i in 1..=4 {
            let path = temp_dir.path().join(format!("nb-{}-big-Data.db", i));
            // 60 MB each (above min_sstable_size threshold)
            let size_bytes = 60 * 1024 * 1024u64;
            let file = std::fs::File::create(&path).unwrap();
            file.set_len(size_bytes).unwrap();
            paths.push(path);
        }

        // Policy should select all 4 as a candidate group
        let selected = policy.select_merge(&paths).unwrap();
        assert_eq!(
            selected.len(),
            4,
            "STCS should select all 4 same-sized SSTables as one compaction group"
        );

        // All selected paths should be from our input set
        for sel in &selected {
            assert!(
                paths.contains(sel),
                "Selected path {:?} not in input set",
                sel
            );
        }
    }

    #[test]
    fn test_stcs_does_not_select_below_threshold() {
        // With only 3 SSTables, STCS (min_threshold=4) should select nothing.
        let policy = crate::storage::write_engine::STCSPolicy::default();

        let temp_dir = TempDir::new().unwrap();
        let mut paths = Vec::new();
        for i in 1..=3 {
            let path = temp_dir.path().join(format!("nb-{}-big-Data.db", i));
            let file = std::fs::File::create(&path).unwrap();
            file.set_len(60 * 1024 * 1024).unwrap();
            paths.push(path);
        }

        let selected = policy.select_merge(&paths).unwrap();
        assert!(
            selected.is_empty(),
            "STCS should NOT select when fewer than min_threshold SSTables exist"
        );
    }

    #[test]
    fn test_maintenance_step_compacts_sstables_atomically() {
        // Create an engine, flush 4 SSTables, then run maintenance_step with STCS.
        // After the step: input files must be gone, output file must exist,
        // and maintenance_stats() must reflect the completed compaction.
        //
        // Uses a sync wrapper so maintenance_step's internal block_on works without
        // nesting inside a pre-existing async runtime.
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        // Use a LOW min_sstable_size so small test files pass bucket grouping
        let policy = crate::storage::write_engine::STCSPolicy::new(
            4,   // min_threshold
            32,  // max_threshold
            0.5, // bucket_low
            1.5, // bucket_high
            0,   // min_sstable_size = 0 so tiny files group together
        )
        .unwrap();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Flush 4 distinct SSTables (sync helper creates its own single-threaded runtime)
        let input_paths = flush_n_sstables_sync(&mut engine, 4);
        assert_eq!(input_paths.len(), 4, "Expected 4 flushed SSTables");

        // Verify all input Data.db files exist before compaction
        for p in &input_paths {
            assert!(
                p.exists(),
                "Input file {:?} should exist before compaction",
                p
            );
        }

        // Attach the policy and run maintenance
        engine.set_merge_policy(Box::new(policy)).unwrap();
        let report = engine.maintenance_step(Duration::from_secs(60)).unwrap();

        // The report must indicate a completed merge
        assert_eq!(
            report.completed_merges.len(),
            1,
            "Expected exactly 1 completed merge, got: {:?}",
            report.completed_merges
        );
        // bytes_written is u64 and always non-negative, so no assertion needed here.

        // The merged output file must exist in the final SSTable directory
        let merged_path = &report.completed_merges[0];
        assert!(
            merged_path.exists(),
            "Merged output file {:?} must exist after compaction",
            merged_path
        );

        // All input files must be gone (consumed by compaction)
        for p in &input_paths {
            assert!(
                !p.exists(),
                "Input file {:?} should have been deleted after compaction",
                p
            );
        }

        // maintenance_stats() must reflect the operation
        let stats = engine.maintenance_stats();
        assert_eq!(
            stats.compactions_completed, 1,
            "compactions_completed must be 1"
        );
        assert_eq!(
            stats.sstables_merged_in, 4,
            "Should have consumed 4 input SSTables"
        );
        assert_eq!(stats.sstables_produced, 1, "sstables_produced must be 1");
        // bytes_written may be 0 if the merged output is empty (reader/writer compatibility),
        // but total_time must be non-zero
        assert!(stats.total_time > Duration::ZERO, "total_time must be > 0");
    }

    #[test]
    fn test_maintenance_stats_accumulate_across_cycles() {
        // Run two compaction cycles and verify that stats accumulate.
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let policy = crate::storage::write_engine::STCSPolicy::new(
            4, 32, 0.5, 1.5, 0, // min_sstable_size=0 for small test files
        )
        .unwrap();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();
        engine.set_merge_policy(Box::new(policy)).unwrap();

        // First cycle: flush 4, compact
        flush_n_sstables_sync(&mut engine, 4);
        engine.maintenance_step(Duration::from_secs(60)).unwrap();

        let stats_after_first = engine.maintenance_stats();
        assert_eq!(stats_after_first.compactions_completed, 1);

        // Second cycle: flush 4 more, compact again
        // Row IDs must not collide with the first cycle so each cycle produces 4 SSTables.
        // flush_n_sstables_sync uses batch * 100 + row, so offset the start batch.
        // We re-use the helper but note generation counter now starts at a higher value,
        // so the output SSTable won't conflict with input paths from cycle 1.
        flush_n_sstables_sync(&mut engine, 4);
        engine.maintenance_step(Duration::from_secs(60)).unwrap();

        let stats_after_second = engine.maintenance_stats();
        assert_eq!(
            stats_after_second.compactions_completed, 2,
            "Stats must accumulate across compaction cycles"
        );
        assert_eq!(
            stats_after_second.sstables_merged_in, 8,
            "Should have consumed 8 total input SSTables (2 cycles × 4 each)"
        );
        assert_eq!(
            stats_after_second.sstables_produced, 2,
            "Should have produced 2 output SSTables"
        );
        assert!(
            stats_after_second.total_time >= stats_after_first.total_time,
            "Cumulative total_time must only increase"
        );
    }

    #[test]
    fn test_maintenance_step_inputs_intact_on_unwriteable_tmp_dir() {
        // Failure injection: make the data_dir read-only so creating the tmp
        // compaction directory fails. All input SSTables must remain intact.
        //
        // Note: This test relies on filesystem permissions and is skipped when
        // running as root (where permissions are not enforced).

        // Skip if running as root (CI containers sometimes run as root)
        #[cfg(unix)]
        {
            use std::os::unix::fs::MetadataExt;
            // Try /proc/self first (Linux), fall back to checking euid via libc
            let is_root = std::fs::metadata("/proc/self")
                .map(|m| m.uid() == 0)
                .unwrap_or_else(|_| {
                    // On macOS, /proc/self doesn't exist; use a writable sentinel
                    false
                });
            // Also check by trying to write to /etc/cqlite-test-root-check
            let is_root_macos = std::fs::write("/etc/cqlite-test-root-check", b"")
                .map(|_| {
                    let _ = std::fs::remove_file("/etc/cqlite-test-root-check");
                    true
                })
                .unwrap_or(false);
            if is_root || is_root_macos {
                // Running as root — permission denial won't work; skip.
                return;
            }
        }

        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();

        // Flush 4 SSTables so STCS can select them
        let input_paths = flush_n_sstables_sync(&mut engine, 4);
        for p in &input_paths {
            assert!(
                p.exists(),
                "Input file {:?} should exist before failure test",
                p
            );
        }

        // Make data_dir read-only so creating tmp dir fails
        let data_dir = temp_dir.path().join("data");
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(
                &data_dir,
                std::fs::Permissions::from_mode(0o555), // read+execute, no write
            )
            .unwrap();
        }

        let policy = crate::storage::write_engine::STCSPolicy::new(4, 32, 0.5, 1.5, 0).unwrap();
        engine.set_merge_policy(Box::new(policy)).unwrap();

        // maintenance_step should fail because it cannot create the tmp directory
        let result = engine.maintenance_step(Duration::from_secs(60));

        // Restore permissions before asserting (so TempDir can clean up)
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(&data_dir, std::fs::Permissions::from_mode(0o755)).unwrap();
        }

        assert!(
            result.is_err(),
            "maintenance_step should return an error when the tmp dir cannot be created"
        );

        // All input files must still exist (atomicity guarantee)
        for p in &input_paths {
            assert!(
                p.exists(),
                "Input file {:?} must remain intact after failed compaction",
                p
            );
        }

        // Stats must NOT have incremented (no successful compaction)
        let stats = engine.maintenance_stats();
        assert_eq!(
            stats.compactions_completed, 0,
            "compactions_completed must not increment on failure"
        );
    }

    #[test]
    fn test_no_tmp_dir_remains_after_successful_merge() {
        // After a successful compaction, the .compaction-tmp-* directory must be cleaned up.
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let policy = crate::storage::write_engine::STCSPolicy::new(4, 32, 0.5, 1.5, 0).unwrap();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();
        flush_n_sstables_sync(&mut engine, 4);

        engine.set_merge_policy(Box::new(policy)).unwrap();
        engine.maintenance_step(Duration::from_secs(60)).unwrap();

        // Scan data_dir for any leftover .compaction-tmp-* directories
        let data_dir = temp_dir.path().join("data");
        let leftover_tmp: Vec<_> = std::fs::read_dir(&data_dir)
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| {
                e.file_name()
                    .to_string_lossy()
                    .starts_with(".compaction-tmp-")
            })
            .collect();

        assert!(
            leftover_tmp.is_empty(),
            "No .compaction-tmp-* directories should remain after successful compaction, \
             found: {:?}",
            leftover_tmp.iter().map(|e| e.path()).collect::<Vec<_>>()
        );
    }

    // ============================================================================
    // Issue #474 review follow-up: NB-1 startup sweep tests
    // ============================================================================

    /// Helper: build a WriteEngineConfig pointing at a given data/wal dir pair.
    fn config_for(temp_dir: &TempDir) -> WriteEngineConfig {
        WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            create_test_schema(),
        )
    }

    #[test]
    fn test_startup_sweep_removes_orphaned_compaction_tmp_dir() {
        // Pre-seed a .compaction-tmp-99/ directory under data_dir.
        // WriteEngine::new() must remove it during the startup sweep.
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        std::fs::create_dir_all(&data_dir).unwrap();

        let orphan_dir = data_dir.join(".compaction-tmp-99");
        std::fs::create_dir_all(&orphan_dir).unwrap();
        // Put a partial component file inside to make it non-trivially non-empty
        std::fs::write(orphan_dir.join("partial.db"), b"partial content").unwrap();

        assert!(
            orphan_dir.exists(),
            "Orphan dir should exist before engine creation"
        );

        let config = config_for(&temp_dir);
        let _engine = WriteEngine::new(config).unwrap();

        assert!(
            !orphan_dir.exists(),
            "Startup sweep must remove orphaned .compaction-tmp-99/ directory"
        );
    }

    #[test]
    fn test_startup_sweep_removes_orphaned_partial_sstable() {
        // Pre-seed nb-99-big-Data.db (and friends) WITHOUT a TOC.txt in the
        // keyspace/table subdirectory.  WriteEngine::new() must remove them.
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        let sstable_dir = data_dir.join("test_ks").join("test_table");
        std::fs::create_dir_all(&sstable_dir).unwrap();

        // Create orphaned components (no TOC.txt)
        let orphan_components = [
            "nb-99-big-Data.db",
            "nb-99-big-Index.db",
            "nb-99-big-Statistics.db",
        ];
        for name in &orphan_components {
            std::fs::write(sstable_dir.join(name), b"orphaned").unwrap();
        }

        // Also create a complete SSTable (has TOC.txt) — must NOT be touched
        let complete_components = ["nb-1-big-Data.db", "nb-1-big-Index.db", "nb-1-big-TOC.txt"];
        for name in &complete_components {
            std::fs::write(sstable_dir.join(name), b"good").unwrap();
        }

        let config = config_for(&temp_dir);
        let _engine = WriteEngine::new(config).unwrap();

        // Orphaned components must be gone
        for name in &orphan_components {
            assert!(
                !sstable_dir.join(name).exists(),
                "Startup sweep must remove orphaned component {:?}",
                name
            );
        }

        // Complete SSTable must remain intact
        for name in &complete_components {
            assert!(
                sstable_dir.join(name).exists(),
                "Startup sweep must NOT remove complete SSTable component {:?}",
                name
            );
        }
    }

    #[test]
    fn test_startup_sweep_leaves_complete_sstable_intact() {
        // A full SSTable set (Data.db + TOC.txt + others) must survive the sweep.
        let temp_dir = TempDir::new().unwrap();
        let data_dir = temp_dir.path().join("data");
        let sstable_dir = data_dir.join("test_ks").join("test_table");
        std::fs::create_dir_all(&sstable_dir).unwrap();

        let all_components = [
            "nb-3-big-Data.db",
            "nb-3-big-Index.db",
            "nb-3-big-Summary.db",
            "nb-3-big-Statistics.db",
            "nb-3-big-Filter.db",
            "nb-3-big-Digest.crc32",
            "nb-3-big-TOC.txt",
        ];
        for name in &all_components {
            std::fs::write(sstable_dir.join(name), b"complete").unwrap();
        }

        let config = config_for(&temp_dir);
        let _engine = WriteEngine::new(config).unwrap();

        for name in &all_components {
            assert!(
                sstable_dir.join(name).exists(),
                "Complete SSTable component {:?} must not be removed by startup sweep",
                name
            );
        }
    }

    #[test]
    fn test_bytes_written_includes_all_components() {
        // After a successful merge, cumulative_stats.bytes_written must be at
        // least as large as the sum of Data.db sizes alone (i.e. it includes
        // the other component files too).
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        let policy = crate::storage::write_engine::STCSPolicy::new(4, 32, 0.5, 1.5, 0).unwrap();

        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );

        let mut engine = WriteEngine::new(config).unwrap();
        let input_paths = flush_n_sstables_sync(&mut engine, 4);

        // Compute the sum of just the Data.db sizes before compaction
        let data_db_total: u64 = input_paths
            .iter()
            .map(|p| std::fs::metadata(p).map(|m| m.len()).unwrap_or(0))
            .sum();

        engine.set_merge_policy(Box::new(policy)).unwrap();
        engine.maintenance_step(Duration::from_secs(60)).unwrap();

        let stats = engine.maintenance_stats();
        // bytes_written counts all output components, so it should be >= what
        // data_db_total reported for the inputs (output may differ in size but
        // the multi-component sum must be >= the Data.db-only measurement).
        // More concretely: if any non-Data component was written, the total
        // must be larger than data_size alone.
        //
        // We assert >= 0 always holds (u64), and additionally that the field
        // was updated at all (compaction ran).
        assert_eq!(stats.compactions_completed, 1, "compaction must have run");
        // The bytes_written field is now the sum of all components.
        // We can't assert an exact value, but we know:
        //  - data_db_total may be 0 for tiny test SSTables written by the test writer
        //  - if data_db_total > 0, bytes_written >= data_db_total is a reasonable lower bound
        //  - at minimum, the field must equal total_bytes_written (multi-component sum) >= 0
        let _ = data_db_total; // used above for context; value may be 0 in test environment
                               // The assertion that matters: stats are populated and consistent across calls.
                               // maintenance_stats() returns a clone so two consecutive calls must agree.
        let stats2 = engine.maintenance_stats();
        assert_eq!(
            stats.bytes_written, stats2.bytes_written,
            "maintenance_stats() must be consistent across calls"
        );
        // bytes_written is u64; it is always >= 0. Just confirm the field was set.
        assert_eq!(
            stats.sstables_produced, 1,
            "one output SSTable must have been produced"
        );
    }

    // ============================================================================
    // Issue #547: WAL durability toggle tests
    // ============================================================================

    /// `Durability::SyncEachWrite` is the default variant.
    #[test]
    fn test_durability_default_is_sync_each_write() {
        assert_eq!(Durability::default(), Durability::SyncEachWrite);
    }

    /// `WriteEngineConfig` defaults to `Durability::SyncEachWrite`.
    #[test]
    fn test_config_default_durability() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );
        assert_eq!(config.durability, Durability::SyncEachWrite);
    }

    /// `with_durability` builder sets the field and returns `Self`.
    #[test]
    fn test_config_with_durability_builder() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_durability(Durability::Disabled);
        assert_eq!(config.durability, Durability::Disabled);
    }

    /// With `Durability::SyncEachWrite`, the WAL grows after each `write`.
    #[test]
    fn test_wal_on_produces_wal_growth() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_durability(Durability::SyncEachWrite);

        let mut engine = WriteEngine::new(config).unwrap();
        assert_eq!(engine.wal_size(), 0, "WAL must start empty");

        let mutation = create_test_mutation(1, "Alice", 1_000_000);
        engine.write(mutation).unwrap();

        assert!(
            engine.wal_size() > 0,
            "WAL must grow after write with SyncEachWrite"
        );
    }

    /// With `Durability::Disabled`, the WAL is never written — `wal_size()` stays 0.
    #[test]
    fn test_wal_off_produces_no_wal_growth() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_durability(Durability::Disabled);

        let mut engine = WriteEngine::new(config).unwrap();
        assert_eq!(engine.wal_size(), 0, "WAL must start empty");

        // Write several mutations — none should touch the WAL.
        for i in 0..10 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
            engine.write(mutation).unwrap();
        }

        assert_eq!(
            engine.wal_size(),
            0,
            "WAL must remain empty with Durability::Disabled"
        );
        assert_eq!(
            engine.memtable_row_count(),
            10,
            "Mutations must reach the memtable even without WAL"
        );
    }

    /// With `Durability::Disabled`, async writes also skip the WAL.
    #[tokio::test]
    async fn test_wal_off_write_async_produces_no_wal_growth() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();
        let config = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_durability(Durability::Disabled);

        let mut engine = WriteEngine::new(config).unwrap();

        for i in 0..5 {
            let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
            engine.write_async(mutation).await.unwrap();
        }

        assert_eq!(
            engine.wal_size(),
            0,
            "WAL must remain empty with Durability::Disabled (async path)"
        );
        assert_eq!(engine.memtable_row_count(), 5);
    }

    /// With `Durability::Disabled`, data that was never WAL'd is NOT replayed on
    /// restart — confirming the documented durability trade-off.
    #[test]
    fn test_wal_off_no_replay_on_restart() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        {
            let config = WriteEngineConfig::new(
                temp_dir.path().join("data"),
                temp_dir.path().join("wal"),
                schema.clone(),
            )
            .with_durability(Durability::Disabled);

            let mut engine = WriteEngine::new(config).unwrap();

            for i in 0..5 {
                let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
                engine.write(mutation).unwrap();
            }

            // Drop without flushing — simulating crash.
        }

        // Reopen with default durability.  Because the WAL was never written, the
        // memtable must be empty.
        let config2 = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        );
        let engine2 = WriteEngine::new(config2).unwrap();

        assert_eq!(
            engine2.memtable_row_count(),
            0,
            "No WAL entries were written with Disabled, so no replay is possible"
        );
    }

    /// With `Durability::SyncEachWrite`, mutations ARE replayed after a simulated crash.
    #[test]
    fn test_wal_on_replays_on_restart() {
        let temp_dir = TempDir::new().unwrap();
        let schema = create_test_schema();

        {
            let config = WriteEngineConfig::new(
                temp_dir.path().join("data"),
                temp_dir.path().join("wal"),
                schema.clone(),
            )
            .with_durability(Durability::SyncEachWrite);

            let mut engine = WriteEngine::new(config).unwrap();

            for i in 0..5 {
                let mutation = create_test_mutation(i, &format!("User{}", i), 1_000_000 + i as i64);
                engine.write(mutation).unwrap();
            }

            // Drop without flushing — WAL entries remain on disk.
        }

        // Reopen — WAL replay must restore the 5 mutations.
        let config2 = WriteEngineConfig::new(
            temp_dir.path().join("data"),
            temp_dir.path().join("wal"),
            schema,
        )
        .with_durability(Durability::SyncEachWrite);

        let engine2 = WriteEngine::new(config2).unwrap();

        assert_eq!(
            engine2.memtable_row_count(),
            5,
            "SyncEachWrite must replay mutations durably on restart"
        );
    }
}