use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use crate::Result;
const MAX_SNAPSHOTS: usize = 20;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotEntry {
pub relative_path: String,
pub had_content: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotManifest {
pub id: String,
pub label: String,
pub root: String,
pub entries: Vec<SnapshotEntry>,
pub created_at: u64,
}
#[derive(Debug, Serialize)]
pub struct RestoreResult {
pub snapshot_id: String,
pub label: String,
pub files_restored: usize,
pub files_removed: usize,
pub errors: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct SnapshotSummary {
pub id: String,
pub label: String,
pub root: String,
pub file_count: usize,
pub created_at: u64,
pub age: String,
}
#[derive(Debug, Clone, Default)]
pub struct InMemoryRollback {
entries: Vec<RollbackEntry>,
}
#[derive(Debug, Clone)]
struct RollbackEntry {
path: PathBuf,
original: Option<Vec<u8>>,
}
impl InMemoryRollback {
pub fn new() -> Self {
Self {
entries: Vec::new(),
}
}
pub fn capture(&mut self, path: &Path) {
if self.entries.iter().any(|e| e.path == path) {
return;
}
let original = if path.is_file() {
std::fs::read(path).ok()
} else {
None
};
self.entries.push(RollbackEntry {
path: path.to_path_buf(),
original,
});
}
pub fn restore_all(&self) {
for entry in &self.entries {
match &entry.original {
Some(content) => {
let _ = std::fs::write(&entry.path, content);
}
None => {
let _ = std::fs::remove_file(&entry.path);
}
}
}
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
pub struct UndoSnapshot {
root: PathBuf,
label: String,
entries: Vec<SnapshotEntry>,
contents: Vec<(String, Vec<u8>)>,
}
impl UndoSnapshot {
pub fn new(root: &Path, label: &str) -> Self {
Self {
root: root.to_path_buf(),
label: label.to_string(),
entries: Vec::new(),
contents: Vec::new(),
}
}
pub fn capture_file(&mut self, relative_path: &str) {
if self
.entries
.iter()
.any(|e| e.relative_path == relative_path)
{
return;
}
let abs = self.root.join(relative_path);
let had_content = abs.is_file();
if had_content {
if let Ok(content) = std::fs::read(&abs) {
self.contents.push((relative_path.to_string(), content));
}
}
self.entries.push(SnapshotEntry {
relative_path: relative_path.to_string(),
had_content,
});
}
pub fn save(self) -> Result<String> {
save_to_dir(self, &snapshots_dir())
}
}
fn save_to_dir(snap: UndoSnapshot, base_dir: &Path) -> Result<String> {
if snap.entries.is_empty() {
return Err(crate::Error::validation_invalid_argument(
"undo",
"No files to snapshot",
None,
None,
));
}
let id = generate_snapshot_id();
let snapshot_dir = base_dir.join(&id);
let files_dir = snapshot_dir.join("files");
std::fs::create_dir_all(&files_dir).map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to create snapshot dir: {}", e))
})?;
for (relative_path, content) in &snap.contents {
let safe_name = sanitize_path(relative_path);
let dest = files_dir.join(&safe_name);
if let Some(parent) = dest.parent() {
std::fs::create_dir_all(parent).ok();
}
std::fs::write(&dest, content).map_err(|e| {
crate::Error::internal_unexpected(format!(
"Failed to write snapshot file {}: {}",
relative_path, e
))
})?;
}
let manifest = SnapshotManifest {
id: id.clone(),
label: snap.label,
root: snap.root.to_string_lossy().to_string(),
entries: snap.entries,
created_at: now_unix(),
};
let manifest_json = serde_json::to_string_pretty(&manifest).map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to serialize manifest: {}", e))
})?;
std::fs::write(snapshot_dir.join("manifest.json"), manifest_json).map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to write manifest: {}", e))
})?;
log_status!(
"undo",
"Snapshot saved: {} ({} file(s))",
id,
manifest.entries.len()
);
expire_old_snapshots_in(base_dir, MAX_SNAPSHOTS);
Ok(id)
}
pub fn restore(snapshot_id: Option<&str>) -> Result<RestoreResult> {
restore_from_dir(snapshot_id, &snapshots_dir())
}
fn restore_from_dir(snapshot_id: Option<&str>, base_dir: &Path) -> Result<RestoreResult> {
let id = match snapshot_id {
Some(id) => id.to_string(),
None => latest_snapshot_id_in(base_dir)?,
};
let snapshot_dir = base_dir.join(&id);
let manifest = load_manifest(&snapshot_dir)?;
let files_dir = snapshot_dir.join("files");
let root = Path::new(&manifest.root);
let mut files_restored = 0;
let mut files_removed = 0;
let mut errors = Vec::new();
for entry in &manifest.entries {
let target = root.join(&entry.relative_path);
if entry.had_content {
let safe_name = sanitize_path(&entry.relative_path);
let source = files_dir.join(&safe_name);
match std::fs::read(&source) {
Ok(content) => {
if let Err(e) = std::fs::write(&target, content) {
errors.push(format!("Failed to restore {}: {}", entry.relative_path, e));
} else {
files_restored += 1;
}
}
Err(e) => {
errors.push(format!("Missing backup for {}: {}", entry.relative_path, e));
}
}
} else {
if target.exists() {
if let Err(e) = std::fs::remove_file(&target) {
errors.push(format!("Failed to remove {}: {}", entry.relative_path, e));
} else {
files_removed += 1;
}
}
if let Some(parent) = target.parent() {
remove_empty_parents(parent, root);
}
}
}
log_status!(
"undo",
"Restored {} file(s), removed {} created file(s)",
files_restored,
files_removed
);
if errors.is_empty() {
let _ = std::fs::remove_dir_all(&snapshot_dir);
log_status!("undo", "Snapshot {} consumed", id);
}
Ok(RestoreResult {
snapshot_id: id,
label: manifest.label,
files_restored,
files_removed,
errors,
})
}
pub fn list_snapshots() -> Result<Vec<SnapshotSummary>> {
list_snapshots_in(&snapshots_dir())
}
fn list_snapshots_in(base_dir: &Path) -> Result<Vec<SnapshotSummary>> {
if !base_dir.exists() {
return Ok(vec![]);
}
let mut summaries = Vec::new();
let now = now_unix();
let mut entries: Vec<_> = std::fs::read_dir(base_dir)
.map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to read snapshots dir: {}", e))
})?
.filter_map(|e| e.ok())
.filter(|e| e.path().is_dir())
.collect();
entries.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
for entry in entries {
let manifest_path = entry.path().join("manifest.json");
if let Ok(content) = std::fs::read_to_string(&manifest_path) {
if let Ok(manifest) = serde_json::from_str::<SnapshotManifest>(&content) {
summaries.push(SnapshotSummary {
id: manifest.id,
label: manifest.label,
root: manifest.root,
file_count: manifest.entries.len(),
age: format_age(now.saturating_sub(manifest.created_at)),
created_at: manifest.created_at,
});
}
}
}
Ok(summaries)
}
pub fn delete_snapshot(snapshot_id: &str) -> Result<()> {
delete_snapshot_in(snapshot_id, &snapshots_dir())
}
fn delete_snapshot_in(snapshot_id: &str, base_dir: &Path) -> Result<()> {
let snapshot_dir = base_dir.join(snapshot_id);
if !snapshot_dir.exists() {
return Err(crate::Error::validation_invalid_argument(
"snapshot_id",
format!("Snapshot '{}' not found", snapshot_id),
None,
None,
));
}
std::fs::remove_dir_all(&snapshot_dir).map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to delete snapshot: {}", e))
})?;
log_status!("undo", "Deleted snapshot {}", snapshot_id);
Ok(())
}
fn snapshots_dir() -> PathBuf {
if let Ok(dir) = std::env::var("HOMEBOY_SNAPSHOTS_DIR") {
return PathBuf::from(dir);
}
let cache_base = std::env::var("HOME")
.map(|h| PathBuf::from(h).join(".cache"))
.unwrap_or_else(|_| PathBuf::from("/tmp"));
cache_base.join("homeboy").join("snapshots")
}
fn generate_snapshot_id() -> String {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
format!("{}", now.as_millis())
}
fn now_unix() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn latest_snapshot_id_in(base_dir: &Path) -> Result<String> {
if !base_dir.exists() {
return Err(crate::Error::internal_unexpected(
"No undo snapshots available. Run a --write operation first.",
));
}
let mut entries: Vec<_> = std::fs::read_dir(base_dir)
.map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to read snapshots dir: {}", e))
})?
.filter_map(|e| e.ok())
.filter(|e| e.path().is_dir())
.collect();
entries.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
entries
.first()
.map(|e| e.file_name().to_string_lossy().to_string())
.ok_or_else(|| {
crate::Error::internal_unexpected(
"No undo snapshots available. Run a --write operation first.",
)
})
}
fn load_manifest(snapshot_dir: &Path) -> Result<SnapshotManifest> {
let manifest_path = snapshot_dir.join("manifest.json");
let content = std::fs::read_to_string(&manifest_path).map_err(|e| {
crate::Error::internal_unexpected(format!("Failed to read manifest: {}", e))
})?;
serde_json::from_str(&content)
.map_err(|e| crate::Error::internal_unexpected(format!("Failed to parse manifest: {}", e)))
}
fn sanitize_path(relative_path: &str) -> String {
relative_path.replace('/', "__")
}
fn expire_old_snapshots_in(base_dir: &Path, keep: usize) {
if !base_dir.exists() {
return;
}
let mut entries: Vec<_> = match std::fs::read_dir(base_dir) {
Ok(entries) => entries.filter_map(|e| e.ok()).collect(),
Err(_) => return,
};
if entries.len() <= keep {
return;
}
entries.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
for entry in entries.into_iter().skip(keep) {
let _ = std::fs::remove_dir_all(entry.path());
}
}
fn format_age(seconds: u64) -> String {
if seconds < 60 {
format!("{}s ago", seconds)
} else if seconds < 3600 {
format!("{}m ago", seconds / 60)
} else if seconds < 86400 {
format!("{}h ago", seconds / 3600)
} else {
format!("{}d ago", seconds / 86400)
}
}
fn remove_empty_parents(dir: &Path, root: &Path) {
let mut current = dir;
while current != root {
if current.is_dir() {
match std::fs::read_dir(current) {
Ok(mut entries) => {
if entries.next().is_none() {
let _ = std::fs::remove_dir(current);
} else {
break; }
}
Err(_) => break,
}
} else {
break;
}
match current.parent() {
Some(parent) => current = parent,
None => break,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use std::time::{SystemTime, UNIX_EPOCH};
fn test_root(name: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
std::env::temp_dir().join(format!("homeboy-undo-{}-{}", name, nanos))
}
fn test_snap_dir(name: &str) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
std::env::temp_dir().join(format!("homeboy-snapdir-{}-{}", name, nanos))
}
fn save_isolated(snap: UndoSnapshot, snap_dir: &Path) -> crate::Result<String> {
save_to_dir(snap, snap_dir)
}
#[test]
fn snapshot_captures_existing_file_and_restores() {
let root = test_root("capture-restore");
let snap_dir = test_snap_dir("capture-restore");
fs::create_dir_all(root.join("src")).unwrap();
fs::write(root.join("src/main.rs"), "fn main() {}\n").unwrap();
let mut snap = UndoSnapshot::new(&root, "test fix");
snap.capture_file("src/main.rs");
let id = save_isolated(snap, &snap_dir).unwrap();
fs::write(root.join("src/main.rs"), "fn main() { changed }\n").unwrap();
assert!(fs::read_to_string(root.join("src/main.rs"))
.unwrap()
.contains("changed"));
let result = restore_from_dir(Some(&id), &snap_dir).unwrap();
assert_eq!(result.files_restored, 1);
assert_eq!(result.files_removed, 0);
assert!(result.errors.is_empty());
assert_eq!(
fs::read_to_string(root.join("src/main.rs")).unwrap(),
"fn main() {}\n"
);
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn snapshot_removes_created_files_on_undo() {
let root = test_root("remove-created");
let snap_dir = test_snap_dir("remove-created");
fs::create_dir_all(root.join("src")).unwrap();
let mut snap = UndoSnapshot::new(&root, "test scaffold");
snap.capture_file("tests/new_test.rs");
let id = save_isolated(snap, &snap_dir).unwrap();
fs::create_dir_all(root.join("tests")).unwrap();
fs::write(root.join("tests/new_test.rs"), "#[test]\nfn test_it() {}\n").unwrap();
assert!(root.join("tests/new_test.rs").exists());
let result = restore_from_dir(Some(&id), &snap_dir).unwrap();
assert_eq!(result.files_restored, 0);
assert_eq!(result.files_removed, 1);
assert!(!root.join("tests/new_test.rs").exists());
assert!(!root.join("tests").exists());
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn snapshot_handles_mixed_existing_and_new_files() {
let root = test_root("mixed");
let snap_dir = test_snap_dir("mixed");
fs::create_dir_all(root.join("src")).unwrap();
fs::write(root.join("src/lib.rs"), "pub mod foo;\n").unwrap();
let mut snap = UndoSnapshot::new(&root, "audit fix");
snap.capture_file("src/lib.rs");
snap.capture_file("src/foo.rs");
let id = save_isolated(snap, &snap_dir).unwrap();
fs::write(root.join("src/lib.rs"), "pub mod foo;\npub mod bar;\n").unwrap();
fs::write(root.join("src/foo.rs"), "pub fn foo() {}\n").unwrap();
let result = restore_from_dir(Some(&id), &snap_dir).unwrap();
assert_eq!(result.files_restored, 1);
assert_eq!(result.files_removed, 1);
assert_eq!(
fs::read_to_string(root.join("src/lib.rs")).unwrap(),
"pub mod foo;\n"
);
assert!(!root.join("src/foo.rs").exists());
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn deduplicates_same_file_captured_twice() {
let root = test_root("dedup");
let snap_dir = test_snap_dir("dedup");
fs::create_dir_all(&root).unwrap();
fs::write(root.join("file.rs"), "original\n").unwrap();
let mut snap = UndoSnapshot::new(&root, "test");
snap.capture_file("file.rs");
snap.capture_file("file.rs"); let id = save_isolated(snap, &snap_dir).unwrap();
let snapshots = list_snapshots_in(&snap_dir).unwrap();
let found = snapshots.iter().find(|s| s.id == id).unwrap();
assert_eq!(found.file_count, 1);
delete_snapshot_in(&id, &snap_dir).unwrap();
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn list_snapshots_returns_newest_first() {
let root = test_root("list-order");
let snap_dir = test_snap_dir("list-order");
fs::create_dir_all(&root).unwrap();
fs::write(root.join("a.rs"), "a\n").unwrap();
let mut snap1 = UndoSnapshot::new(&root, "first");
snap1.capture_file("a.rs");
let id1 = save_isolated(snap1, &snap_dir).unwrap();
std::thread::sleep(std::time::Duration::from_millis(5));
let mut snap2 = UndoSnapshot::new(&root, "second");
snap2.capture_file("a.rs");
let id2 = save_isolated(snap2, &snap_dir).unwrap();
let snapshots = list_snapshots_in(&snap_dir).unwrap();
let ids: Vec<&str> = snapshots.iter().map(|s| s.id.as_str()).collect();
let pos1 = ids.iter().position(|id| *id == id1).unwrap();
let pos2 = ids.iter().position(|id| *id == id2).unwrap();
assert!(pos2 < pos1, "newest (id2) should come before oldest (id1)");
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn restore_latest_picks_most_recent() {
let root = test_root("restore-latest");
let snap_dir = test_snap_dir("restore-latest");
fs::create_dir_all(&root).unwrap();
fs::write(root.join("a.rs"), "original\n").unwrap();
let mut snap1 = UndoSnapshot::new(&root, "old");
snap1.capture_file("a.rs");
let _id1 = save_isolated(snap1, &snap_dir).unwrap();
std::thread::sleep(std::time::Duration::from_millis(5));
fs::write(root.join("a.rs"), "after-first\n").unwrap();
let mut snap2 = UndoSnapshot::new(&root, "new");
snap2.capture_file("a.rs");
let _id2 = save_isolated(snap2, &snap_dir).unwrap();
fs::write(root.join("a.rs"), "after-second\n").unwrap();
let result = restore_from_dir(None, &snap_dir).unwrap();
assert_eq!(result.label, "new");
assert_eq!(
fs::read_to_string(root.join("a.rs")).unwrap(),
"after-first\n"
);
let result = restore_from_dir(None, &snap_dir).unwrap();
assert_eq!(result.label, "old");
assert_eq!(fs::read_to_string(root.join("a.rs")).unwrap(), "original\n");
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn empty_snapshot_returns_error() {
let root = test_root("empty");
let snap_dir = test_snap_dir("empty");
fs::create_dir_all(&root).unwrap();
let snap = UndoSnapshot::new(&root, "empty");
assert!(save_isolated(snap, &snap_dir).is_err());
let _ = fs::remove_dir_all(root);
let _ = fs::remove_dir_all(snap_dir);
}
#[test]
fn sanitize_path_flattens_slashes() {
assert_eq!(sanitize_path("src/core/fixer.rs"), "src__core__fixer.rs");
assert_eq!(sanitize_path("simple.rs"), "simple.rs");
}
#[test]
fn format_age_outputs_readable_strings() {
assert_eq!(format_age(30), "30s ago");
assert_eq!(format_age(90), "1m ago");
assert_eq!(format_age(7200), "2h ago");
assert_eq!(format_age(172800), "2d ago");
}
#[test]
fn in_memory_rollback_restores_modified_file() {
let root = test_root("imr-restore");
fs::create_dir_all(&root).unwrap();
let file = root.join("a.rs");
fs::write(&file, "original\n").unwrap();
let mut rollback = InMemoryRollback::new();
rollback.capture(&file);
assert_eq!(rollback.len(), 1);
fs::write(&file, "modified\n").unwrap();
assert!(fs::read_to_string(&file).unwrap().contains("modified"));
rollback.restore_all();
assert_eq!(fs::read_to_string(&file).unwrap(), "original\n");
let _ = fs::remove_dir_all(root);
}
#[test]
fn in_memory_rollback_removes_created_file() {
let root = test_root("imr-remove");
fs::create_dir_all(&root).unwrap();
let file = root.join("new.rs");
let mut rollback = InMemoryRollback::new();
rollback.capture(&file);
fs::write(&file, "created\n").unwrap();
assert!(file.exists());
rollback.restore_all();
assert!(!file.exists());
let _ = fs::remove_dir_all(root);
}
#[test]
fn in_memory_rollback_deduplicates() {
let root = test_root("imr-dedup");
fs::create_dir_all(&root).unwrap();
let file = root.join("dup.rs");
fs::write(&file, "content\n").unwrap();
let mut rollback = InMemoryRollback::new();
rollback.capture(&file);
rollback.capture(&file); assert_eq!(rollback.len(), 1);
let _ = fs::remove_dir_all(root);
}
#[test]
fn in_memory_rollback_handles_mixed_files() {
let root = test_root("imr-mixed");
fs::create_dir_all(&root).unwrap();
let existing = root.join("existing.rs");
let new_file = root.join("new.rs");
fs::write(&existing, "before\n").unwrap();
let mut rollback = InMemoryRollback::new();
rollback.capture(&existing);
rollback.capture(&new_file);
assert_eq!(rollback.len(), 2);
fs::write(&existing, "after\n").unwrap();
fs::write(&new_file, "created\n").unwrap();
rollback.restore_all();
assert_eq!(fs::read_to_string(&existing).unwrap(), "before\n");
assert!(!new_file.exists());
let _ = fs::remove_dir_all(root);
}
}