sochdb 2.0.2

// SPDX-License-Identifier: AGPL-3.0-or-later
// SochDB - LLM-Optimized Embedded Database
// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! ACID Transactions with MVCC
//!
//! Provides full transaction support:
//! - Atomicity: Buffered writes, all-or-nothing commit
//! - Consistency: Schema validation before commit
//! - Isolation: MVCC snapshots, read/write set tracking
//! - Durability: WAL with fsync, group commit
//!
//! ## Isolation Levels
//!
//! - ReadCommitted: See committed changes
//! - SnapshotIsolation: Consistent point-in-time view
//! - Serializable: Strongest isolation with conflict detection

use std::collections::{HashMap, HashSet};

use crate::connection::{Timestamp, SochConnection, TxnId};
use crate::error::{ClientError, Result};

/// Transaction isolation levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum IsolationLevel {
    /// Read committed - see committed changes
    ReadCommitted,
    /// Snapshot isolation - consistent point-in-time view (default)
    #[default]
    SnapshotIsolation,
    /// Serializable - strongest isolation
    Serializable,
}

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

/// Read set entry for conflict detection
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct TxnRead {
    pub table: String,
    pub key: Vec<u8>,
}

/// Write set entry
#[derive(Debug, Clone)]
pub struct TxnWrite {
    pub table: String,
    pub key: Vec<u8>,
    pub value: Option<Vec<u8>>, // None = delete
}

/// Client transaction handle with RAII cleanup
pub struct ClientTransaction<'a> {
    pub(crate) conn: &'a SochConnection,
    /// Transaction ID
    txn_id: TxnId,
    /// Start timestamp
    start_ts: Timestamp,
    /// Current state
    state: TxnState,
    /// Isolation level
    isolation: IsolationLevel,
    /// Buffered writes
    writes: Vec<TxnWrite>,
    /// Read set for conflict detection
    read_set: HashSet<TxnRead>,
    /// Local write cache for read-your-writes
    local_cache: HashMap<(String, Vec<u8>), Option<Vec<u8>>>,
    /// Auto-rollback on drop if not committed
    committed: bool,
}

impl<'a> ClientTransaction<'a> {
    /// Begin new transaction with isolation level
    pub fn begin(conn: &'a SochConnection, isolation: IsolationLevel) -> Result<Self> {
        let txn_id = conn.storage
            .begin_transaction()
            .map_err(|e| ClientError::Storage(e.to_string()))?;
        let start_ts = txn_id; // DurableStorage txn_id serves as logical timestamp

        Ok(Self {
            conn,
            txn_id,
            start_ts,
            state: TxnState::Active,
            isolation,
            writes: Vec::new(),
            read_set: HashSet::new(),
            local_cache: HashMap::new(),
            committed: false,
        })
    }

    /// Get transaction ID
    pub fn id(&self) -> TxnId {
        self.txn_id
    }

    /// Get start timestamp
    pub fn start_ts(&self) -> Timestamp {
        self.start_ts
    }

    /// Get transaction state
    pub fn state(&self) -> TxnState {
        self.state
    }

    /// Get isolation level
    pub fn isolation(&self) -> IsolationLevel {
        self.isolation
    }

    /// Check if read-only (no writes buffered)
    pub fn is_read_only(&self) -> bool {
        self.writes.is_empty()
    }

    /// Read with read-your-writes semantics
    pub fn get(&mut self, table: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        // Check local writes first (read-your-writes)
        let cache_key = (table.to_string(), key.to_vec());
        if let Some(value) = self.local_cache.get(&cache_key) {
            return Ok(value.clone());
        }

        // Record read for conflict detection (serializable)
        if self.isolation == IsolationLevel::Serializable {
            self.read_set.insert(TxnRead {
                table: table.to_string(),
                key: key.to_vec(),
            });
        }

        // Read from storage with MVCC visibility via DurableStorage
        self.conn
            .storage
            .read(self.txn_id, key)
            .map_err(|e| ClientError::Storage(format!("Read failed: {}", e)))
    }

    /// Write (buffered until commit)
    pub fn put(&mut self, table: &str, key: Vec<u8>, value: Vec<u8>) {
        // Cache for read-your-writes
        self.local_cache
            .insert((table.to_string(), key.clone()), Some(value.clone()));

        // Buffer write
        self.writes.push(TxnWrite {
            table: table.to_string(),
            key,
            value: Some(value),
        });
    }

    /// Delete (buffered until commit)
    pub fn delete(&mut self, table: &str, key: Vec<u8>) {
        // Cache tombstone for read-your-writes
        self.local_cache
            .insert((table.to_string(), key.clone()), None);

        // Buffer delete
        self.writes.push(TxnWrite {
            table: table.to_string(),
            key,
            value: None,
        });
    }

    /// Commit transaction with durability guarantee
    pub fn commit(mut self) -> Result<CommitResult> {
        if self.state != TxnState::Active {
            return Err(ClientError::Transaction("Transaction not active".into()));
        }

        // For serializable: check for conflicts
        if self.isolation == IsolationLevel::Serializable && !self.read_set.is_empty() {
            self.check_conflicts()?;
        }

        // Commit via DurableStorage (WAL + MVCC)
        let commit_ts = self.conn.storage
            .commit(self.txn_id)
            .map_err(|e| ClientError::Storage(e.to_string()))?;

        self.committed = true;
        self.state = TxnState::Committed;

        Ok(CommitResult {
            txn_id: self.txn_id,
            commit_ts,
            writes_count: self.writes.len(),
        })
    }

    /// Explicit rollback
    pub fn rollback(mut self) -> Result<()> {
        if self.state != TxnState::Active {
            return Err(ClientError::Transaction("Transaction not active".into()));
        }

        self.conn.storage
            .abort(self.txn_id)
            .map_err(|e| ClientError::Storage(e.to_string()))?;
        self.committed = true; // Prevent double-rollback in Drop
        self.state = TxnState::Aborted;

        Ok(())
    }

    /// Check for serializable conflicts
    fn check_conflicts(&self) -> Result<()> {
        // Placeholder - real impl would:
        // 1. Get concurrent transactions that committed after we started
        // 2. Check if their writes intersect our reads
        Ok(())
    }
}

/// RAII: Auto-rollback on drop if not committed
impl<'a> Drop for ClientTransaction<'a> {
    fn drop(&mut self) {
        if !self.committed && self.state == TxnState::Active {
            // Best-effort rollback via DurableStorage
            let _ = self.conn.storage.abort(self.txn_id);
        }
    }
}

/// Commit result
#[derive(Debug, Clone)]
pub struct CommitResult {
    pub txn_id: TxnId,
    pub commit_ts: Timestamp,
    pub writes_count: usize,
}

/// Read-only snapshot for consistent point-in-time queries
pub struct SnapshotReader<'a> {
    conn: &'a SochConnection,
    /// Snapshot timestamp
    snapshot_ts: Timestamp,
    /// Track visibility for diagnostics
    track_visibility: bool,
    /// Visibility log
    visibility_log: Vec<VisibilityCheck>,
}

/// Visibility check for diagnostics
#[derive(Debug, Clone)]
pub struct VisibilityCheck {
    pub key: Vec<u8>,
    pub visible: bool,
    pub reason: VisibilityReason,
}

/// Why a version was visible or not
#[derive(Debug, Clone)]
pub enum VisibilityReason {
    /// Committed before our start
    CommittedBeforeStart { commit_ts: Timestamp },
    /// Still in progress when we started
    InProgressAtStart { txn_id: TxnId },
    /// Committed after our start (not visible)
    CommittedAfterStart { commit_ts: Timestamp },
    /// Deleted before our start
    DeletedBeforeStart { delete_ts: Timestamp },
    /// Not yet committed
    Uncommitted,
}

impl<'a> SnapshotReader<'a> {
    /// Create snapshot at current timestamp
    pub fn now(conn: &'a SochConnection) -> Result<Self> {
        // Use DurableStorage begin_transaction to get a consistent snapshot point
        let snapshot_ts = conn.storage
            .begin_transaction()
            .map_err(|e| ClientError::Storage(e.to_string()))?;
        Ok(Self {
            conn,
            snapshot_ts,
            track_visibility: false,
            visibility_log: Vec::new(),
        })
    }

    /// Create snapshot at specific timestamp (historical read)
    pub fn at_timestamp(conn: &'a SochConnection, ts: Timestamp) -> Result<Self> {
        Ok(Self {
            conn,
            snapshot_ts: ts,
            track_visibility: false,
            visibility_log: Vec::new(),
        })
    }

    /// Enable visibility tracking for debugging
    pub fn with_visibility_tracking(mut self) -> Self {
        self.track_visibility = true;
        self
    }

    /// Get snapshot timestamp
    pub fn timestamp(&self) -> Timestamp {
        self.snapshot_ts
    }

    /// Read a key with MVCC visibility
    pub fn get(&mut self, table: &str, key: &[u8]) -> Result<Option<Vec<u8>>> {
        // Use the snapshot txn_id for MVCC-consistent reads
        self.conn
            .storage
            .read(self.snapshot_ts, key)
            .map_err(|e| ClientError::Storage(format!("Read failed: {}", e)))
    }

    /// Get visibility diagnostics
    pub fn visibility_diagnostics(&self) -> &[VisibilityCheck] {
        &self.visibility_log
    }
}

/// Batch transaction writer for high-throughput workloads
pub struct BatchWriter<'a> {
    conn: &'a SochConnection,
    /// Pending transactions
    #[allow(dead_code)]
    pending_txns: Vec<TxnId>,
    /// Pending writes
    pending_writes: Vec<TxnWrite>,
}

impl<'a> BatchWriter<'a> {
    /// Create new batch writer
    pub fn new(conn: &'a SochConnection) -> Self {
        Self {
            conn,
            pending_txns: Vec::new(),
            pending_writes: Vec::new(),
        }
    }

    /// Add a write to the batch
    pub fn write(&mut self, table: &str, key: Vec<u8>, value: Vec<u8>) {
        self.pending_writes.push(TxnWrite {
            table: table.to_string(),
            key,
            value: Some(value),
        });
    }

    /// Add a delete to the batch
    pub fn delete(&mut self, table: &str, key: Vec<u8>) {
        self.pending_writes.push(TxnWrite {
            table: table.to_string(),
            key,
            value: None,
        });
    }

    /// Flush all pending writes with single fsync
    pub fn flush(&mut self) -> Result<BatchCommitResult> {
        if self.pending_writes.is_empty() {
            return Ok(BatchCommitResult::default());
        }

        let start = std::time::Instant::now();
        let txn_id = self.conn.storage
            .begin_transaction()
            .map_err(|e| ClientError::Storage(e.to_string()))?;

        // Apply all pending writes within the transaction
        for write in &self.pending_writes {
            match &write.value {
                Some(value) => {
                    self.conn.storage
                        .write(txn_id, write.key.clone(), value.clone())
                        .map_err(|e| ClientError::Storage(e.to_string()))?;
                }
                None => {
                    self.conn.storage
                        .delete(txn_id, write.key.clone())
                        .map_err(|e| ClientError::Storage(e.to_string()))?;
                }
            }
        }

        let _commit_ts = self.conn.storage
            .commit(txn_id)
            .map_err(|e| ClientError::Storage(e.to_string()))?;
        let duration = start.elapsed();

        let count = self.pending_writes.len();
        self.pending_writes.clear();

        Ok(BatchCommitResult {
            transactions_committed: 1,
            writes_committed: count,
            fsync_latency: duration,
        })
    }

    /// Get pending write count
    pub fn pending_count(&self) -> usize {
        self.pending_writes.len()
    }
}

impl<'a> Drop for BatchWriter<'a> {
    fn drop(&mut self) {
        if !self.pending_writes.is_empty() {
            let _ = self.flush();
        }
    }
}

/// Batch commit result
#[derive(Debug, Clone, Default)]
pub struct BatchCommitResult {
    pub transactions_committed: usize,
    pub writes_committed: usize,
    pub fsync_latency: std::time::Duration,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_transaction_lifecycle() {
        let conn = SochConnection::open("./test").unwrap();

        let mut txn = ClientTransaction::begin(&conn, IsolationLevel::SnapshotIsolation).unwrap();
        assert_eq!(txn.state(), TxnState::Active);
        assert!(txn.is_read_only());

        txn.put("test", b"key".to_vec(), b"value".to_vec());
        assert!(!txn.is_read_only());

        let result = txn.commit().unwrap();
        assert!(result.writes_count > 0);
    }

    #[test]
    fn test_read_your_writes() {
        let conn = SochConnection::open("./test").unwrap();

        let mut txn = ClientTransaction::begin(&conn, IsolationLevel::SnapshotIsolation).unwrap();

        txn.put("test", b"key".to_vec(), b"value".to_vec());

        // Should see our own write
        let value = txn.get("test", b"key").unwrap();
        assert_eq!(value, Some(b"value".to_vec()));
    }

    #[test]
    fn test_rollback() {
        let conn = SochConnection::open("./test").unwrap();

        let mut txn = ClientTransaction::begin(&conn, IsolationLevel::SnapshotIsolation).unwrap();
        txn.put("test", b"key".to_vec(), b"value".to_vec());

        txn.rollback().unwrap();
        // Transaction should be aborted
    }

    #[test]
    fn test_snapshot_reader() {
        let conn = SochConnection::open("./test").unwrap();

        let snapshot = SnapshotReader::now(&conn).unwrap();
        assert!(snapshot.timestamp() > 0);
    }

    #[test]
    fn test_batch_writer() {
        let conn = SochConnection::open("./test").unwrap();

        let mut batch = BatchWriter::new(&conn);
        batch.write("test", b"k1".to_vec(), b"v1".to_vec());
        batch.write("test", b"k2".to_vec(), b"v2".to_vec());

        assert_eq!(batch.pending_count(), 2);

        let result = batch.flush().unwrap();
        assert_eq!(result.writes_committed, 2);
        assert_eq!(batch.pending_count(), 0);
    }

    #[test]
    fn test_isolation_levels() {
        assert_eq!(IsolationLevel::default(), IsolationLevel::SnapshotIsolation);
    }
}