claude_utils/turbo/

rollback.rs

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::path::PathBuf;
use tokio::fs;
use tokio::process::Command;
use tracing::info;
use rand::Rng;

use super::checkpoint_store::CheckpointStore;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Checkpoint {
    pub id: String,
    pub timestamp: DateTime<Utc>,
    pub description: String,
    pub git_stash: Option<String>,
    pub operations: Vec<Operation>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Operation {
    pub op_type: OperationType,
    pub target: String,
    pub backup_path: Option<PathBuf>,
    pub original_content: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OperationType {
    FileEdit,
    FileCreate,
    FileDelete,
    CommandExecute,
}

pub struct RollbackManager {
    checkpoints: tokio::sync::RwLock<VecDeque<Checkpoint>>,
    max_checkpoints: usize,
    store: CheckpointStore,
}

impl RollbackManager {
    pub async fn new() -> crate::Result<Self> {
        let store = CheckpointStore::new();
        let checkpoints = store.load_checkpoints().await.unwrap_or_else(|e| {
            info!("Failed to load checkpoints: {}, starting fresh", e);
            VecDeque::new()
        });
        
        Ok(Self {
            checkpoints: tokio::sync::RwLock::new(checkpoints),
            max_checkpoints: 10,
            store,
        })
    }

    /// Create a new checkpoint
    pub async fn create_checkpoint(&self, description: String) -> crate::Result<String> {
        let checkpoint_id = format!("{:x}", rand::thread_rng().gen::<u128>());
        
        // Try to create git stash
        let git_stash = match Command::new("git")
            .args(["stash", "push", "-m", &format!("turbo-checkpoint-{checkpoint_id}")])
            .output()
            .await
        {
            Ok(output) if output.status.success() => {
                let stash_ref = String::from_utf8_lossy(&output.stdout);
                Some(stash_ref.trim().to_string())
            }
            _ => {
                info!("Git stash failed, using file backups only");
                None
            }
        };

        let checkpoint = Checkpoint {
            id: checkpoint_id.clone(),
            timestamp: Utc::now(),
            description,
            git_stash,
            operations: Vec::new(),
        };

        let mut checkpoints = self.checkpoints.write().await;
        checkpoints.push_front(checkpoint);
        
        // Keep only max_checkpoints
        while checkpoints.len() > self.max_checkpoints {
            checkpoints.pop_back();
        }
        
        // Save to disk
        if let Err(e) = self.store.save_checkpoints(&checkpoints).await {
            info!("Failed to save checkpoints to disk: {}", e);
        }

        info!("✅ Created checkpoint: {}", checkpoint_id);
        Ok(checkpoint_id)
    }

    /// Record an operation for rollback
    pub async fn record_operation(&self, operation: Operation) -> crate::Result<()> {
        let mut checkpoints = self.checkpoints.write().await;
        if let Some(current) = checkpoints.front_mut() {
            current.operations.push(operation);
        }
        Ok(())
    }

    /// Rollback to the last checkpoint
    pub async fn rollback_last(&self) -> crate::Result<()> {
        let checkpoint = {
            let mut checkpoints = self.checkpoints.write().await;
            let cp = checkpoints.pop_front();
            
            // Save updated list to disk
            if cp.is_some() {
                if let Err(e) = self.store.save_checkpoints(&checkpoints).await {
                    info!("Failed to save checkpoints after rollback: {}", e);
                }
            }
            
            cp
        };

        if let Some(checkpoint) = checkpoint {
            self.rollback_to_checkpoint(&checkpoint).await?;
            info!("✅ Rolled back to checkpoint: {}", checkpoint.id);
        } else {
            return Err(crate::ClaudeUtilsError::Turbo("No checkpoints available".into()));
        }

        Ok(())
    }

    /// Rollback to a specific checkpoint
    async fn rollback_to_checkpoint(&self, checkpoint: &Checkpoint) -> crate::Result<()> {
        info!("🔄 Rolling back {} operations...", checkpoint.operations.len());

        // First try git stash pop if available
        if let Some(ref stash_ref) = checkpoint.git_stash {
            match Command::new("git")
                .args(["stash", "pop", stash_ref])
                .output()
                .await
            {
                Ok(output) if output.status.success() => {
                    info!("✅ Git stash restored successfully");
                    return Ok(());
                }
                _ => {
                    info!("Git stash restore failed, falling back to manual rollback");
                }
            }
        }

        // Manual rollback of operations
        for operation in checkpoint.operations.iter().rev() {
            match operation.op_type {
                OperationType::FileEdit => {
                    if let Some(ref content) = operation.original_content {
                        fs::write(&operation.target, content).await?;
                        info!("Restored: {}", operation.target);
                    }
                }
                OperationType::FileCreate => {
                    // Delete the created file
                    let _ = fs::remove_file(&operation.target).await;
                    info!("Removed: {}", operation.target);
                }
                OperationType::FileDelete => {
                    if let Some(ref backup_path) = operation.backup_path {
                        fs::copy(backup_path, &operation.target).await?;
                        info!("Restored: {}", operation.target);
                    }
                }
                OperationType::CommandExecute => {
                    // Commands can't be rolled back directly
                    info!("⚠️  Cannot rollback command: {}", operation.target);
                }
            }
        }

        Ok(())
    }

    /// Get list of available checkpoints
    pub async fn list_checkpoints(&self) -> Vec<(String, String, DateTime<Utc>)> {
        self.checkpoints
            .read()
            .await
            .iter()
            .map(|cp| (cp.id.clone(), cp.description.clone(), cp.timestamp))
            .collect()
    }
}