rexis_rag/incremental/

rollback.rs

//! # Rollback System
//!
//! Provides comprehensive rollback capabilities for failed operations.
//! Includes operation logging, state snapshots, and recovery mechanisms.

use crate::incremental::index_manager::{IndexUpdate, UpdateResult};
use crate::RragResult;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

/// Rollback system configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackConfig {
    /// Maximum operations to keep in log
    pub max_operation_log_size: usize,

    /// Enable state snapshots
    pub enable_snapshots: bool,

    /// Snapshot interval (number of operations)
    pub snapshot_interval: usize,

    /// Maximum snapshots to retain
    pub max_snapshots: usize,

    /// Enable automatic rollback on failures
    pub enable_auto_rollback: bool,

    /// Rollback timeout in seconds
    pub rollback_timeout_secs: u64,

    /// Enable rollback verification
    pub enable_verification: bool,

    /// Rollback strategy
    pub rollback_strategy: RollbackStrategy,
}

/// Rollback strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RollbackStrategy {
    /// Roll back to last known good state
    LastKnownGood,
    /// Roll back to specific snapshot
    SpecificSnapshot,
    /// Selective rollback of failed operations
    Selective,
    /// Complete system rollback
    Complete,
    /// Custom rollback logic
    Custom(String),
}

impl Default for RollbackConfig {
    fn default() -> Self {
        Self {
            max_operation_log_size: 10000,
            enable_snapshots: true,
            snapshot_interval: 100,
            max_snapshots: 50,
            enable_auto_rollback: true,
            rollback_timeout_secs: 300,
            enable_verification: true,
            rollback_strategy: RollbackStrategy::LastKnownGood,
        }
    }
}

/// Rollback operation types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RollbackOperation {
    /// Restore from snapshot
    RestoreSnapshot {
        snapshot_id: String,
        target_state: SystemState,
    },

    /// Undo specific operations
    UndoOperations { operation_ids: Vec<String> },

    /// Revert to timestamp
    RevertToTimestamp {
        timestamp: chrono::DateTime<chrono::Utc>,
    },

    /// Selective document rollback
    SelectiveRollback {
        document_ids: Vec<String>,
        target_versions: HashMap<String, String>,
    },

    /// Complete system reset
    SystemReset { reset_to_snapshot: String },
}

/// System state snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemState {
    /// Snapshot ID
    pub snapshot_id: String,

    /// Snapshot timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,

    /// Document states at time of snapshot
    pub document_states: HashMap<String, DocumentState>,

    /// Index states
    pub index_states: HashMap<String, IndexState>,

    /// System metadata at snapshot time
    pub system_metadata: HashMap<String, serde_json::Value>,

    /// Operations count at snapshot time
    pub operations_count: u64,

    /// Snapshot size in bytes
    pub size_bytes: u64,

    /// Snapshot compression ratio
    pub compression_ratio: f64,
}

/// Document state in snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocumentState {
    /// Document ID
    pub document_id: String,

    /// Document version at snapshot time
    pub version_id: String,

    /// Content hash
    pub content_hash: String,

    /// Metadata hash
    pub metadata_hash: String,

    /// Chunk states
    pub chunk_states: Vec<ChunkState>,

    /// Embedding states
    pub embedding_states: Vec<EmbeddingState>,
}

/// Chunk state in snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChunkState {
    /// Chunk index
    pub chunk_index: usize,

    /// Chunk hash
    pub chunk_hash: String,

    /// Chunk size
    pub size_bytes: usize,
}

/// Embedding state in snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingState {
    /// Embedding ID
    pub embedding_id: String,

    /// Source chunk or document
    pub source_id: String,

    /// Embedding vector (simplified)
    pub vector_hash: String,

    /// Embedding metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Index state in snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexState {
    /// Index name
    pub index_name: String,

    /// Index type
    pub index_type: String,

    /// Document count in index
    pub document_count: usize,

    /// Index metadata
    pub metadata: HashMap<String, serde_json::Value>,

    /// Index health status
    pub health_status: String,
}

/// Operation log entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperationLogEntry {
    /// Entry ID
    pub entry_id: String,

    /// Operation that was performed
    pub operation: IndexUpdate,

    /// Operation result
    pub result: Option<UpdateResult>,

    /// Pre-operation state hash
    pub pre_state_hash: String,

    /// Post-operation state hash
    pub post_state_hash: Option<String>,

    /// Operation timestamp
    pub timestamp: chrono::DateTime<chrono::Utc>,

    /// Operation source
    pub source: String,

    /// Rollback information for this operation
    pub rollback_info: RollbackOperationInfo,
}

/// Rollback information for an operation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackOperationInfo {
    /// Can this operation be rolled back?
    pub is_rollbackable: bool,

    /// Rollback priority (higher = more critical)
    pub rollback_priority: u8,

    /// Dependencies that must be rolled back first
    pub rollback_dependencies: Vec<String>,

    /// Custom rollback steps
    pub custom_rollback_steps: Vec<CustomRollbackStep>,

    /// Estimated rollback time
    pub estimated_rollback_time_ms: u64,
}

/// Custom rollback step
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CustomRollbackStep {
    /// Step name
    pub step_name: String,

    /// Step type
    pub step_type: RollbackStepType,

    /// Step parameters
    pub parameters: HashMap<String, serde_json::Value>,

    /// Step order
    pub order: u32,
}

/// Types of rollback steps
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RollbackStepType {
    /// Delete entries
    Delete,
    /// Restore entries
    Restore,
    /// Update entries
    Update,
    /// Rebuild index
    RebuildIndex,
    /// Custom operation
    Custom(String),
}

/// Rollback point for grouped operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackPoint {
    /// Rollback point ID
    pub rollback_point_id: String,

    /// Point creation timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,

    /// Description of the rollback point
    pub description: String,

    /// Operations included in this point
    pub operation_ids: Vec<String>,

    /// System state at this point
    pub system_state: SystemState,

    /// Point metadata
    pub metadata: HashMap<String, serde_json::Value>,

    /// Whether this point can be automatically used
    pub auto_rollback_eligible: bool,
}

/// Recovery result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecoveryResult {
    /// Recovery operation ID
    pub recovery_id: String,

    /// Whether recovery was successful
    pub success: bool,

    /// Operations that were rolled back
    pub rolled_back_operations: Vec<String>,

    /// Final system state after recovery
    pub final_state: Option<SystemState>,

    /// Recovery time in milliseconds
    pub recovery_time_ms: u64,

    /// Recovery verification results
    pub verification_results: Vec<VerificationResult>,

    /// Errors encountered during recovery
    pub errors: Vec<String>,

    /// Recovery metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Verification result for rollback operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VerificationResult {
    /// Verification check name
    pub check_name: String,

    /// Whether verification passed
    pub passed: bool,

    /// Verification details
    pub details: String,

    /// Expected vs actual values
    pub comparison: Option<HashMap<String, serde_json::Value>>,
}

/// Operation log for tracking changes
pub struct OperationLog {
    /// Log entries
    entries: VecDeque<OperationLogEntry>,

    /// Maximum log size
    max_size: usize,

    /// Total operations logged
    total_operations: u64,
}

impl OperationLog {
    /// Create new operation log
    pub fn new(max_size: usize) -> Self {
        Self {
            entries: VecDeque::new(),
            max_size,
            total_operations: 0,
        }
    }

    /// Add operation to log
    pub fn log_operation(
        &mut self,
        operation: IndexUpdate,
        result: Option<UpdateResult>,
        pre_state_hash: String,
        post_state_hash: Option<String>,
    ) {
        let entry = OperationLogEntry {
            entry_id: Uuid::new_v4().to_string(),
            operation,
            result,
            pre_state_hash,
            post_state_hash,
            timestamp: chrono::Utc::now(),
            source: "operation_log".to_string(),
            rollback_info: RollbackOperationInfo {
                is_rollbackable: true,
                rollback_priority: 5,
                rollback_dependencies: Vec::new(),
                custom_rollback_steps: Vec::new(),
                estimated_rollback_time_ms: 1000,
            },
        };

        self.entries.push_back(entry);
        self.total_operations += 1;

        // Limit log size
        while self.entries.len() > self.max_size {
            self.entries.pop_front();
        }
    }

    /// Get recent operations
    pub fn get_recent_operations(&self, count: usize) -> Vec<&OperationLogEntry> {
        self.entries.iter().rev().take(count).collect()
    }

    /// Find operations by criteria
    pub fn find_operations<F>(&self, predicate: F) -> Vec<&OperationLogEntry>
    where
        F: Fn(&OperationLogEntry) -> bool,
    {
        self.entries
            .iter()
            .filter(|entry| predicate(entry))
            .collect()
    }
}

/// Main rollback manager
pub struct RollbackManager {
    /// Configuration
    config: RollbackConfig,

    /// Operation log
    operation_log: Arc<RwLock<OperationLog>>,

    /// System snapshots
    snapshots: Arc<RwLock<VecDeque<SystemState>>>,

    /// Rollback points
    rollback_points: Arc<RwLock<HashMap<String, RollbackPoint>>>,

    /// Recovery history
    recovery_history: Arc<RwLock<VecDeque<RecoveryResult>>>,

    /// Rollback statistics
    stats: Arc<RwLock<RollbackStats>>,

    /// Background task handles
    task_handles: Arc<tokio::sync::Mutex<Vec<tokio::task::JoinHandle<()>>>>,
}

/// Rollback statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackStats {
    /// Total operations logged
    pub total_operations_logged: u64,

    /// Total rollbacks performed
    pub total_rollbacks: u64,

    /// Successful rollbacks
    pub successful_rollbacks: u64,

    /// Failed rollbacks
    pub failed_rollbacks: u64,

    /// Average rollback time
    pub avg_rollback_time_ms: f64,

    /// Total snapshots created
    pub total_snapshots: u64,

    /// Current storage usage
    pub storage_usage_bytes: u64,

    /// Last snapshot timestamp
    pub last_snapshot_at: Option<chrono::DateTime<chrono::Utc>>,

    /// Last updated
    pub last_updated: chrono::DateTime<chrono::Utc>,
}

impl RollbackManager {
    /// Create new rollback manager
    pub async fn new(config: RollbackConfig) -> RragResult<Self> {
        let manager = Self {
            config: config.clone(),
            operation_log: Arc::new(RwLock::new(OperationLog::new(
                config.max_operation_log_size,
            ))),
            snapshots: Arc::new(RwLock::new(VecDeque::new())),
            rollback_points: Arc::new(RwLock::new(HashMap::new())),
            recovery_history: Arc::new(RwLock::new(VecDeque::new())),
            stats: Arc::new(RwLock::new(RollbackStats {
                total_operations_logged: 0,
                total_rollbacks: 0,
                successful_rollbacks: 0,
                failed_rollbacks: 0,
                avg_rollback_time_ms: 0.0,
                total_snapshots: 0,
                storage_usage_bytes: 0,
                last_snapshot_at: None,
                last_updated: chrono::Utc::now(),
            })),
            task_handles: Arc::new(tokio::sync::Mutex::new(Vec::new())),
        };

        manager.start_background_tasks().await?;
        Ok(manager)
    }

    /// Log an operation for potential rollback
    pub async fn log_operation(
        &self,
        operation: IndexUpdate,
        result: Option<UpdateResult>,
        pre_state_hash: String,
        post_state_hash: Option<String>,
    ) -> RragResult<()> {
        let mut log = self.operation_log.write().await;
        log.log_operation(operation, result, pre_state_hash, post_state_hash);

        // Update statistics
        {
            let mut stats = self.stats.write().await;
            stats.total_operations_logged += 1;
            stats.last_updated = chrono::Utc::now();
        }

        // Check if snapshot is needed
        if self.config.enable_snapshots
            && log.total_operations % self.config.snapshot_interval as u64 == 0
        {
            drop(log);
            self.create_snapshot("auto_snapshot".to_string()).await?;
        }

        Ok(())
    }

    /// Create a system state snapshot
    pub async fn create_snapshot(&self, _description: String) -> RragResult<String> {
        let snapshot_id = Uuid::new_v4().to_string();

        // Collect current system state
        let snapshot = SystemState {
            snapshot_id: snapshot_id.clone(),
            created_at: chrono::Utc::now(),
            document_states: self.collect_document_states().await?,
            index_states: self.collect_index_states().await?,
            system_metadata: HashMap::new(),
            operations_count: {
                let log = self.operation_log.read().await;
                log.total_operations
            },
            size_bytes: 0, // Would be calculated
            compression_ratio: 1.0,
        };

        // Add to snapshots
        {
            let mut snapshots = self.snapshots.write().await;
            snapshots.push_back(snapshot);

            // Limit snapshot count
            while snapshots.len() > self.config.max_snapshots {
                snapshots.pop_front();
            }
        }

        // Update statistics
        {
            let mut stats = self.stats.write().await;
            stats.total_snapshots += 1;
            stats.last_snapshot_at = Some(chrono::Utc::now());
            stats.last_updated = chrono::Utc::now();
        }

        Ok(snapshot_id)
    }

    /// Create a rollback point
    pub async fn create_rollback_point(
        &self,
        description: String,
        operation_ids: Vec<String>,
        auto_eligible: bool,
    ) -> RragResult<String> {
        let rollback_point_id = Uuid::new_v4().to_string();

        // Create snapshot for rollback point
        let snapshot_id = self
            .create_snapshot(format!("rollback_point_{}", description))
            .await?;

        let snapshot = {
            let snapshots = self.snapshots.read().await;
            snapshots
                .iter()
                .find(|s| s.snapshot_id == snapshot_id)
                .unwrap()
                .clone()
        };

        let rollback_point = RollbackPoint {
            rollback_point_id: rollback_point_id.clone(),
            created_at: chrono::Utc::now(),
            description,
            operation_ids,
            system_state: snapshot,
            metadata: HashMap::new(),
            auto_rollback_eligible: auto_eligible,
        };

        {
            let mut points = self.rollback_points.write().await;
            points.insert(rollback_point_id.clone(), rollback_point);
        }

        Ok(rollback_point_id)
    }

    /// Perform rollback operation
    pub async fn rollback(&self, rollback_op: RollbackOperation) -> RragResult<RecoveryResult> {
        let start_time = std::time::Instant::now();
        let recovery_id = Uuid::new_v4().to_string();

        let mut recovery_result = RecoveryResult {
            recovery_id: recovery_id.clone(),
            success: false,
            rolled_back_operations: Vec::new(),
            final_state: None,
            recovery_time_ms: 0,
            verification_results: Vec::new(),
            errors: Vec::new(),
            metadata: HashMap::new(),
        };

        match rollback_op {
            RollbackOperation::RestoreSnapshot { snapshot_id, .. } => {
                match self.restore_from_snapshot(&snapshot_id).await {
                    Ok(operations) => {
                        recovery_result.rolled_back_operations = operations;
                        recovery_result.success = true;
                    }
                    Err(e) => {
                        recovery_result.errors.push(e.to_string());
                    }
                }
            }

            RollbackOperation::UndoOperations { operation_ids } => {
                match self.undo_operations(&operation_ids).await {
                    Ok(operations) => {
                        recovery_result.rolled_back_operations = operations;
                        recovery_result.success = true;
                    }
                    Err(e) => {
                        recovery_result.errors.push(e.to_string());
                    }
                }
            }

            RollbackOperation::RevertToTimestamp { timestamp } => {
                match self.revert_to_timestamp(timestamp).await {
                    Ok(operations) => {
                        recovery_result.rolled_back_operations = operations;
                        recovery_result.success = true;
                    }
                    Err(e) => {
                        recovery_result.errors.push(e.to_string());
                    }
                }
            }

            _ => {
                recovery_result
                    .errors
                    .push("Rollback operation not implemented".to_string());
            }
        }

        recovery_result.recovery_time_ms = start_time.elapsed().as_millis() as u64;

        // Perform verification if enabled
        if self.config.enable_verification {
            recovery_result.verification_results = self.verify_rollback(&recovery_result).await?;
        }

        // Store recovery result
        {
            let mut history = self.recovery_history.write().await;
            history.push_back(recovery_result.clone());

            // Limit history size
            if history.len() > 100 {
                history.pop_front();
            }
        }

        // Update statistics
        {
            let mut stats = self.stats.write().await;
            stats.total_rollbacks += 1;
            if recovery_result.success {
                stats.successful_rollbacks += 1;
            } else {
                stats.failed_rollbacks += 1;
            }
            stats.avg_rollback_time_ms =
                (stats.avg_rollback_time_ms + recovery_result.recovery_time_ms as f64) / 2.0;
            stats.last_updated = chrono::Utc::now();
        }

        Ok(recovery_result)
    }

    /// Get rollback statistics
    pub async fn get_stats(&self) -> RollbackStats {
        self.stats.read().await.clone()
    }

    /// Get available snapshots
    pub async fn get_snapshots(&self) -> RragResult<Vec<SystemState>> {
        let snapshots = self.snapshots.read().await;
        Ok(snapshots.iter().cloned().collect())
    }

    /// Get rollback points
    pub async fn get_rollback_points(&self) -> RragResult<Vec<RollbackPoint>> {
        let points = self.rollback_points.read().await;
        Ok(points.values().cloned().collect())
    }

    /// Health check
    pub async fn health_check(&self) -> RragResult<bool> {
        let handles = self.task_handles.lock().await;
        let all_running = handles.iter().all(|handle| !handle.is_finished());

        let stats = self.get_stats().await;
        let healthy_stats = stats.failed_rollbacks < stats.successful_rollbacks * 2; // Arbitrary threshold

        Ok(all_running && healthy_stats)
    }

    /// Start background maintenance tasks
    async fn start_background_tasks(&self) -> RragResult<()> {
        let mut handles = self.task_handles.lock().await;

        if self.config.enable_snapshots {
            handles.push(self.start_snapshot_cleanup_task().await);
        }

        Ok(())
    }

    /// Start snapshot cleanup task
    async fn start_snapshot_cleanup_task(&self) -> tokio::task::JoinHandle<()> {
        let snapshots = Arc::clone(&self.snapshots);
        let config = self.config.clone();

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(3600)); // 1 hour

            loop {
                interval.tick().await;

                let mut snapshots_guard = snapshots.write().await;

                // Remove old snapshots beyond retention limit
                while snapshots_guard.len() > config.max_snapshots {
                    snapshots_guard.pop_front();
                }

                // Could also implement time-based cleanup here
            }
        })
    }

    /// Collect current document states (placeholder)
    async fn collect_document_states(&self) -> RragResult<HashMap<String, DocumentState>> {
        // In production, this would collect actual document states from storage
        Ok(HashMap::new())
    }

    /// Collect current index states (placeholder)
    async fn collect_index_states(&self) -> RragResult<HashMap<String, IndexState>> {
        // In production, this would collect actual index states
        Ok(HashMap::new())
    }

    /// Restore system from snapshot (placeholder)
    async fn restore_from_snapshot(&self, _snapshot_id: &str) -> RragResult<Vec<String>> {
        // In production, this would restore system state from snapshot
        Ok(Vec::new())
    }

    /// Undo specific operations (placeholder)
    async fn undo_operations(&self, operation_ids: &[String]) -> RragResult<Vec<String>> {
        // In production, this would undo the specified operations
        Ok(operation_ids.to_vec())
    }

    /// Revert to timestamp (placeholder)
    async fn revert_to_timestamp(
        &self,
        _timestamp: chrono::DateTime<chrono::Utc>,
    ) -> RragResult<Vec<String>> {
        // In production, this would revert to the specified timestamp
        Ok(Vec::new())
    }

    /// Verify rollback results (placeholder)
    async fn verify_rollback(
        &self,
        _result: &RecoveryResult,
    ) -> RragResult<Vec<VerificationResult>> {
        // In production, this would verify the rollback was successful
        Ok(vec![VerificationResult {
            check_name: "system_integrity".to_string(),
            passed: true,
            details: "System integrity verified".to_string(),
            comparison: None,
        }])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::incremental::index_manager::IndexOperation;
    use crate::Document;

    #[tokio::test]
    async fn test_rollback_manager_creation() {
        let config = RollbackConfig::default();
        let manager = RollbackManager::new(config).await.unwrap();
        assert!(manager.health_check().await.unwrap());
    }

    #[tokio::test]
    async fn test_operation_logging() {
        let manager = RollbackManager::new(RollbackConfig::default())
            .await
            .unwrap();

        let doc = Document::new("Test content");
        let operation = IndexOperation::Add {
            document: doc,
            chunks: Vec::new(),
            embeddings: Vec::new(),
        };

        let update = IndexUpdate {
            operation_id: Uuid::new_v4().to_string(),
            operation,
            priority: 5,
            timestamp: chrono::Utc::now(),
            source: "test".to_string(),
            metadata: HashMap::new(),
            dependencies: Vec::new(),
            max_retries: 3,
            retry_count: 0,
        };

        manager
            .log_operation(
                update,
                None,
                "pre_hash".to_string(),
                Some("post_hash".to_string()),
            )
            .await
            .unwrap();

        let stats = manager.get_stats().await;
        assert_eq!(stats.total_operations_logged, 1);
    }

    #[tokio::test]
    async fn test_snapshot_creation() {
        let manager = RollbackManager::new(RollbackConfig::default())
            .await
            .unwrap();

        let snapshot_id = manager
            .create_snapshot("test_snapshot".to_string())
            .await
            .unwrap();
        assert!(!snapshot_id.is_empty());

        let snapshots = manager.get_snapshots().await.unwrap();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].snapshot_id, snapshot_id);

        let stats = manager.get_stats().await;
        assert_eq!(stats.total_snapshots, 1);
    }

    #[tokio::test]
    async fn test_rollback_point_creation() {
        let manager = RollbackManager::new(RollbackConfig::default())
            .await
            .unwrap();

        let point_id = manager
            .create_rollback_point(
                "test_point".to_string(),
                vec!["op1".to_string(), "op2".to_string()],
                true,
            )
            .await
            .unwrap();

        assert!(!point_id.is_empty());

        let points = manager.get_rollback_points().await.unwrap();
        assert_eq!(points.len(), 1);
        assert_eq!(points[0].rollback_point_id, point_id);
        assert_eq!(points[0].operation_ids.len(), 2);
    }

    #[tokio::test]
    async fn test_rollback_operation() {
        let manager = RollbackManager::new(RollbackConfig::default())
            .await
            .unwrap();

        // Create a snapshot first
        let snapshot_id = manager
            .create_snapshot("test_snapshot".to_string())
            .await
            .unwrap();

        // Perform rollback
        let rollback_op = RollbackOperation::RestoreSnapshot {
            snapshot_id,
            target_state: SystemState {
                snapshot_id: "dummy".to_string(),
                created_at: chrono::Utc::now(),
                document_states: HashMap::new(),
                index_states: HashMap::new(),
                system_metadata: HashMap::new(),
                operations_count: 0,
                size_bytes: 0,
                compression_ratio: 1.0,
            },
        };

        let result = manager.rollback(rollback_op).await.unwrap();
        assert!(result.success);
        assert!(result.recovery_time_ms > 0);

        let stats = manager.get_stats().await;
        assert_eq!(stats.total_rollbacks, 1);
        assert_eq!(stats.successful_rollbacks, 1);
    }

    #[test]
    fn test_rollback_strategies() {
        let strategies = vec![
            RollbackStrategy::LastKnownGood,
            RollbackStrategy::SpecificSnapshot,
            RollbackStrategy::Selective,
            RollbackStrategy::Complete,
            RollbackStrategy::Custom("custom".to_string()),
        ];

        // Ensure all strategies are different
        for (i, strategy1) in strategies.iter().enumerate() {
            for (j, strategy2) in strategies.iter().enumerate() {
                if i != j {
                    assert_ne!(format!("{:?}", strategy1), format!("{:?}", strategy2));
                }
            }
        }
    }

    #[test]
    fn test_rollback_step_types() {
        let step_types = vec![
            RollbackStepType::Delete,
            RollbackStepType::Restore,
            RollbackStepType::Update,
            RollbackStepType::RebuildIndex,
            RollbackStepType::Custom("custom".to_string()),
        ];

        // Ensure all step types are different
        for (i, type1) in step_types.iter().enumerate() {
            for (j, type2) in step_types.iter().enumerate() {
                if i != j {
                    assert_ne!(format!("{:?}", type1), format!("{:?}", type2));
                }
            }
        }
    }
}