motedb 0.8.2

//! Transaction Coordinator
//!
//! Manages transaction lifecycle: begin, commit, rollback
//! Provides snapshot isolation through MVCC

use crate::txn::version_store::{Snapshot, Timestamp, TransactionId, VersionStore};
use crate::types::{Row, RowId};
use crate::{Result, StorageError};
use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
use std::sync::Arc;

/// Delta operations for incremental snapshot
///
/// 🚀 P3.3 COW Optimization: Use Arc<Row> to avoid deep cloning
#[derive(Debug, Clone)]
pub enum DeltaOperation {
    /// Insert a new row (RowId didn't exist before)
    Insert(RowId, String, Arc<Row>),
    /// Update an existing row (store old value for rollback)
    Update(RowId, String, Arc<Row>), // old_value
    /// Delete a row (store old value for rollback)
    Delete(RowId, String, Arc<Row>), // old_value
}

/// Savepoint representation (Delta Snapshot optimized)
#[derive(Debug, Clone)]
pub struct Savepoint {
    /// Savepoint name
    pub name: String,

    /// 🚀 Delta: operations performed AFTER this savepoint was created
    /// Rolling back means undoing these operations in reverse order
    pub write_deltas: Vec<DeltaOperation>,

    /// 🚀 Delta: read_set additions AFTER this savepoint
    /// Only track new reads, not entire read_set
    pub read_deltas: HashSet<RowId>,
}

/// Transaction isolation level
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IsolationLevel {
    /// Read uncommitted data (not recommended)
    ReadUncommitted = 0,
    /// Read only committed data
    ReadCommitted = 1,
    /// Repeatable reads within transaction
    RepeatableRead = 2,
    /// Full serializable isolation
    Serializable = 3,
}

/// Transaction state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum TransactionState {
    Active = 0,
    Committed = 1,
    Aborted = 2,
}

/// Transaction context
pub struct TransactionContext {
    /// Transaction ID
    pub txn_id: TransactionId,

    /// Start timestamp
    pub start_ts: Timestamp,

    /// Isolation level
    pub isolation_level: IsolationLevel,

    /// Transaction state
    pub state: AtomicU8,

    /// Write set (local cache of writes) with table metadata
    /// Format: (table_name, RowId) → row_data
    /// 🔑 Keyed by (table_name, row_id) — NOT just row_id — because different
    /// tables can have the same row_id (e.g. both have PK=2 → row_id=2).
    /// Keying by row_id alone caused cross-table INSERTs to overwrite each
    /// other, silently losing transactional writes.
    pub write_set: RwLock<HashMap<(String, RowId), Row>>,

    /// Read set (for conflict detection in Serializable)
    pub read_set: RwLock<HashSet<RowId>>,

    /// Undo log: operations that directly modified storage (UPDATE/DELETE)
    /// and need to be reversed on ROLLBACK. Each entry stores the old value
    /// so rollback can restore it. Unlike write_set (which buffers INSERTs),
    /// these ops already hit the store — rollback must actively undo them.
    pub undo_log: RwLock<Vec<DeltaOperation>>,

    /// Snapshot for this transaction
    pub snapshot: Snapshot,

    /// Savepoint stack (for partial rollback)
    /// Savepoints are stacked: [sp1, sp2, sp3] where sp3 is the most recent
    pub savepoints: RwLock<Vec<Savepoint>>,
}

impl TransactionContext {}

/// Transaction Coordinator
pub struct TransactionCoordinator {
    /// Version store
    version_store: Arc<VersionStore>,

    /// Active transactions
    active_txns: Arc<DashMap<TransactionId, Arc<TransactionContext>>>,

    /// Transaction ID generator
    txn_id_gen: Arc<AtomicU64>,
}

impl TransactionCoordinator {
    /// Create a new transaction coordinator
    pub fn new(version_store: Arc<VersionStore>) -> Self {
        Self {
            version_store,
            active_txns: Arc::new(DashMap::new()),
            txn_id_gen: Arc::new(AtomicU64::new(1)),
        }
    }

    /// Begin a new transaction
    pub fn begin(&self, isolation_level: IsolationLevel) -> Result<TransactionId> {
        let txn_id = self.txn_id_gen.fetch_add(1, Ordering::SeqCst);
        let start_ts = self.version_store.allocate_timestamp();

        // Create snapshot
        let snapshot = self.create_snapshot_internal(txn_id, start_ts)?;

        let ctx = Arc::new(TransactionContext {
            txn_id,
            start_ts,
            isolation_level,
            state: AtomicU8::new(TransactionState::Active as u8),
            write_set: RwLock::new(HashMap::new()),
            read_set: RwLock::new(HashSet::new()),
            undo_log: RwLock::new(Vec::new()),
            snapshot,
            savepoints: RwLock::new(Vec::new()), // Initialize empty savepoint stack
        });

        self.active_txns.insert(txn_id, ctx);

        Ok(txn_id)
    }

    /// Commit a transaction
    ///
    /// Returns the commit timestamp for MVCC visibility
    pub fn commit(&self, txn_id: TransactionId) -> Result<Timestamp> {
        let ctx = self.get_context(txn_id)?;

        // Check if already committed or aborted
        let state = ctx.state.load(Ordering::Acquire);
        if state != TransactionState::Active as u8 {
            return Err(StorageError::Transaction(format!(
                "Transaction {} is not active",
                txn_id
            )));
        }

        // Get commit timestamp
        let commit_ts = self.version_store.allocate_timestamp();

        // For Serializable isolation: validate read-set (no write locks needed —
        // this is a read-only consistency check and doesn't have a TOCTOU issue
        // because reads don't modify state).
        if ctx.isolation_level == IsolationLevel::Serializable {
            self.validate_read_set(&ctx)?;
        }

        // Sort rows by row_id to avoid deadlocks when acquiring write locks
        let write_set = ctx.write_set.read();
        let mut sorted_rows: Vec<((String, RowId), Row)> = write_set
            .iter()
            .map(|(key, data)| (key.clone(), data.clone()))
            .collect();
        drop(write_set);
        sorted_rows.sort_by_key(|((_, rid), _)| *rid);

        // Insert each row version with atomic validation (validation + insertion
        // happen under a single write lock, eliminating the TOCTOU window).
        for ((_, row_id), data) in &sorted_rows {
            self.version_store.insert_version_atomic(
                *row_id,
                data.clone(),
                txn_id,
                commit_ts,
                &ctx.snapshot,
            )?;
        }

        // Mark as committed
        ctx.state
            .store(TransactionState::Committed as u8, Ordering::Release);

        // Remove from active transactions
        self.active_txns.remove(&txn_id);

        Ok(commit_ts)
    }

    /// Rollback a transaction
    pub fn rollback(&self, txn_id: TransactionId) -> Result<()> {
        let ctx = self.get_context(txn_id)?;

        // 🔑 Undo log replay is handled by the executor (which has db access)
        // before calling this method. Clear the undo log here for cleanliness.
        ctx.undo_log.write().clear();

        // Clear write set
        ctx.write_set.write().clear();

        // Mark as aborted
        ctx.state
            .store(TransactionState::Aborted as u8, Ordering::Release);

        // Remove from active transactions
        self.active_txns.remove(&txn_id);

        Ok(())
    }

    /// Create a savepoint in the current transaction (Delta Snapshot optimized)
    ///
    /// 🚀 Memory Optimization: Instead of cloning entire write_set,
    /// we create an empty delta tracker. Future operations will append to this delta.
    /// Memory usage: O(1) at creation time, O(k) for k operations after savepoint.
    pub fn create_savepoint(&self, txn_id: TransactionId, name: String) -> Result<()> {
        let ctx = self.get_context(txn_id)?;

        // Check if transaction is active
        let state = ctx.state.load(Ordering::Acquire);
        if state != TransactionState::Active as u8 {
            return Err(StorageError::Transaction(format!(
                "Transaction {} is not active",
                txn_id
            )));
        }

        // 🚀 Delta Snapshot: Start with empty deltas
        // Operations after this point will be tracked incrementally
        let savepoint = Savepoint {
            name: name.clone(),
            write_deltas: Vec::new(), // No memory allocation at creation
            read_deltas: HashSet::new(),
        };

        // Push to savepoint stack
        ctx.savepoints.write().push(savepoint);

        debug_log!(
            "[Savepoint] Created delta savepoint '{}' for txn {} (mem: 0 bytes)",
            name,
            txn_id
        );

        Ok(())
    }

    /// Record a write delta to the most recent savepoint (if any exist).
    /// This enables rollback_to_savepoint to undo operations.
    pub fn record_write_delta(&self, txn_id: TransactionId, delta: DeltaOperation) -> Result<()> {
        let ctx = self.get_context(txn_id)?;
        // Record in undo_log for full-transaction rollback.
        ctx.undo_log.write().push(delta.clone());
        // Also record in the latest savepoint if one exists.
        let mut savepoints = ctx.savepoints.write();
        if let Some(last) = savepoints.last_mut() {
            last.write_deltas.push(delta);
        }
        Ok(())
    }

    /// 🔑 Update the write_set entry for a row that was INSERTed in this
    /// transaction and is now being UPDATEd. Without this, COMMIT flushes
    /// the stale INSERT value (overwriting the UPDATE). Returns true if
    /// the write_set was updated (row was found in it).
    pub fn update_write_set_row(
        &self,
        txn_id: TransactionId,
        table_name: &str,
        row_id: RowId,
        new_row: Row,
    ) -> Result<bool> {
        let ctx = self.get_context(txn_id)?;
        let mut write_set = ctx.write_set.write();
        let key = (table_name.to_string(), row_id);
        if write_set.contains_key(&key) {
            write_set.insert(key, new_row);
            Ok(true)
        } else {
            Ok(false)
        }
    }
    /// Rollback to a savepoint (Delta Snapshot optimized)
    ///
    /// 🚀 Memory Optimization: Instead of restoring full snapshot,
    /// we undo operations in reverse order using deltas.
    ///
    /// Algorithm:
    /// 1. Collect all deltas from savepoint[position+1..end] in reverse
    /// 2. Apply undo operations:
    ///    - Insert → Remove from write_set
    ///    - Update → Restore old value
    ///    - Delete → Restore old value
    /// 3. Remove savepoint[position..end] from stack
    pub fn rollback_to_savepoint(&self, txn_id: TransactionId, name: &str) -> Result<()> {
        let ctx = self.get_context(txn_id)?;

        // Check if transaction is active
        let state = ctx.state.load(Ordering::Acquire);
        if state != TransactionState::Active as u8 {
            return Err(StorageError::Transaction(format!(
                "Transaction {} is not active",
                txn_id
            )));
        }

        // Phase 1: Collect undo data from savepoints, then DROP savepoints lock.
        // This prevents deadlock with record_write_delta() which acquires
        // savepoints.write() (callers may hold write_set.read()).
        let (position, all_deltas, read_deltas) = {
            let savepoints = ctx.savepoints.write();

            // Find the savepoint by name
            let position = savepoints
                .iter()
                .position(|sp| sp.name == name)
                .ok_or_else(|| {
                    StorageError::Transaction(format!(
                        "Savepoint '{}' not found in transaction {}",
                        name, txn_id
                    ))
                })?;

            // Collect all deltas from savepoints AFTER this one (in reverse order)
            let mut all_deltas = Vec::new();
            for sp in savepoints.iter().skip(position + 1).rev() {
                for delta in sp.write_deltas.iter().rev() {
                    all_deltas.push(delta.clone());
                }
            }

            // Also include deltas from the target savepoint itself
            for delta in savepoints[position].write_deltas.iter().rev() {
                all_deltas.push(delta.clone());
            }

            // Collect read deltas for undo
            let mut read_deltas = Vec::new();
            for sp in savepoints.iter().skip(position).rev() {
                for row_id in &sp.read_deltas {
                    read_deltas.push(*row_id);
                }
            }

            (position, all_deltas, read_deltas)
        }; // savepoints lock dropped here — safe to acquire write_set/write

        // Phase 2: Apply undo operations (P3.3 COW: Arc clone is cheap)
        let mut write_set = ctx.write_set.write();
        let mut read_set = ctx.read_set.write();

        let mut undo_count = 0;
        for delta in all_deltas {
            match delta {
                DeltaOperation::Insert(row_id, table_name, _new_value) => {
                    // Undo insert: remove from write_set
                    write_set.remove(&(table_name, row_id));
                    undo_count += 1;
                }
                DeltaOperation::Update(row_id, table_name, old_value) => {
                    // Undo update: restore old value (Arc clone is O(1))
                    write_set.insert(
                        (table_name, row_id),
                        Arc::try_unwrap(old_value).unwrap_or_else(|arc| (*arc).clone()),
                    );
                    undo_count += 1;
                }
                DeltaOperation::Delete(row_id, table_name, old_value) => {
                    // Undo delete: restore old value (Arc clone is O(1))
                    write_set.insert(
                        (table_name, row_id),
                        Arc::try_unwrap(old_value).unwrap_or_else(|arc| (*arc).clone()),
                    );
                    undo_count += 1;
                }
            }
        }

        // Undo read_set changes
        for row_id in &read_deltas {
            read_set.remove(row_id);
        }

        drop(write_set);
        drop(read_set);

        // Phase 3: Truncate savepoints (re-acquire after dropping write_set/read_set)
        {
            let mut savepoints = ctx.savepoints.write();
            savepoints.truncate(position);
        }

        debug_log!(
            "[Savepoint] Rolled back to '{}' in txn {} (undid {} ops)",
            name,
            txn_id,
            undo_count
        );

        Ok(())
    }

    /// Release a savepoint
    ///
    /// Removes the savepoint but keeps all changes made after it.
    /// Useful for cleaning up savepoints you no longer need.
    pub fn release_savepoint(&self, txn_id: TransactionId, name: &str) -> Result<()> {
        let ctx = self.get_context(txn_id)?;

        // Check if transaction is active
        let state = ctx.state.load(Ordering::Acquire);
        if state != TransactionState::Active as u8 {
            return Err(StorageError::Transaction(format!(
                "Transaction {} is not active",
                txn_id
            )));
        }

        let mut savepoints = ctx.savepoints.write();

        // Find and remove the savepoint
        let position = savepoints
            .iter()
            .position(|sp| sp.name == name)
            .ok_or_else(|| {
                StorageError::Transaction(format!(
                    "Savepoint '{}' not found in transaction {}",
                    name, txn_id
                ))
            })?;

        savepoints.remove(position);

        debug_log!(
            "[Savepoint] Released savepoint '{}' in txn {}",
            name,
            txn_id
        );

        Ok(())
    }

    /// Get transaction context
    pub fn get_context(&self, txn_id: TransactionId) -> Result<Arc<TransactionContext>> {
        self.active_txns
            .get(&txn_id)
            .map(|entry| entry.value().clone())
            .ok_or_else(|| StorageError::Transaction(format!("Transaction {} not found", txn_id)))
    }

    /// Create snapshot for a transaction
    fn create_snapshot_internal(
        &self,
        txn_id: TransactionId,
        timestamp: Timestamp,
    ) -> Result<Snapshot> {
        // Collect all active transaction IDs except this one
        let active_txns: HashSet<TransactionId> = self
            .active_txns
            .iter()
            .filter(|entry| *entry.key() != txn_id)
            .map(|entry| *entry.key())
            .collect();

        Ok(Snapshot {
            timestamp,
            active_txns,
        })
    }

    /// Validate write set for conflicts
    /// Validate read-set for Serializable isolation (read-write conflict detection).
    /// Write-write conflict detection is handled atomically by insert_version_atomic.
    fn validate_read_set(&self, ctx: &TransactionContext) -> Result<()> {
        let read_set = ctx.read_set.read();
        for row_id in read_set.iter() {
            if let Some(chain) = self.version_store.versions.get(row_id) {
                let head = chain.head.read();
                if let Some(version) = head.as_ref() {
                    if version.begin_ts > ctx.snapshot.timestamp && version.txn_id != ctx.txn_id {
                        return Err(StorageError::Transaction(format!(
                            "Read-write conflict on row {} in txn {}",
                            row_id, ctx.txn_id
                        )));
                    }
                }
            }
        }
        Ok(())
    }

    /// Get minimum active timestamp (for vacuum)
    pub fn get_min_active_timestamp(&self) -> Timestamp {
        self.active_txns
            .iter()
            .map(|entry| entry.value().start_ts)
            .min()
            .unwrap_or(self.version_store.current_timestamp())
    }

    /// Get statistics
    pub fn stats(&self) -> TransactionCoordinatorStats {
        let next_txn_id = self.txn_id_gen.load(Ordering::Relaxed);
        let active = self.active_txns.len() as u64;
        // total_committed = (next_txn_id - 1) - active
        // next_txn_id - 1 because we start from 1 and fetch_add returns old value
        let total_committed = if next_txn_id > 1 {
            (next_txn_id - 1) - active
        } else {
            0
        };

        TransactionCoordinatorStats {
            active_transactions: active,
            total_committed,
        }
    }
}

/// Transaction coordinator statistics
#[derive(Debug, Clone)]
pub struct TransactionCoordinatorStats {
    pub active_transactions: u64,
    pub total_committed: u64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Timestamp, Value};

    fn create_test_coordinator() -> TransactionCoordinator {
        let version_store = Arc::new(VersionStore::new());
        TransactionCoordinator::new(version_store)
    }

    #[test]
    fn test_begin_transaction() {
        let coord = create_test_coordinator();
        let txn_id = coord.begin(IsolationLevel::ReadCommitted).unwrap();

        assert!(coord.active_txns.contains_key(&txn_id));

        let ctx = coord.get_context(txn_id).unwrap();
        assert_eq!(ctx.txn_id, txn_id);
        assert_eq!(ctx.isolation_level, IsolationLevel::ReadCommitted);
    }

    #[test]
    fn test_commit_transaction() {
        let coord = create_test_coordinator();
        let txn_id = coord.begin(IsolationLevel::ReadCommitted).unwrap();

        // Add some writes (with table metadata)
        let ctx = coord.get_context(txn_id).unwrap();
        ctx.write_set.write().insert(
            ("test_table".to_string(), 1),
            vec![Value::Timestamp(Timestamp::from_micros(100))],
        );

        coord.commit(txn_id).unwrap();

        assert!(!coord.active_txns.contains_key(&txn_id));
    }

    #[test]
    fn test_rollback_transaction() {
        let coord = create_test_coordinator();
        let txn_id = coord.begin(IsolationLevel::ReadCommitted).unwrap();

        // Add some writes (with table metadata)
        let ctx = coord.get_context(txn_id).unwrap();
        ctx.write_set.write().insert(
            ("test_table".to_string(), 1),
            vec![Value::Timestamp(Timestamp::from_micros(100))],
        );

        coord.rollback(txn_id).unwrap();

        assert!(!coord.active_txns.contains_key(&txn_id));
    }

    #[test]
    fn test_snapshot_isolation() {
        let coord = create_test_coordinator();

        let txn1 = coord.begin(IsolationLevel::RepeatableRead).unwrap();
        let txn2 = coord.begin(IsolationLevel::RepeatableRead).unwrap();

        let ctx1 = coord.get_context(txn1).unwrap();
        let ctx2 = coord.get_context(txn2).unwrap();

        // txn1 was created first, so it doesn't see txn2 in its snapshot
        assert!(!ctx1.snapshot.active_txns.contains(&txn2));
        // txn2 was created second, so it sees txn1 as active
        assert!(ctx2.snapshot.active_txns.contains(&txn1));
    }

    #[test]
    fn test_concurrent_transactions() {
        let coord = Arc::new(create_test_coordinator());
        let mut handles = vec![];

        for _ in 0..10 {
            let coord = coord.clone();
            let handle = std::thread::spawn(move || {
                let txn = coord.begin(IsolationLevel::ReadCommitted).unwrap();
                std::thread::sleep(std::time::Duration::from_millis(10));
                coord.commit(txn).unwrap();
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().unwrap();
        }

        let stats = coord.stats();
        assert_eq!(stats.total_committed, 10);
        assert_eq!(stats.active_transactions, 0);
    }
}