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,
})
}
pub async fn create_checkpoint(&self, description: String) -> crate::Result<String> {
let checkpoint_id = format!("{:x}", rand::thread_rng().gen::<u128>());
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);
while checkpoints.len() > self.max_checkpoints {
checkpoints.pop_back();
}
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)
}
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(())
}
pub async fn rollback_last(&self) -> crate::Result<()> {
let checkpoint = {
let mut checkpoints = self.checkpoints.write().await;
let cp = checkpoints.pop_front();
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(())
}
async fn rollback_to_checkpoint(&self, checkpoint: &Checkpoint) -> crate::Result<()> {
info!("🔄 Rolling back {} operations...", checkpoint.operations.len());
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");
}
}
}
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 => {
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 => {
info!("⚠️ Cannot rollback command: {}", operation.target);
}
}
}
Ok(())
}
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()
}
}