heliosdb-nano 4.6.3

//! Transaction implementation
//!
//! Basic 2PL (Two-Phase Locking) with standard MVCC snapshot isolation.
//! Optimized with lock-free concurrent data structures for improved performance.
//!
//! v3.1.0: Enhanced with session-aware locking and isolation level support

use super::conflict::WriteConflictRegistry;
use super::dirty_tracker::DirtyTracker;
use super::lock_manager::{LockGuard, LockManager, LockType};
use super::time_travel::SnapshotManager;
use super::{Key, Snapshot, SnapshotId};
use crate::session::{IsolationLevel, SessionId};
use crate::{Error, Result, Tuple};
use dashmap::{DashMap, DashSet};
use parking_lot::RwLock;
use rocksdb::{WriteOptions, DB};
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::Arc;
use tracing::{debug, trace, warn};

/// Transaction state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransactionState {
    /// Active (can read/write)
    Active,
    /// Committed
    Committed,
    /// Aborted
    Aborted,
}

impl TransactionState {
    /// Convert state to u8 for atomic storage
    const fn to_u8(self) -> u8 {
        match self {
            Self::Active => 0,
            Self::Committed => 1,
            Self::Aborted => 2,
        }
    }

    /// Convert u8 to state
    const fn from_u8(value: u8) -> Self {
        match value {
            0 => Self::Active,
            1 => Self::Committed,
            _ => Self::Aborted,
        }
    }
}

/// Transaction
///
/// Provides ACID guarantees using standard snapshot isolation.
///
/// Optimized with lock-free data structures:
/// - DashMap for concurrent write_set access without mutex contention
/// - AtomicU8 for lock-free state checking
/// This significantly reduces lock contention on the read path.
///
/// v3.1.0 enhancements:
/// - Session-aware transactions with isolation levels
/// - Lock manager integration for deadlock detection
/// - Dirty state tracking for dump operations
pub struct Transaction {
    /// Database handle
    db: Arc<DB>,
    /// Snapshot for reads
    snapshot: Snapshot,
    /// Snapshot timestamp (read timestamp for MVCC)
    snapshot_ts: u64,
    /// Unique transaction ID for locking and tracking
    transaction_id: u64,
    /// Snapshot manager for versioned reads
    snapshot_manager: Arc<SnapshotManager>,
    /// Write set (buffered writes) - uses DashMap for lock-free concurrent access
    write_set: Arc<DashMap<Key, Option<Vec<u8>>>>,
    /// Append-only data-row inserts staged by SQL fast paths.
    ///
    /// Large explicit-transaction INSERT loops are single-writer in the embedded
    /// API and do not need per-row DashMap hashing. Updates/deletes still go to
    /// `write_set`; if both contain the same key, `write_set` is the final value.
    insert_log: Arc<RwLock<Vec<(Key, Vec<u8>)>>>,
    /// Latest row-counter values staged by fast transactional inserts.
    ///
    /// Only the final counter value per table needs to be persisted at commit,
    /// so keeping it out of `write_set` avoids a redundant DashMap overwrite on
    /// every inserted row in large explicit transactions.
    row_counter_stages: Arc<RwLock<std::collections::HashMap<String, u64>>>,
    /// R2.3: tables touched by this transaction's staged writes (parsed from
    /// `data:{table}:{row_id}` keys in `put`/`put_insert_fast`/`delete`, plus
    /// `stage_row_counter` tables). Lets the executor answer
    /// [`has_writes_for_table`](Self::has_writes_for_table) in O(1) instead of
    /// scanning `write_set`/`insert_log` on every SELECT gate check.
    ///
    /// Deliberately monotonic: ROLLBACK TO SAVEPOINT does not shrink it. An
    /// over-approximation only keeps a table on the slow (write-set-merging)
    /// read path, which is always correct.
    written_tables: DashSet<String>,
    /// R2.3: set when a staged write's key is not a parseable `data:` key
    /// (e.g. raw `counter:`/meta keys via `put`). From then on
    /// `has_writes_for_table` answers `true` for every table — conservative,
    /// keeps all reads on the slow path.
    has_unattributed_writes: AtomicBool,
    /// Transaction state - uses AtomicU8 for lock-free state checking
    state: AtomicU8,
    /// Session ID (for multi-user support)
    session_id: Option<SessionId>,
    /// Isolation level for this transaction
    isolation_level: IsolationLevel,
    /// Lock manager for concurrency control
    lock_manager: Option<Arc<LockManager>>,
    /// Acquired locks (RAII guards for automatic release)
    acquired_locks: Arc<RwLock<Vec<LockGuard>>>,
    /// Dirty tracker for dump operations
    dirty_tracker: Option<Arc<DirtyTracker>>,
    /// Whether to emit MVCC version-history keys (`v:` / `v_idx:`) at commit.
    /// Mirrors `StorageConfig::time_travel_enabled`. When false, commit writes
    /// only the `data:` key per row (no version value + index, no double
    /// serialization) — for read-committed workloads that don't need
    /// AS OF / snapshot-history reads.
    versioning_enabled: bool,
    /// Whether RocksDB's own WAL should be used for the final WriteBatch.
    /// Disabled only for `memory_only` databases, where durability is not a
    /// contract and the temporary RocksDB directory is discarded on close.
    rocksdb_wal_enabled: bool,
    /// R1.3: fsync the commit WriteBatch (power-loss durable commits).
    sync_commit: bool,
    /// R1.3 phase 2: engine-wide leader/follower group committer. When set
    /// (always, for engine/session-constructed transactions), a `sync_commit`
    /// writes its batch UNSYNCED and then waits on ONE cohort-wide
    /// `flush_wal(true)` instead of paying a private fsync. When absent,
    /// falls back to the phase-1 per-commit synced WriteBatch.
    group_committer: Option<Arc<super::group_commit::GroupCommitter>>,
    /// Engine row cache (R1.3-p2 correctness fix). Written rows are
    /// invalidated INSIDE commit — after the batch applies, BEFORE
    /// `end_commit` lifts the snapshot barrier — so a fresh snapshot that
    /// passes the barrier can never hit a stale cached pre-commit row (the
    /// R0.2 lost-update vector via the UPDATE arm's PK point-lookup; the
    /// caller-side invalidation after commit returns was too late, and the
    /// R1.3-p2 group-fsync wait would have widened that window to a whole
    /// fsync).
    row_cache: Option<Arc<super::RowCache>>,
    /// W2.5: engine per-table committed-write watermark map. When wired,
    /// commit raises each written table's watermark to `commit_ts` BEFORE the
    /// batch is applied (the same fence rationale as `row_cache`, but the
    /// opposite direction: the watermark must be visible to a racing reader
    /// *before* the write is, so a snapshot older than this commit keeps
    /// taking the snapshot read path for these tables). Branch commits carry
    /// `bdata:` keys that never enter `written_tables`, so this branch-blind
    /// map only records main-branch tables.
    write_watermarks: Option<Arc<DashMap<String, u64>>>,
    /// Write-write conflict registry (R0.2). When present, commits record
    /// their write-set keys; when `conflict_validation` is also set, commit
    /// aborts with a serialization failure if any key was committed after
    /// this transaction's snapshot (first-committer-wins).
    conflict_registry: Option<Arc<WriteConflictRegistry>>,
    conflict_validation: bool,
    /// R4.3: id of this transaction's version-GC snapshot pin in the
    /// conflict registry. EVERY transaction with a registry pins its
    /// snapshot (ReadCommitted session transactions also scan at their
    /// snapshot via `scan_table_at_snapshot`), so the version GC never
    /// reclaims a version an open transaction may still read.
    gc_pin_id: Option<u64>,
}

impl Drop for Transaction {
    fn drop(&mut self) {
        // Idempotent: deregister the snapshot however the transaction ends
        // (commit, rollback, or dropped on an error path) so registry
        // pruning is never blocked by a dead snapshot. Only validating
        // transactions register (only they read `recent_writes`).
        if self.conflict_validation {
            if let Some(registry) = &self.conflict_registry {
                registry.deregister_txn(self.transaction_id);
            }
        }
        // R4.3: release the version-GC pin on every exit path.
        if let Some(pin_id) = self.gc_pin_id.take() {
            if let Some(registry) = &self.conflict_registry {
                registry.unpin_snapshot(pin_id);
            }
        }
    }
}

/// Transaction state captured by a SAVEPOINT.
#[derive(Clone, Default)]
pub struct TransactionSavepointSnapshot {
    write_set: Vec<(Key, Option<Vec<u8>>)>,
    insert_log_len: usize,
}

impl Transaction {
    /// Create a new transaction (backwards compatible)
    pub fn new(db: Arc<DB>, snapshot_id: SnapshotId, snapshot_manager: Arc<SnapshotManager>) -> Result<Self> {
        static TXN_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
        let transaction_id = TXN_COUNTER.fetch_add(1, Ordering::SeqCst);

        debug!(
            txn_id = transaction_id,
            snapshot_id = snapshot_id,
            "Transaction started (legacy mode)"
        );

        Ok(Self {
            db,
            snapshot: Snapshot::new(snapshot_id),
            snapshot_ts: snapshot_id,
            transaction_id,
            snapshot_manager,
            write_set: Arc::new(DashMap::new()),
            insert_log: Arc::new(RwLock::new(Vec::new())),
            row_counter_stages: Arc::new(RwLock::new(std::collections::HashMap::new())),
            written_tables: DashSet::new(),
            has_unattributed_writes: AtomicBool::new(false),
            state: AtomicU8::new(TransactionState::Active.to_u8()),
            session_id: None,
            isolation_level: IsolationLevel::ReadCommitted,
            lock_manager: None,
            acquired_locks: Arc::new(RwLock::new(Vec::new())),
            dirty_tracker: None,
            versioning_enabled: true,
            rocksdb_wal_enabled: true,
            sync_commit: false,
            group_committer: None,
            row_cache: None,
            write_watermarks: None,
            conflict_registry: None,
            conflict_validation: false,
            gc_pin_id: None,
        })
    }

    /// Enable/disable MVCC version-history emission at commit (see field docs).
    /// Set by the engine from `StorageConfig::time_travel_enabled`.
    pub fn set_versioning_enabled(&mut self, enabled: bool) {
        self.versioning_enabled = enabled;
    }

    /// R1.3: request a synced (power-loss durable) commit WriteBatch.
    pub fn set_sync_commit(&mut self, enabled: bool) {
        self.sync_commit = enabled;
    }

    /// R1.3 phase 2: route this transaction's durability fsync through the
    /// engine's leader/follower group committer (see field docs).
    pub fn set_group_committer(&mut self, committer: Arc<super::group_commit::GroupCommitter>) {
        self.group_committer = Some(committer);
    }

    /// Attach the engine row cache so commit can invalidate written rows
    /// BEFORE lifting the snapshot barrier (see field docs).
    pub fn set_row_cache(&mut self, cache: Arc<super::RowCache>) {
        self.row_cache = Some(cache);
    }

    /// W2.5: attach the engine per-table committed-write watermark map so
    /// commit can raise each written table's watermark to the commit
    /// timestamp before the batch becomes visible (see field docs).
    pub fn set_write_watermarks(&mut self, watermarks: Arc<DashMap<String, u64>>) {
        self.write_watermarks = Some(watermarks);
    }

    /// True when the engine row cache is wired into this transaction, i.e.
    /// commit itself runs the stale-row fence and the caller-side
    /// `invalidate_row_cache_for` sweep would be redundant work on the
    /// commit hot path (R1.3-p2).
    pub fn has_row_cache(&self) -> bool {
        self.row_cache.is_some()
    }

    pub fn set_rocksdb_wal_enabled(&mut self, enabled: bool) {
        self.rocksdb_wal_enabled = enabled;
    }

    /// Create a new transaction with session and lock manager support
    pub fn new_with_session(
        db: Arc<DB>,
        snapshot_id: SnapshotId,
        snapshot_manager: Arc<SnapshotManager>,
        session_id: SessionId,
        isolation_level: IsolationLevel,
        lock_manager: Arc<LockManager>,
        dirty_tracker: Arc<DirtyTracker>,
    ) -> Result<Self> {
        static TXN_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
        let transaction_id = TXN_COUNTER.fetch_add(1, Ordering::SeqCst);

        debug!(
            txn_id = transaction_id,
            session_id = ?session_id,
            snapshot_id = snapshot_id,
            isolation_level = ?isolation_level,
            "Transaction started with session"
        );

        Ok(Self {
            db,
            snapshot: Snapshot::new(snapshot_id),
            snapshot_ts: snapshot_id,
            transaction_id,
            snapshot_manager,
            write_set: Arc::new(DashMap::new()),
            insert_log: Arc::new(RwLock::new(Vec::new())),
            row_counter_stages: Arc::new(RwLock::new(std::collections::HashMap::new())),
            written_tables: DashSet::new(),
            has_unattributed_writes: AtomicBool::new(false),
            state: AtomicU8::new(TransactionState::Active.to_u8()),
            session_id: Some(session_id),
            isolation_level,
            lock_manager: Some(lock_manager),
            acquired_locks: Arc::new(RwLock::new(Vec::new())),
            dirty_tracker: Some(dirty_tracker),
            versioning_enabled: true,
            rocksdb_wal_enabled: true,
            sync_commit: false,
            group_committer: None,
            row_cache: None,
            write_watermarks: None,
            conflict_registry: None,
            conflict_validation: false,
            gc_pin_id: None,
        })
    }

    /// Read a value
    ///
    /// Standard MVCC read: see snapshot-consistent version.
    /// Optimized with lock-free atomic state check and DashMap write_set lookup.
    ///
    /// v3.1.0: Enhanced with isolation level support
    /// - READ COMMITTED: No read locks (fresh snapshot per statement)
    /// - REPEATABLE READ: Acquires read locks on accessed rows
    /// - SERIALIZABLE: Acquires read locks on accessed rows (prevents phantom reads)
    pub fn get(&self, key: &Key) -> Result<Option<Vec<u8>>> {
        // Lock-free state check using atomic
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        // Acquire read lock based on isolation level
        if let Some(ref lock_mgr) = self.lock_manager {
            match self.isolation_level {
                IsolationLevel::RepeatableRead | IsolationLevel::Serializable => {
                    // Acquire read lock to prevent non-repeatable reads
                    let key_str = String::from_utf8_lossy(key).to_string();
                    let lock_guard = lock_mgr.acquire_lock(&key_str, self.transaction_id, LockType::Read)?;

                    // Store lock guard for automatic release on commit/rollback
                    self.acquired_locks.write().push(lock_guard);
                }
                IsolationLevel::ReadCommitted => {
                    // No read locks for READ COMMITTED
                }
            }
        }

        // Lock-free write set lookup using DashMap
        if let Some(entry) = self.write_set.get(key) {
            return Ok(entry.value().clone());
        }

        for (logged_key, logged_value) in self.insert_log.read().iter().rev() {
            if logged_key == key {
                return Ok(Some(logged_value.clone()));
            }
        }

        // Read from database at snapshot using MVCC
        self.read_at_version(key, self.snapshot_ts)
    }

    /// Read a versioned value at the transaction's snapshot timestamp
    ///
    /// Implements MVCC snapshot isolation by reading the latest version
    /// that is visible to this transaction (version timestamp <= snapshot_ts).
    ///
    /// Optimized with zero-copy key parsing to avoid allocations.
    ///
    /// Returns None if:
    /// - The key doesn't exist
    /// - All versions are newer than the snapshot timestamp
    /// - The key was deleted before the snapshot
    fn read_at_version(&self, key: &Key, snapshot_ts: u64) -> Result<Option<Vec<u8>>> {
        // Zero-copy key parsing optimization
        // Expected key format: "data:{table_name}:{row_id}"

        // Fast path: check prefix without UTF-8 conversion
        const DATA_PREFIX: &[u8] = b"data:";
        if !key.starts_with(DATA_PREFIX) {
            // Not a versioned data key, fallback to simple read
            return self
                .db
                .get(key)
                .map_err(|e| Error::storage(format!("Transaction get failed: {}", e)));
        }

        // Parse key with minimal allocations
        let key_str = std::str::from_utf8(key).map_err(|e| Error::storage(format!("Invalid key encoding: {}", e)))?;

        // Manual parsing without allocation - skip "data:" prefix
        let rest = &key_str[5..];

        // Find first colon position for table name
        let colon_pos = match rest.find(':') {
            Some(pos) => pos,
            None => {
                // Invalid format, fallback to simple read
                return self
                    .db
                    .get(key)
                    .map_err(|e| Error::storage(format!("Transaction get failed: {}", e)));
            }
        };

        let table_name = &rest[..colon_pos];
        let row_id_str = &rest[colon_pos + 1..];

        // Parse row ID directly from slice
        let row_id = match row_id_str.parse::<u64>() {
            Ok(id) => id,
            Err(_) => {
                // Invalid row ID format, fallback to simple read
                return self
                    .db
                    .get(key)
                    .map_err(|e| Error::storage(format!("Transaction get failed: {}", e)));
            }
        };

        // P0#1 read-side gate: when versioning is disabled, version history is
        // not maintained, so stale pre-existing v_idx: must not be consulted —
        // read the current committed `data:` value directly (mirrors the commit
        // write-gate and the engine's scan_table_at_snapshot gate).
        if !self.versioning_enabled {
            return self
                .db
                .get(key)
                .map_err(|e| Error::storage(format!("Transaction get failed: {}", e)));
        }

        // Use snapshot manager to read the versioned value
        // This implements the core MVCC logic: find the latest version <= snapshot_ts
        self.snapshot_manager.read_at_snapshot(table_name, row_id, snapshot_ts)
    }

    /// Write a value
    ///
    /// Buffered in write set until commit.
    /// Uses lock-free DashMap for concurrent access.
    ///
    /// v3.1.0: Enhanced with write lock acquisition for all isolation levels
    pub fn put(&self, key: Key, value: Vec<u8>) -> Result<()> {
        // Lock-free state check
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        // Acquire write lock (all isolation levels require write locks)
        if let Some(ref lock_mgr) = self.lock_manager {
            let key_str = String::from_utf8_lossy(&key).to_string();
            let lock_guard = lock_mgr.acquire_lock(&key_str, self.transaction_id, LockType::Write)?;

            // Store lock guard for automatic release on commit/rollback
            self.acquired_locks.write().push(lock_guard);
        }

        // Track dirty state with proper table/row granularity
        if let Some(ref tracker) = self.dirty_tracker {
            // Parse key to extract table_name and row_key for proper tracking
            // Key format: "data:{table_name}:{row_id}"
            const DATA_PREFIX: &[u8] = b"data:";
            if key.starts_with(DATA_PREFIX) {
                if let Ok(key_str) = std::str::from_utf8(&key) {
                    let rest = &key_str[5..]; // Skip "data:" prefix
                    if let Some(colon_pos) = rest.find(':') {
                        let table_name = &rest[..colon_pos];
                        let row_key = &rest[colon_pos + 1..];
                        let _ = tracker.track_insert(table_name, row_key, &value);
                    }
                }
            }
        }

        // R2.3: attribute the staged write to its table for the executor's
        // write-set-empty fast-path gate.
        self.note_written_key(&key);

        // Lock-free write_set insert using DashMap
        self.write_set.insert(key, Some(value));
        Ok(())
    }

    /// Stage a fast INSERT data row without the per-row DashMap overhead.
    ///
    /// This is used only for embedded explicit-transaction fast INSERT paths.
    /// Session transactions still use `put()` so lock and dirty tracking remain
    /// unchanged.
    pub fn put_insert_fast(&self, key: Key, value: Vec<u8>) -> Result<()> {
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        if self.lock_manager.is_some() || self.dirty_tracker.is_some() {
            return self.put(key, value);
        }

        // R2.3: attribute the staged insert to its table (see `written_tables`).
        self.note_written_key(&key);
        self.insert_log.write().push((key, value));
        Ok(())
    }

    /// Stage a table row-counter update to be written once at commit.
    ///
    /// Fast transactional INSERTs allocate row IDs from the engine's volatile
    /// counter. Persisting the counter through `write_set` on every row creates
    /// avoidable per-row work because each insert overwrites the same key. This
    /// side staging map keeps only the highest row ID per table and emits one
    /// `counter:<table>` write in the final WriteBatch.
    pub fn stage_row_counter(&self, table_name: &str, row_id: u64) -> Result<()> {
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        // R2.3: a staged row counter implies staged rows for this table.
        if !self.written_tables.contains(table_name) {
            self.written_tables.insert(table_name.to_string());
        }

        let mut stages = self.row_counter_stages.write();
        if let Some(current) = stages.get_mut(table_name) {
            if row_id > *current {
                *current = row_id;
            }
        } else {
            stages.insert(table_name.to_string(), row_id);
        }
        Ok(())
    }

    /// Delete a key
    ///
    /// Buffered as tombstone until commit.
    /// Uses lock-free DashMap for concurrent access.
    pub fn delete(&self, key: Key) -> Result<()> {
        // Lock-free state check
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        // R2.3: attribute the staged tombstone to its table (see `written_tables`).
        self.note_written_key(&key);

        // Lock-free write_set insert (tombstone) using DashMap
        self.write_set.insert(key, None);
        Ok(())
    }

    /// Update tuples within transaction
    ///
    /// Buffers updates in the transaction's write set instead of writing directly to storage.
    /// This ensures ACID guarantees - writes are only visible after commit.
    pub fn update_tuples(&self, table_name: &str, updates: Vec<(u64, crate::Tuple)>) -> Result<u64> {
        // Lock-free state check
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        let mut update_count = 0u64;

        for (row_id, tuple) in updates {
            // Format key for data table (using main branch format)
            let key = format!("data:{}:{}", table_name, row_id).into_bytes();

            // Serialize tuple
            let value = bincode::serialize(&tuple)
                .map_err(|e| crate::Error::storage(format!("Failed to serialize tuple: {}", e)))?;

            // Buffer in write set (will be applied on commit)
            self.put(key, value)?;
            update_count += 1;
        }

        Ok(update_count)
    }

    /// Delete tuples within transaction
    ///
    /// Buffers deletions as tombstones in the transaction's write set.
    /// This ensures ACID guarantees - deletions are only visible after commit.
    pub fn delete_tuples(&self, table_name: &str, row_ids: Vec<u64>) -> Result<u64> {
        // Lock-free state check
        let state_value = self.state.load(Ordering::Acquire);
        let state = TransactionState::from_u8(state_value);
        if state != TransactionState::Active {
            return Err(Error::transaction("Transaction is not active"));
        }

        let mut delete_count = 0u64;

        for row_id in row_ids {
            // Format key for data table (using main branch format)
            let key = format!("data:{}:{}", table_name, row_id).into_bytes();

            // Buffer tombstone in write set (will be applied on commit)
            self.delete(key)?;
            delete_count += 1;
        }

        Ok(delete_count)
    }

    /// Check if transaction is active
    ///
    /// Lock-free atomic check for maximum performance.
    pub fn is_active(&self) -> bool {
        let state_value = self.state.load(Ordering::Acquire);
        TransactionState::from_u8(state_value) == TransactionState::Active
    }

    /// Commit the transaction with a specific timestamp
    pub fn commit_with_timestamp(self, commit_ts: u64) -> Result<()> {
        let commit_start = std::time::Instant::now();
        let insert_count = self.insert_log.read().len();
        let row_counter_count = self.row_counter_stages.read().len();
        let write_count = self.write_set.len() + insert_count + row_counter_count;
        let lock_count = self.acquired_locks.read().len();

        trace!(
            txn_id = self.transaction_id,
            session_id = ?self.session_id,
            commit_ts = commit_ts,
            write_count = write_count,
            lock_count = lock_count,
            "Committing transaction"
        );

        // Check and transition state atomically
        let current = self.state.compare_exchange(
            TransactionState::Active.to_u8(),
            TransactionState::Committed.to_u8(),
            Ordering::AcqRel,
            Ordering::Acquire,
        );

        if current.is_err() {
            warn!(txn_id = self.transaction_id, "Commit failed: transaction not active");
            return Err(Error::transaction("Transaction is not active"));
        }

        // R0.2: first-committer-wins. Validate the write set against commits
        // newer than our snapshot and record ours — atomically w.r.t. other
        // committers. Runs before the batch is built; insert_log rows use
        // never-reused engine row ids and skip the registry by design.
        let inflight = self.has_tracked_writes();
        if let Some(registry) = &self.conflict_registry {
            if let Err((key, committed_ts)) =
                registry.validate_and_record(&self.write_set, self.conflict_validation, self.snapshot_ts, commit_ts)
            {
                if inflight {
                    registry.end_commit(commit_ts);
                }
                self.acquired_locks.write().clear();
                self.write_set.clear();
                self.insert_log.write().clear();
                self.row_counter_stages.write().clear();
                self.state.store(TransactionState::Aborted.to_u8(), Ordering::Release);
                warn!(
                    txn_id = self.transaction_id,
                    snapshot_ts = self.snapshot_ts,
                    committed_ts = committed_ts,
                    "Commit aborted: write-write conflict"
                );
                return Err(Error::transaction(format!(
                    "serialization failure: write-write conflict on key '{}' (committed at ts {}, transaction snapshot ts {}); retry the transaction",
                    String::from_utf8_lossy(&key),
                    committed_ts,
                    self.snapshot_ts
                )));
            }
        }

        let reverse_ts = u64::MAX - commit_ts;
        let commit_ts_text = if self.versioning_enabled {
            let mut buf = itoa::Buffer::new();
            buf.format(commit_ts).to_owned()
        } else {
            String::new()
        };
        let reverse_ts_text = if self.versioning_enabled {
            let mut buf = itoa::Buffer::new();
            let digits = buf.format(reverse_ts);
            let mut padded = String::with_capacity(20);
            for _ in digits.len()..20 {
                padded.push('0');
            }
            padded.push_str(digits);
            padded
        } else {
            String::new()
        };
        let commit_ts_bytes = commit_ts.to_be_bytes();

        // Apply write set atomically using RocksDB batch
        let mut batch = rocksdb::WriteBatch::default();
        let write_set_has_entries = !self.write_set.is_empty();
        // Item #2: hoist the "any COPY markers live?" check out of the per-entry
        // loop so a no-COPY commit pays one atomic load, not a `data:`-key parse
        // per write. Only when markers exist do we parse keys + materialize.
        let check_copy_markers = self.versioning_enabled && self.snapshot_manager.has_any_copy_markers();
        // W3.2: `elide` decides whether NEW versions written this commit skip the
        // `v:` copy (flagged `v_idx:` event). `check_elided` decides whether an
        // OVERWRITE must first materialize a prior flagged event — armed whenever
        // any flagged row might exist, including ones a prior session left when
        // this session has elision off (`maybe_elided_rows` ⊇ `elide_latest_version`).
        let elide = self.versioning_enabled && self.snapshot_manager.elide_latest_version();
        let check_elided = self.versioning_enabled && self.snapshot_manager.maybe_elided_rows();
        let mut version_key_buf = Vec::with_capacity(128);
        let mut version_index_key_buf = Vec::with_capacity(128);

        {
            let insert_log = self.insert_log.read();
            for (key, val) in insert_log.iter() {
                if write_set_has_entries && self.write_set.contains_key(key) {
                    continue;
                }
                Self::put_versioned_batch(
                    &mut batch,
                    key,
                    val,
                    self.versioning_enabled,
                    elide,
                    commit_ts_text.as_bytes(),
                    reverse_ts_text.as_bytes(),
                    &commit_ts_bytes,
                    &mut version_key_buf,
                    &mut version_index_key_buf,
                );
            }
        }

        // Iterate over DashMap entries
        let mut touches_columnar = false;
        for entry in self.write_set.iter() {
            let (key, value) = (entry.key(), entry.value());
            // R3.1/R3.3: columnar batches stage zone-stats (`colz:`) and
            // live-row presence (`colp:`) sidecars alongside the batch.
            // Applying them under the zone-stats write lock (below) keeps
            // lazy stats/presence backfill from racing this commit.
            if !touches_columnar && (key.starts_with(b"col:") || key.starts_with(b"colz:") || key.starts_with(b"colp:"))
            {
                touches_columnar = true;
            }
            // This commit is about to overwrite/delete this entry's `data:` row.
            // BEFORE it does, stage the pre-image version(s) into the SAME batch
            // (crash-atomic), so AS-OF reads predating this commit still resolve:
            //   * Item #2 — a COPY marker-covered row (its per-row `v:`/`v_idx:`
            //     were elided at COPY) gets its insert version from the pre-commit
            //     `data:`; else [copy_ts, commit_ts) would find no version.
            //   * W3.2 — a prior FLAGGED (elided) version of the row is
            //     materialized to a real `v:` with the flag cleared.
            // Both gates are hoisted above, so a plain commit parses no key.
            if check_copy_markers || check_elided {
                if let Some(rest) = key.strip_prefix(b"data:".as_slice()) {
                    if let Ok(text) = std::str::from_utf8(rest) {
                        if let Some(pos) = text.rfind(':') {
                            if let Ok(row_id) = text[pos + 1..].parse::<u64>() {
                                if check_copy_markers {
                                    self.snapshot_manager.materialize_copy_marker_row(
                                        &mut batch,
                                        &text[..pos],
                                        row_id,
                                        commit_ts,
                                    )?;
                                }
                                if check_elided {
                                    self.snapshot_manager.materialize_elided_latest_version(
                                        &mut batch,
                                        &text[..pos],
                                        row_id,
                                    )?;
                                }
                            }
                        }
                    }
                }
            }
            match value {
                Some(val) => {
                    Self::put_versioned_batch(
                        &mut batch,
                        key,
                        val,
                        self.versioning_enabled,
                        elide,
                        commit_ts_text.as_bytes(),
                        reverse_ts_text.as_bytes(),
                        &commit_ts_bytes,
                        &mut version_key_buf,
                        &mut version_index_key_buf,
                    );
                }
                None => {
                    batch.delete(key);
                    // W3.2: buffered DELETE tombstone (`data:` key bytes).
                    if crate::write_volume::enabled() {
                        crate::write_volume::add_row();
                        crate::write_volume::add(crate::write_volume::Category::Data, key.len() as u64);
                    }
                }
            }
        }

        for (table_name, row_id) in self.row_counter_stages.read().iter() {
            let key = format!("counter:{}", table_name);
            let value = match bincode::serialize(row_id) {
                Ok(value) => value,
                Err(e) => {
                    // Must NOT early-return with `?` here: an announced
                    // in-flight commit that never reaches `end_commit` would
                    // stall the snapshot barrier forever.
                    if inflight {
                        if let Some(registry) = &self.conflict_registry {
                            registry.end_commit(commit_ts);
                        }
                    }
                    self.acquired_locks.write().clear();
                    self.state.store(TransactionState::Aborted.to_u8(), Ordering::Release);
                    return Err(Error::storage(format!("Failed to serialize counter: {}", e)));
                }
            };
            batch.put(key.as_bytes(), value);
        }

        // R3.1: exclude zone-stats backfill while columnar batches + their
        // stats sidecars are applied (no-op guard for non-columnar commits).
        let _columnar_stats_guard = touches_columnar.then(super::columnar::stats_write_lock);

        // W2.5: raise the committed-write watermark for every table this
        // transaction wrote to `commit_ts` BEFORE the batch is applied. The
        // ordering is load-bearing (opposite of the row-cache fence below): a
        // concurrent in-transaction reader whose snapshot predates this commit
        // must see the raised watermark no later than it could observe this
        // commit's `data:` bytes, so it keeps taking the snapshot path for
        // these tables (fast-path readers also re-check the watermark after
        // their scan and discard a scan that raced a commit). `written_tables`
        // holds only main-branch `data:` tables — branch `bdata:` writes and
        // raw non-`data:` keys never enter it — matching the branch-blind map,
        // which is cleared on branch switch. Over-raising on a later failed
        // `db.write` only forces the safe snapshot path.
        if let Some(watermarks) = &self.write_watermarks {
            for table in self.written_tables.iter() {
                watermarks
                    .entry(table.key().clone())
                    .and_modify(|w| {
                        if commit_ts > *w {
                            *w = commit_ts;
                        }
                    })
                    .or_insert(commit_ts);
            }
        }

        // R1.3 phase 2: a durable commit with a group committer writes its
        // batch UNSYNCED here and pays durability via ONE cohort-wide
        // `flush_wal(true)` below — after the commit is applied, visible,
        // and its locks are released, so the fsync wait never convoys the
        // snapshot barrier or row-lock contenders. RocksDB's WAL is
        // sequential: the group fsync covers every byte written before it.
        let group_fsync = self.sync_commit && self.rocksdb_wal_enabled && self.group_committer.is_some();
        let result = if self.rocksdb_wal_enabled {
            if self.sync_commit && !group_fsync {
                // R1.3 phase 1 fallback (no group committer wired): one
                // synced WriteBatch per commit; RocksDB write groups
                // amortize it across concurrent committers.
                let mut write_opts = WriteOptions::default();
                write_opts.set_sync(true);
                self.db.write_opt(batch, &write_opts)
            } else {
                self.db.write(batch)
            }
        } else {
            let mut write_opts = WriteOptions::default();
            write_opts.set_sync(false);
            write_opts.disable_wal(true);
            self.db.write_opt(batch, &write_opts)
        };

        // Stale-row-cache fence: written rows must leave the row cache
        // BEFORE end_commit lifts the snapshot barrier, or a fresh snapshot
        // could pass the barrier and the UPDATE arm's PK point-lookup would
        // serve the cached pre-commit value — both transactions then pass
        // first-committer-wins validation and an update is lost (observed
        // ~10-50% per contended bench run on a loaded 32-core host). Callers
        // skip their post-commit `invalidate_row_cache_for` sweep when this
        // fence is wired (`has_row_cache`), so this is the only invalidation
        // pass on the hot path; the table name is borrowed straight out of
        // the key to avoid `written_data_keys`'s per-row String allocation.
        if result.is_ok() && !self.write_set.is_empty() {
            if let Some(cache) = &self.row_cache {
                for entry in self.write_set.iter() {
                    if let Some(rest) = entry.key().strip_prefix(b"data:".as_slice()) {
                        if let Ok(text) = std::str::from_utf8(rest) {
                            if let Some(pos) = text.rfind(':') {
                                if let Ok(row_id) = text[pos + 1..].parse::<u64>() {
                                    cache.invalidate(&text[..pos], row_id);
                                }
                            }
                        }
                    }
                }
            }
        }

        // The write is applied (or definitively failed): new snapshots may
        // proceed past this commit timestamp.
        if inflight {
            if let Some(registry) = &self.conflict_registry {
                registry.end_commit(commit_ts);
            }
        }

        // Release all acquired locks
        self.acquired_locks.write().clear();

        // Register the snapshot so it's visible to future queries
        if result.is_ok() {
            let _ = self.snapshot_manager.register_snapshot(commit_ts);
        }

        result.map_err(|e| {
            warn!(
                txn_id = self.transaction_id,
                error = %e,
                "Transaction commit failed, aborting"
            );
            self.state.store(TransactionState::Aborted.to_u8(), Ordering::Release);
            Error::transaction(format!("Commit failed: {}", e))
        })?;

        // R1.3 phase 2: group durability barrier. The commit is already
        // applied and visible (memtable + unsynced WAL); block until a
        // cohort fsync covering it completes. On failure the commit is NOT
        // power-loss durable: report the error (RocksDB records the WAL-sync
        // failure as a background error and rejects subsequent writes, so a
        // lost generation cannot be silently followed by a "durable" one).
        if group_fsync {
            if let Some(committer) = &self.group_committer {
                committer
                    .wait_durable(|| self.db.flush_wal(true).map_err(|e| e.to_string()))
                    .map_err(|e| {
                        warn!(
                            txn_id = self.transaction_id,
                            error = %e,
                            "Group WAL fsync failed; commit applied but not durable"
                        );
                        Error::transaction(format!("Commit failed: WAL group fsync failed: {}", e))
                    })?;
            }
        }

        debug!(
            txn_id = self.transaction_id,
            session_id = ?self.session_id,
            commit_ts = commit_ts,
            write_count = write_count,
            duration_us = commit_start.elapsed().as_micros() as u64,
            "Transaction committed successfully"
        );

        Ok(())
    }

    fn put_versioned_batch(
        batch: &mut rocksdb::WriteBatch,
        key: &[u8],
        val: &[u8],
        versioning_enabled: bool,
        elide: bool,
        commit_ts_text: &[u8],
        reverse_ts_text: &[u8],
        commit_ts_bytes: &[u8; 8],
        version_key_buf: &mut Vec<u8>,
        version_index_key_buf: &mut Vec<u8>,
    ) {
        batch.put(key, val);
        // W3.2: `data:` bytes for the buffered commit path (explicit txns and
        // autocommit-implicit txns). Class is the ambient scope — `other` for a
        // multi-statement txn whose writes land at COMMIT, decoupled from the
        // staging statement (documented in `W3_2_DESIGN.md`).
        if crate::write_volume::enabled() {
            crate::write_volume::add_row();
            crate::write_volume::add(crate::write_volume::Category::Data, (key.len() + val.len()) as u64);
        }

        // Create version history (value + reverse-ts index) for AS OF /
        // snapshot-history reads. P0#1: skip entirely when versioning is
        // disabled — saves 2 extra keys/row and a second serialization of the
        // row value. Read-committed reads use the `data:` key directly, so they
        // are unaffected.
        if versioning_enabled {
            Self::put_version_index_batch(
                batch,
                key,
                val,
                elide,
                commit_ts_text,
                reverse_ts_text,
                commit_ts_bytes,
                version_key_buf,
                version_index_key_buf,
            );
        }
    }

    fn put_version_index_batch(
        batch: &mut rocksdb::WriteBatch,
        key: &[u8],
        val: &[u8],
        elide: bool,
        commit_ts_text: &[u8],
        reverse_ts_text: &[u8],
        commit_ts_bytes: &[u8; 8],
        version_key_buf: &mut Vec<u8>,
        version_index_key_buf: &mut Vec<u8>,
    ) {
        let Some(rest) = key.strip_prefix(b"data:") else {
            return;
        };
        let Some(colon_pos) = rest.iter().position(|b| *b == b':') else {
            return;
        };
        let table_name = &rest[..colon_pos];
        let row_id = &rest[colon_pos + 1..];
        if table_name.is_empty() || row_id.is_empty() || !row_id.iter().all(u8::is_ascii_digit) {
            return;
        }

        // W3.2: when eliding, skip the `v:` copy entirely and flag the `v_idx:`
        // event value's high bit; the value is served from `data:` (this commit's
        // `data:` put, same batch) until the row's first later overwrite
        // materializes the real `v:`. `elide` mirrors the value the row-store /
        // `append_version_snapshot_to_batch` funnel writes for the fast path.
        if !elide {
            version_key_buf.clear();
            version_key_buf.extend_from_slice(b"v:");
            version_key_buf.extend_from_slice(table_name);
            version_key_buf.push(b':');
            version_key_buf.extend_from_slice(row_id);
            version_key_buf.push(b':');
            version_key_buf.extend_from_slice(commit_ts_text);
            batch.put(&version_key_buf, val);
        }

        version_index_key_buf.clear();
        version_index_key_buf.extend_from_slice(b"v_idx:");
        version_index_key_buf.extend_from_slice(table_name);
        version_index_key_buf.push(b':');
        version_index_key_buf.extend_from_slice(row_id);
        version_index_key_buf.push(b':');
        version_index_key_buf.extend_from_slice(reverse_ts_text);
        // Big-endian `commit_ts`, high bit = elided flag (see
        // `time_travel::VERSION_VALUE_ELIDED_FLAG`). Byte 0 is the MSB.
        let mut index_value = *commit_ts_bytes;
        if elide {
            if let Some(b0) = index_value.first_mut() {
                *b0 |= 0x80;
            }
        }
        batch.put(&version_index_key_buf, index_value);

        // W3.2: version-chain bytes for the buffered commit path. With elision on
        // the `v:` copy is gone, so this collapses to the `v_idx:` event size.
        if crate::write_volume::enabled() {
            let version_bytes = if elide {
                (version_index_key_buf.len() + index_value.len()) as u64
            } else {
                (version_key_buf.len() + val.len() + version_index_key_buf.len() + index_value.len()) as u64
            };
            crate::write_volume::add(crate::write_volume::Category::Version, version_bytes);
        }
    }

    /// Commit the transaction
    ///
    /// Atomically apply all buffered writes and create version index entries.
    /// Uses atomic state transition for consistency.
    ///
    /// v3.1.0: Enhanced with lock release and dirty tracker integration
    pub fn commit(self) -> Result<()> {
        let ts = self.snapshot_ts;
        self.commit_with_timestamp(ts)
    }

    /// Rollback the transaction
    ///
    /// Discard all buffered writes.
    /// Uses atomic state transition.
    ///
    /// v3.1.0: Enhanced with lock release
    pub fn rollback(self) -> Result<()> {
        let write_count = self.write_set.len() + self.insert_log.read().len();
        let lock_count = self.acquired_locks.read().len();

        debug!(
            txn_id = self.transaction_id,
            session_id = ?self.session_id,
            write_count = write_count,
            lock_count = lock_count,
            "Rolling back transaction"
        );

        // Check and transition state atomically
        let current = self.state.compare_exchange(
            TransactionState::Active.to_u8(),
            TransactionState::Aborted.to_u8(),
            Ordering::AcqRel,
            Ordering::Acquire,
        );

        if current.is_err() {
            warn!(txn_id = self.transaction_id, "Rollback failed: transaction not active");
            return Err(Error::transaction("Transaction is not active"));
        }

        // Release all acquired locks (RAII will handle this automatically when guards are dropped)
        // Explicitly clear the lock vector to trigger drops
        self.acquired_locks.write().clear();

        // Clear write set (DashMap handles concurrency)
        self.write_set.clear();
        self.insert_log.write().clear();
        self.row_counter_stages.write().clear();

        debug!(txn_id = self.transaction_id, "Transaction rolled back successfully");

        Ok(())
    }

    /// Get transaction state
    ///
    /// Lock-free atomic read.
    pub fn state(&self) -> TransactionState {
        let state_value = self.state.load(Ordering::Acquire);
        TransactionState::from_u8(state_value)
    }

    /// Get snapshot ID
    pub fn snapshot_id(&self) -> SnapshotId {
        self.snapshot_ts
    }

    /// Session that owns this transaction, if it was created through the
    /// per-session API (`new_with_session`). `None` for the embedded
    /// global-slot transaction path.
    pub fn session_id(&self) -> Option<SessionId> {
        self.session_id
    }

    /// Attach the engine's write-conflict registry (R0.2). `validate`
    /// enables first-committer-wins abort at commit; recording always
    /// happens so other transactions can validate against this one.
    ///
    /// Applies the snapshot barrier: this transaction's snapshot timestamp
    /// must not cover commits whose writes have not been applied yet, or
    /// its reads would be stale while validating clean.
    pub fn set_conflict_registry(&mut self, registry: Arc<WriteConflictRegistry>, validate: bool) {
        registry.snapshot_barrier(self.snapshot_ts);
        if validate {
            registry.register_txn(self.transaction_id, self.snapshot_ts);
        }
        // R4.3: pin this snapshot against MVCC version GC regardless of
        // isolation level — snapshot scans (`scan_table_at_snapshot`) read
        // version history for every transaction kind. Released in Drop.
        self.gc_pin_id = Some(registry.pin_snapshot(self.snapshot_ts));
        self.conflict_registry = Some(registry);
        self.conflict_validation = validate;
    }

    /// True when this commit must be announced in flight (it will record
    /// write-set keys in the conflict registry). INSERT-only transactions
    /// (`insert_log`) skip the registry by design.
    pub fn has_tracked_writes(&self) -> bool {
        self.conflict_registry.is_some() && !self.write_set.is_empty()
    }

    /// (table, row_id) pairs written through the write set, parsed from
    /// `data:{table}:{row_id}` keys. Commit callers use this to invalidate
    /// row caches after the batch is applied — a transaction commit writes
    /// RocksDB directly, and a stale row-cache entry would otherwise serve
    /// pre-commit values to later snapshot-clean transactions (R0.2
    /// lost-update vector via the UPDATE arm's PK point-lookup read).
    pub fn written_data_keys(&self) -> Vec<(String, u64)> {
        if self.write_set.is_empty() {
            return Vec::new();
        }
        let mut out = Vec::with_capacity(self.write_set.len());
        for entry in self.write_set.iter() {
            let key = entry.key();
            if let Some(rest) = key.strip_prefix(b"data:".as_slice()) {
                if let Ok(text) = std::str::from_utf8(rest) {
                    if let Some(pos) = text.rfind(':') {
                        if let Ok(row_id) = text[pos + 1..].parse::<u64>() {
                            out.push((text[..pos].to_string(), row_id));
                        }
                    }
                }
            }
        }
        out
    }

    /// R2.3: record which table a staged write belongs to.
    ///
    /// Keys shaped `data:{table}:{suffix}` are attributed to `{table}`; any
    /// other key flips `has_unattributed_writes`, which conservatively makes
    /// [`has_writes_for_table`](Self::has_writes_for_table) report `true` for
    /// every table (the staged write could affect anything, so every read in
    /// this transaction must stay on the write-set-merging slow path).
    fn note_written_key(&self, key: &[u8]) {
        if let Some(rest) = key.strip_prefix(b"data:".as_slice()) {
            if let Ok(text) = std::str::from_utf8(rest) {
                if let Some(pos) = text.find(':') {
                    let table = &text[..pos];
                    if !self.written_tables.contains(table) {
                        self.written_tables.insert(table.to_string());
                    }
                    return;
                }
            }
        }
        self.has_unattributed_writes.store(true, Ordering::Release);
    }

    /// R2.3: does this transaction hold staged writes that could affect reads
    /// of `table`?
    ///
    /// O(1): answered from the `written_tables` side-set maintained by
    /// `put`/`put_insert_fast`/`delete`/`stage_row_counter`, never by scanning
    /// `write_set` or `insert_log`. May over-approximate after
    /// ROLLBACK TO SAVEPOINT (the set is monotonic) and answers `true` for
    /// all tables once any non-`data:` key was staged — both directions of
    /// imprecision only force the (always-correct) slow read path.
    pub fn has_writes_for_table(&self, table: &str) -> bool {
        self.has_unattributed_writes.load(Ordering::Acquire) || self.written_tables.contains(table)
    }

    /// Isolation level this transaction runs under (R2.3: the executor's
    /// fast-path read gates are isolation-aware).
    pub fn isolation_level(&self) -> IsolationLevel {
        self.isolation_level
    }

    /// Refresh the snapshot timestamp to the current database state
    ///
    /// Useful for READ COMMITTED isolation level where each statement
    /// should see a fresh snapshot of the database.
    pub fn refresh_snapshot(&mut self, new_ts: u64) {
        if let Some(registry) = &self.conflict_registry {
            registry.snapshot_barrier(new_ts);
            if self.conflict_validation {
                registry.refresh_txn(self.transaction_id, new_ts);
            }
        }
        self.snapshot_ts = new_ts;
        self.snapshot = Snapshot::new(new_ts);
    }

    /// Create a snapshot of the current write set for savepoint support.
    ///
    /// Returns a copy of all key-value pairs currently in the write set.
    /// This snapshot can later be passed to `rollback_to_savepoint()` to
    /// restore the write set to this state.
    pub fn savepoint_snapshot(&self) -> TransactionSavepointSnapshot {
        TransactionSavepointSnapshot {
            write_set: self
                .write_set
                .iter()
                .map(|entry| (entry.key().clone(), entry.value().clone()))
                .collect(),
            insert_log_len: self.insert_log.read().len(),
        }
    }

    /// Rollback the write set to a previously captured savepoint snapshot.
    ///
    /// This removes any keys added after the savepoint was created and
    /// restores any keys that were modified to their values at savepoint time.
    /// Keys that existed at savepoint time but were deleted after are restored.
    ///
    /// # Arguments
    /// * `snapshot` - The write set snapshot captured at savepoint creation time
    pub fn rollback_to_savepoint(&self, snapshot: &TransactionSavepointSnapshot) {
        // Build a set of keys that existed at savepoint time for quick lookup
        let snapshot_keys: std::collections::HashSet<&Key> = snapshot.write_set.iter().map(|(k, _)| k).collect();

        // Remove keys that were added after the savepoint
        let current_keys: Vec<Key> = self.write_set.iter().map(|entry| entry.key().clone()).collect();

        for key in &current_keys {
            if !snapshot_keys.contains(key) {
                self.write_set.remove(key);
            }
        }

        // Restore keys to their savepoint values (handles modifications and re-inserts)
        for (key, value) in &snapshot.write_set {
            self.write_set.insert(key.clone(), value.clone());
        }

        self.insert_log.write().truncate(snapshot.insert_log_len);
    }

    /// Merge a set of tuples with the transaction's write set
    ///
    /// This ensures "read-your-own-writes" consistency for scans.
    /// It replaces tuples with newer versions in the write set and adds new tuples.
    pub fn merge_with_write_set(&self, table_name: &str, mut tuples: Vec<Tuple>) -> Result<Vec<Tuple>> {
        let prefix = format!("data:{}:", table_name);

        // Track which row IDs we've already handled from the base set
        let mut handled_row_ids = std::collections::HashSet::new();

        // 1. Update existing tuples from write set and handle tombstones
        let mut i = 0;
        while i < tuples.len() {
            let current_row_id = tuples.get(i).and_then(|t| t.row_id);
            if let Some(row_id) = current_row_id {
                handled_row_ids.insert(row_id);
                let key = format!("{}{}", prefix, row_id).into_bytes();

                if let Some(entry) = self.write_set.get(&key) {
                    match entry.value() {
                        Some(data) => {
                            // Replace with updated version
                            let mut updated_tuple: Tuple = bincode::deserialize(data)
                                .map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
                            updated_tuple.row_id = Some(row_id);
                            if let Some(slot) = tuples.get_mut(i) {
                                *slot = updated_tuple;
                            }
                            i += 1;
                        }
                        None => {
                            // Remove deleted tuple
                            tuples.remove(i);
                            // Don't increment i
                        }
                    }
                } else {
                    i += 1;
                }
            } else {
                i += 1;
            }
        }

        // 2. Add new tuples from write set that weren't in the base set
        for entry in self.write_set.iter() {
            let key = entry.key();
            if let Ok(key_str) = std::str::from_utf8(key) {
                if let Some(row_id_str) = key_str.strip_prefix(&prefix) {
                    if let Ok(row_id) = row_id_str.parse::<u64>() {
                        if !handled_row_ids.contains(&row_id) {
                            if let Some(data) = entry.value() {
                                let mut new_tuple: Tuple = bincode::deserialize(data)
                                    .map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
                                new_tuple.row_id = Some(row_id);
                                tuples.push(new_tuple);
                            }
                        }
                    }
                }
            }
        }

        let write_set_has_entries = !self.write_set.is_empty();
        for (key, data) in self.insert_log.read().iter() {
            if write_set_has_entries && self.write_set.contains_key(key) {
                continue;
            }
            if let Ok(key_str) = std::str::from_utf8(key) {
                if let Some(row_id_str) = key_str.strip_prefix(&prefix) {
                    if let Ok(row_id) = row_id_str.parse::<u64>() {
                        if !handled_row_ids.contains(&row_id) {
                            let mut new_tuple: Tuple = bincode::deserialize(data)
                                .map_err(|e| Error::storage(format!("Failed to deserialize tuple: {}", e)))?;
                            new_tuple.row_id = Some(row_id);
                            tuples.push(new_tuple);
                            handled_row_ids.insert(row_id);
                        }
                    }
                }
            }
        }

        Ok(tuples)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use crate::storage::StorageEngine;
    use crate::Config;

    #[test]
    fn test_transaction_commit() {
        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        let tx = engine.begin_transaction().expect("Failed to begin transaction");

        let key = b"test_key".to_vec();
        let value = b"test_value".to_vec();

        tx.put(key.clone(), value.clone()).expect("Failed to put value");
        tx.commit().expect("Failed to commit transaction");

        // Verify committed
        let result = engine.get(&key).expect("Failed to get value");
        assert_eq!(result, Some(value));
    }

    #[test]
    fn test_transaction_rollback() {
        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        let tx = engine.begin_transaction().expect("Failed to begin transaction");

        let key = b"test_key".to_vec();
        let value = b"test_value".to_vec();

        tx.put(key.clone(), value.clone()).expect("Failed to put value");
        tx.rollback().expect("Failed to rollback transaction");

        // Verify not committed
        let result = engine.get(&key).expect("Failed to get value");
        assert_eq!(result, None);
    }

    #[test]
    fn test_read_your_own_writes() {
        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        let tx = engine.begin_transaction().expect("Failed to begin transaction");

        let key = b"test_key".to_vec();
        let value = b"test_value".to_vec();

        tx.put(key.clone(), value.clone()).expect("Failed to put value");

        // Should see own writes before commit
        let result = tx.get(&key).expect("Failed to get value");
        assert_eq!(result, Some(value));
    }

    #[test]
    fn test_row_counter_staging_commits_latest_without_write_set_overwrites() {
        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        let tx = engine.begin_transaction().expect("Failed to begin transaction");
        tx.stage_row_counter("users", 3).unwrap();
        tx.stage_row_counter("users", 9).unwrap();
        tx.stage_row_counter("users", 7).unwrap();
        tx.stage_row_counter("orders", 4).unwrap();

        assert_eq!(tx.write_set.len(), 0);
        assert_eq!(tx.row_counter_stages.read().len(), 2);

        tx.commit().expect("Failed to commit transaction");

        let users_counter = engine
            .get(&b"counter:users".to_vec())
            .expect("Failed to read users counter")
            .expect("users counter missing");
        let users_counter: u64 = bincode::deserialize(&users_counter).unwrap();
        assert_eq!(users_counter, 9);

        let orders_counter = engine
            .get(&b"counter:orders".to_vec())
            .expect("Failed to read orders counter")
            .expect("orders counter missing");
        let orders_counter: u64 = bincode::deserialize(&orders_counter).unwrap();
        assert_eq!(orders_counter, 4);
    }

    #[test]
    fn test_mvcc_snapshot_isolation() {
        use crate::{Column, DataType, Schema, Tuple, Value};

        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        // Create a test table
        let schema = Schema {
            columns: vec![
                Column {
                    name: "id".to_string(),
                    data_type: DataType::Int4,
                    nullable: false,
                    primary_key: true,
                    source_table: None,
                    source_table_name: None,
                    default_expr: None,
                    unique: false,
                    storage_mode: crate::ColumnStorageMode::Default,
                },
                Column {
                    name: "value".to_string(),
                    data_type: DataType::Text,
                    nullable: false,
                    primary_key: false,
                    source_table: None,
                    source_table_name: None,
                    default_expr: None,
                    unique: false,
                    storage_mode: crate::ColumnStorageMode::Default,
                },
            ],
        };

        let catalog = engine.catalog();
        catalog
            .create_table("mvcc_test", schema)
            .expect("Failed to create table");

        // Insert initial data and create version
        let row_id = engine
            .insert_tuple_versioned(
                "mvcc_test",
                Tuple {
                    values: vec![Value::Int4(1), Value::String("initial".to_string())],
                    row_id: None,
                    branch_id: None,
                },
            )
            .expect("Failed to insert");

        let version1_ts = engine.current_timestamp();

        // Create snapshot manager reference for version writes
        let snapshot_mgr = engine.snapshot_manager();

        // Write a version at timestamp 1
        let key1 = format!("data:mvcc_test:{}", row_id).into_bytes();
        let value1 = bincode::serialize(&Tuple {
            values: vec![Value::Int4(1), Value::String("version1".to_string())],
            row_id: None,
            branch_id: None,
        })
        .expect("Failed to serialize");
        snapshot_mgr
            .write_version("mvcc_test", row_id, version1_ts, &value1)
            .expect("Failed to write version 1");

        // Start transaction 1 at this snapshot
        let tx1 = engine.begin_transaction().expect("Failed to begin tx1");

        // Update to version 2 (after tx1 started)
        let version2_ts = engine.current_timestamp();
        let value2 = bincode::serialize(&Tuple {
            values: vec![Value::Int4(1), Value::String("version2".to_string())],
            row_id: None,
            branch_id: None,
        })
        .expect("Failed to serialize");
        snapshot_mgr
            .write_version("mvcc_test", row_id, version2_ts, &value2)
            .expect("Failed to write version 2");

        // tx1 should still see version1 (snapshot isolation)
        let result = tx1.get(&key1).expect("Failed to read in tx1");
        assert!(result.is_some(), "tx1 should see version1");

        // Start transaction 2 at current snapshot (should see version2)
        let tx2 = engine.begin_transaction().expect("Failed to begin tx2");
        let result2 = tx2.get(&key1).expect("Failed to read in tx2");
        assert!(result2.is_some(), "tx2 should see version2");

        // Verify isolation: tx1 still sees version1 even after tx2 started
        let result_again = tx1.get(&key1).expect("Failed to read in tx1 again");
        assert_eq!(result, result_again, "tx1 should see consistent snapshot");
    }

    #[test]
    fn test_mvcc_concurrent_reads() {
        use crate::{Column, DataType, Schema, Tuple, Value};

        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        // Create table
        let schema = Schema {
            columns: vec![Column {
                name: "id".to_string(),
                data_type: DataType::Int4,
                nullable: false,
                primary_key: true,
                source_table: None,
                source_table_name: None,
                default_expr: None,
                unique: false,
                storage_mode: crate::ColumnStorageMode::Default,
            }],
        };

        let catalog = engine.catalog();
        catalog
            .create_table("concurrent_test", schema)
            .expect("Failed to create table");

        // Insert data
        let row_id = engine
            .insert_tuple_versioned(
                "concurrent_test",
                Tuple {
                    values: vec![Value::Int4(1)],
                    row_id: None,
                    branch_id: None,
                },
            )
            .expect("Failed to insert");

        let ts1 = engine.current_timestamp();
        let snapshot_mgr = engine.snapshot_manager();

        // Create version at ts1
        let value = bincode::serialize(&Tuple {
            values: vec![Value::Int4(100)],
            row_id: None,
            branch_id: None,
        })
        .expect("Failed to serialize");
        snapshot_mgr
            .write_version("concurrent_test", row_id, ts1, &value)
            .expect("Failed to write version");

        // Multiple concurrent transactions reading at same snapshot
        let tx1 = engine.begin_transaction().expect("Failed to begin tx1");
        let tx2 = engine.begin_transaction().expect("Failed to begin tx2");
        let tx3 = engine.begin_transaction().expect("Failed to begin tx3");

        let key = format!("data:concurrent_test:{}", row_id).into_bytes();

        // All should see same data (no phantom reads)
        let r1 = tx1.get(&key).expect("tx1 read failed");
        let r2 = tx2.get(&key).expect("tx2 read failed");
        let r3 = tx3.get(&key).expect("tx3 read failed");

        assert!(r1.is_some());
        assert_eq!(r1, r2);
        assert_eq!(r2, r3);
    }

    #[test]
    fn test_mvcc_deleted_row_visibility() {
        use crate::{Column, DataType, Schema, Tuple, Value};

        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        // Create table
        let schema = Schema {
            columns: vec![Column {
                name: "id".to_string(),
                data_type: DataType::Int4,
                nullable: false,
                primary_key: true,
                source_table: None,
                source_table_name: None,
                default_expr: None,
                unique: false,
                storage_mode: crate::ColumnStorageMode::Default,
            }],
        };

        let catalog = engine.catalog();
        catalog
            .create_table("delete_test", schema)
            .expect("Failed to create table");

        // Insert and version data
        let row_id = engine
            .insert_tuple_versioned(
                "delete_test",
                Tuple {
                    values: vec![Value::Int4(1)],
                    row_id: None,
                    branch_id: None,
                },
            )
            .expect("Failed to insert");

        let ts1 = engine.current_timestamp();
        let snapshot_mgr = engine.snapshot_manager();

        let value = bincode::serialize(&Tuple {
            values: vec![Value::Int4(42)],
            row_id: None,
            branch_id: None,
        })
        .expect("Failed to serialize");
        snapshot_mgr
            .write_version("delete_test", row_id, ts1, &value)
            .expect("Failed to write version");

        // Start tx1 (should see the row)
        let tx1 = engine.begin_transaction().expect("Failed to begin tx1");

        // Simulate delete by not creating a new version
        // (In real implementation, deletes would create tombstone versions)

        // tx1 should still see the row at its snapshot
        let key = format!("data:delete_test:{}", row_id).into_bytes();
        let result = tx1.get(&key).expect("Failed to read");
        assert!(result.is_some(), "tx1 should see row at its snapshot");
    }

    #[test]
    fn test_mvcc_read_at_version_parsing() {
        let config = Config::in_memory();
        let engine = StorageEngine::open_in_memory(&config).expect("Failed to open in-memory storage");

        let tx = engine.begin_transaction().expect("Failed to begin transaction");

        // Test non-versioned key (should fallback to simple read)
        let meta_key = b"meta:some_key".to_vec();
        let result = tx.get(&meta_key);
        assert!(result.is_ok(), "Should handle non-versioned keys");

        // Test invalid key format (should fallback to simple read)
        let invalid_key = b"data:invalid".to_vec();
        let result = tx.get(&invalid_key);
        assert!(result.is_ok(), "Should handle invalid key format");
    }
}