agent-manager/src/backup.rs

//! Pre-update backups (issue #64, axe 8): a snapshot of the install
//! directory, the state database and the user config is taken before every
//! real update, kept under <state>/backups/<timestamp>/, and restored by
//! 'am update --rollback <point>'. Retention is bounded by
//! settings.backups_keep (default 5).

use crate::app::App;
use anyhow::{anyhow, bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

/// Default number of backups to keep (settings.backups_keep default).
pub const DEFAULT_KEEP: u32 = 5;

/// Manifest of one backup point (JSON file inside the backup directory).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BackupInfo {
    /// Backup id = the directory name (RFC 3339 timestamp).
    pub id: String,
    /// RFC 3339 creation timestamp.
    pub created_at: String,
    /// Reason: "update:<agent>" or "update:--all".
    pub reason: String,
    /// Agents whose install directories were snapshotted (empty = all).
    pub agents: Vec<String>,
    /// True when the state database was snapshotted.
    pub state: bool,
    /// True when the user config file was snapshotted.
    pub config: bool,
}

impl BackupInfo {
    /// Load the manifest of one backup directory.
    fn load(dir: &Path) -> Option<BackupInfo> {
        let text = std::fs::read_to_string(dir.join("manifest.json")).ok()?;
        serde_json::from_str(&text).ok()
    }
}

/// Directory holding every backup point. Lives next to the state database
/// so tests and --config overrides stay fully sandboxed.
pub fn backups_dir(app: &App) -> PathBuf {
    app.paths
        .state_file
        .parent()
        .unwrap_or_else(|| std::path::Path::new("."))
        .join("backups")
}

/// Take a snapshot before an update. `agents` lists the agents being
/// updated; an empty list means "the whole install directory".
pub fn create_backup(app: &App, reason: &str, agents: &[String]) -> Result<BackupInfo> {
    let now = crate::installers::now_rfc3339();
    // Directory names must be sortable and filesystem-safe: the RFC 3339
    // timestamp with colons replaced by dashes.
    let id = now.replace(':', "-").replace('Z', "").replace('+', "_");
    let root = backups_dir(app).join(&id);
    std::fs::create_dir_all(&root)
        .with_context(|| format!("cannot create backup directory {}", root.display()))?;

    let mut info = BackupInfo {
        id: id.clone(),
        created_at: now,
        reason: reason.to_string(),
        agents: agents.to_vec(),
        state: false,
        config: false,
    };

    // 1. State database.
    if let Ok(sf) = app.state.load() {
        let text = serde_json::to_string_pretty(&sf)?;
        std::fs::write(root.join("state.json"), text)
            .with_context(|| "cannot snapshot state.json".to_string())?;
        info.state = true;
    }

    // 2. User config file (the active origin when it is a user file).
    if let Some(origin) = &app.config_origin {
        if origin.is_file() {
            std::fs::copy(origin, root.join("config.yaml")).with_context(|| {
                format!("cannot snapshot user config {}", origin.display())
            })?;
            info.config = true;
        }
    }

    // 3. Install directories of the updated agents (or everything).
    let install_root = &app.paths.install_dir;
    let dest_root = root.join("install");
    std::fs::create_dir_all(&dest_root)?;
    if agents.is_empty() {
        copy_tree(install_root, &dest_root)
            .with_context(|| format!("cannot snapshot {}", install_root.display()))?;
    } else {
        for name in agents {
            let entry = app.state.get(name).ok().flatten();
            if let Some(e) = entry {
                let src = PathBuf::from(&e.install_dir);
                if src.is_dir() {
                    let dest = dest_root.join(&e.name);
                    copy_tree(&src, &dest).with_context(|| {
                        format!("cannot snapshot {} for {name}", src.display())
                    })?;
                }
            }
            // Also snapshot the shim binaries (bin/<name>).
            let shim = app.paths.bin_dir.join(name);
            if shim.exists() {
                copy_tree(&shim, &dest_root.join("bin").join(name))?;
            }
        }
    }

    // Write the manifest last: a backup without a manifest is invalid.
    let manifest = serde_json::to_string_pretty(&info)?;
    std::fs::write(root.join("manifest.json"), manifest)
        .with_context(|| "cannot write backup manifest".to_string())?;

    prune(app)?;
    Ok(info)
}

/// List every backup point, oldest first.
pub fn list_backups(app: &App) -> Vec<BackupInfo> {
    let mut out: Vec<BackupInfo> = Vec::new();
    if let Ok(entries) = std::fs::read_dir(backups_dir(app)) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                if let Some(info) = BackupInfo::load(&path) {
                    out.push(info);
                }
            }
        }
    }
    out.sort_by(|a, b| a.id.cmp(&b.id));
    out
}

/// Resolve a rollback point: "latest"/empty = most recent backup,
/// otherwise the backup id (prefix matching allowed).
pub fn resolve_point(app: &App, point: Option<&str>) -> Result<BackupInfo> {
    let all = list_backups(app);
    if all.is_empty() {
        bail!("no backup found — nothing to roll back to");
    }
    match point {
        None | Some("latest") => all
            .last()
            .cloned()
            .ok_or_else(|| anyhow!("no backup found")),
        Some(p) => {
            let exact = all.iter().find(|b| b.id == p);
            if let Some(b) = exact {
                return Ok(b.clone());
            }
            let prefix: Vec<BackupInfo> = all
                .iter()
                .filter(|b| b.id.starts_with(p))
                .cloned()
                .collect();
            match prefix.len() {
                0 => bail!("backup '{p}' not found — run 'am update --rollback --list'"),
                1 => Ok(prefix[0].clone()),
                _ => bail!(
                    "backup prefix '{p}' is ambiguous ({} matches: {})",
                    prefix.len(),
                    prefix
                        .iter()
                        .map(|b| b.id.as_str())
                        .collect::<Vec<_>>()
                        .join(", ")
                ),
            }
        }
    }
}

/// Restore a backup point over the live install: state, user config and
/// install directories. The agent processes of the restored agents are
/// stopped first (they would hold the old binaries open).
pub fn restore_backup(app: &App, point: Option<&str>) -> Result<BackupInfo> {
    let info = resolve_point(app, point)?;
    let root = backups_dir(app).join(&info.id);
    if !root.join("manifest.json").is_file() {
        bail!("backup {} is corrupted (no manifest)", info.id);
    }

    // Stop the agents being restored (only managed ones we know).
    let names: Vec<String> = if info.agents.is_empty() {
        app.state
            .all()?
            .keys()
            .cloned()
            .collect()
    } else {
        info.agents.clone()
    };
    for name in &names {
        if let Ok(Some(entry)) = app.state.get(name) {
            if let Some(pid) = entry.pid {
                if crate::process::is_running(pid) {
                    app.log
                        .info(&format!("stopping {name} (pid {pid}) before rollback"));
                    let _ = crate::commands::run_cmd::stop(app, name, false, None);
                }
            }
        }
    }

    // 1. Restore the state database.
    if info.state {
        let src = root.join("state.json");
        if src.is_file() {
            std::fs::copy(&src, app.state.path())
                .with_context(|| "cannot restore state.json".to_string())?;
        }
    }

    // 2. Restore the user config.
    if info.config {
        if let Some(origin) = &app.config_origin {
            if let Some(parent) = origin.parent() {
                std::fs::create_dir_all(parent)?;
            }
            std::fs::copy(root.join("config.yaml"), origin)
                .with_context(|| format!("cannot restore {}", origin.display()))?;
        }
    }

    // 3. Restore the install directories.
    let src_root = root.join("install");
    if src_root.is_dir() {
        if info.agents.is_empty() {
            if app.paths.install_dir.exists() {
                std::fs::remove_dir_all(&app.paths.install_dir)
                    .with_context(|| "cannot clear the install directory".to_string())?;
            }
            copy_tree(&src_root, &app.paths.install_dir)?;
        } else {
            for name in &info.agents {
                let src = src_root.join(name);
                if !src.exists() {
                    continue;
                }
                // Remove the agent's current install dir (known from the
                // restored state) and copy the snapshot back.
                if let Ok(Some(entry)) = app.state.get(name) {
                    let cur = PathBuf::from(&entry.install_dir);
                    if cur.is_dir() {
                        let _ = std::fs::remove_dir_all(&cur);
                    }
                }
                let dest = app.paths.install_dir.join(name);
                copy_tree(&src, &dest)?;
            }
            // Restore the shims.
            let bin_src = src_root.join("bin");
            if bin_src.is_dir() {
                copy_tree(&bin_src, &app.paths.bin_dir)?;
            }
        }
    }

    Ok(info)
}

/// Enforce settings.backups_keep: delete the oldest backups beyond the cap.
pub fn prune(app: &App) -> Result<()> {
    let keep = app.config.settings.backups_keep.unwrap_or(DEFAULT_KEEP) as usize;
    let all = list_backups(app);
    if all.len() <= keep {
        return Ok(());
    }
    for old in all.iter().take(all.len() - keep) {
        let dir = backups_dir(app).join(&old.id);
        if dir.is_dir() {
            let _ = std::fs::remove_dir_all(&dir);
        }
    }
    Ok(())
}

/// Recursively copy a directory tree (files and directories only; symlinks
/// are copied as-is when possible).
fn copy_tree(src: &Path, dest: &Path) -> Result<()> {
    if !src.exists() {
        return Ok(());
    }
    let meta = std::fs::metadata(src)?;
    if meta.is_file() {
        if let Some(parent) = dest.parent() {
            std::fs::create_dir_all(parent)?;
        }
        std::fs::copy(src, dest)?;
        return Ok(());
    }
    std::fs::create_dir_all(dest)?;
    let mut entries: Vec<_> = std::fs::read_dir(src)?
        .flatten()
        .map(|e| e.path())
        .collect();
    entries.sort();
    for child in entries {
        let name = child
            .file_name()
            .ok_or_else(|| anyhow!("invalid path {}", child.display()))?;
        copy_tree(&child, &dest.join(name))?;
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::Cli;
    use crate::state::StateFile;
    use clap::Parser;
    use std::io::Write;

    fn test_app(tag: &str) -> App {
        let dir = std::env::temp_dir().join(format!("am-backup-{tag}-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let cfg = dir.join("config.yaml");
        std::fs::write(&cfg, "version: \"1.0\"\nagents: []\n").unwrap();
        let cli = Cli::parse_from([
            "am",
            "--config",
            cfg.to_str().unwrap(),
            "--yes",
            "--no-color",
        ]);
        let mut app = App::from_cli(cli).expect("app builds");
        // Point the state/install dirs into the temp dir so nothing touches
        // the real user data.
        let state_file = dir.join("state.json");
        let install = dir.join("agents");
        std::fs::create_dir_all(&install).unwrap();
        app.paths.state_file = state_file.clone();
        app.paths.install_dir = install.clone();
        app.paths.bin_dir = install.join("bin");
        app.state = crate::state::StateStore::new(state_file);
        app.config_origin = Some(cfg);
        // Fresh backups dir for this test run.
        let _ = std::fs::remove_dir_all(backups_dir(&app));
        app
    }

    fn write_agent_file(app: &App, name: &str, content: &str) -> PathBuf {
        let dir = app.paths.install_dir.join(name);
        std::fs::create_dir_all(&dir).unwrap();
        let f = dir.join("agent.txt");
        let mut h = std::fs::File::create(&f).unwrap();
        h.write_all(content.as_bytes()).unwrap();
        f
    }

    #[test]
    fn backup_create_list_restore_roundtrip() {
        let app = test_app("roundtrip");
        write_agent_file(&app, "alpha", "v1");
        write_agent_file(&app, "beta", "b1");
        // Simulate a managed entry for alpha.
        let mut sf = StateFile::default();
        sf.installed.insert(
            "alpha".to_string(),
            crate::state::InstalledEntry {
                name: "alpha".to_string(),
                version: Some("1.0.0".to_string()),
                method: "binary".to_string(),
                install_dir: app.paths.install_dir.join("alpha").display().to_string(),
                bins: vec![],
                run: "alpha".to_string(),
                installed_at: "2026-01-01T00:00:00Z".to_string(),
                updated_at: None,
                pid: None,
                started_at: None,
            },
        );
        app.state.save(&sf).unwrap();

        let info = create_backup(&app, "update:alpha", &["alpha".to_string()]).unwrap();
        assert!(info.state && info.config);

        // Corrupt the live install + state.
        write_agent_file(&app, "alpha", "BROKEN");
        let mut bad = StateFile::default();
        bad.sessions_count = 999;
        app.state.save(&bad).unwrap();

        restore_backup(&app, Some(&info.id)).unwrap();
        let text = std::fs::read_to_string(app.paths.install_dir.join("alpha/agent.txt")).unwrap();
        assert_eq!(text, "v1", "install dir restored");
        let sf = app.state.load().unwrap();
        assert_eq!(sf.installed.get("alpha").unwrap().version.as_deref(), Some("1.0.0"));
        assert_eq!(sf.sessions_count, 0, "state restored");
    }

    #[test]
    fn resolve_point_prefix_and_latest() {
        let app = test_app("resolve");
        create_backup(&app, "update:a", &[]).unwrap();
        let info = resolve_point(&app, None).unwrap();
        assert_eq!(info.agents, Vec::<String>::new());
        let prefix = &info.id[..info.id.len().min(10)];
        let by_prefix = resolve_point(&app, Some(prefix)).unwrap();
        assert_eq!(by_prefix.id, info.id);
    }

    #[test]
    fn prune_keeps_only_backups_keep() {
        let mut app = test_app("prune");
        app.config.settings.backups_keep = Some(1);
        create_backup(&app, "update:one", &[]).unwrap();
        create_backup(&app, "update:two", &[]).unwrap();
        let all = list_backups(&app);
        assert_eq!(all.len(), 1, "only the newest backup survives");
    }
}