archy/core/archipelago/src/update.rs

//! Update system: check for updates, download deltas, apply with rollback.

use anyhow::{Context, Result};
use chrono::Timelike;
use serde::{Deserialize, Serialize};
use std::path::Path;
use tokio::fs;
use tracing::{debug, info, warn};

const DEFAULT_UPDATE_MANIFEST_URL: &str =
    "https://raw.githubusercontent.com/archipelago-os/releases/main/manifest.json";
const UPDATE_STATE_FILE: &str = "update_state.json";

fn update_manifest_url() -> String {
    std::env::var("ARCHIPELAGO_UPDATE_URL").unwrap_or_else(|_| DEFAULT_UPDATE_MANIFEST_URL.to_string())
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateManifest {
    pub version: String,
    pub release_date: String,
    pub changelog: Vec<String>,
    pub components: Vec<ComponentUpdate>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComponentUpdate {
    pub name: String,
    pub current_version: String,
    pub new_version: String,
    pub download_url: String,
    pub sha256: String,
    pub size_bytes: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UpdateSchedule {
    Manual,
    DailyCheck,
    AutoApply,
}

impl Default for UpdateSchedule {
    fn default() -> Self {
        Self::DailyCheck
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateState {
    pub current_version: String,
    pub last_check: Option<String>,
    pub available_update: Option<UpdateManifest>,
    pub update_in_progress: bool,
    pub rollback_available: bool,
    #[serde(default)]
    pub schedule: UpdateSchedule,
}

impl Default for UpdateState {
    fn default() -> Self {
        Self {
            current_version: env!("CARGO_PKG_VERSION").to_string(),
            last_check: None,
            available_update: None,
            update_in_progress: false,
            rollback_available: false,
            schedule: UpdateSchedule::DailyCheck,
        }
    }
}

pub async fn load_state(data_dir: &Path) -> Result<UpdateState> {
    let path = data_dir.join(UPDATE_STATE_FILE);
    if !path.exists() {
        let state = UpdateState::default();
        save_state(data_dir, &state).await?;
        return Ok(state);
    }
    let data = fs::read_to_string(&path)
        .await
        .context("Reading update state")?;
    serde_json::from_str(&data).context("Parsing update state")
}

pub async fn save_state(data_dir: &Path, state: &UpdateState) -> Result<()> {
    let path = data_dir.join(UPDATE_STATE_FILE);
    let data = serde_json::to_string_pretty(state)?;
    fs::write(&path, data)
        .await
        .context("Writing update state")
}

/// Check for available updates by fetching the release manifest.
pub async fn check_for_updates(data_dir: &Path) -> Result<UpdateState> {
    let mut state = load_state(data_dir).await?;

    info!("Checking for updates...");
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(15))
        .build()
        .context("Failed to create HTTP client")?;

    let manifest_url = update_manifest_url();
    match client.get(&manifest_url).send().await {
        Ok(resp) if resp.status().is_success() => {
            let manifest: UpdateManifest = resp
                .json()
                .await
                .context("Failed to parse update manifest")?;

            if manifest.version != state.current_version {
                info!(
                    current = %state.current_version,
                    available = %manifest.version,
                    "Update available"
                );
                state.available_update = Some(manifest);
            } else {
                debug!("Already on latest version: {}", state.current_version);
                state.available_update = None;
            }
        }
        Ok(resp) => {
            debug!("Update check returned status: {}", resp.status());
        }
        Err(e) => {
            debug!("Update check failed (offline?): {}", e);
        }
    }

    state.last_check = Some(chrono::Utc::now().to_rfc3339());
    save_state(data_dir, &state).await?;
    Ok(state)
}

/// Get current update status without checking remote.
pub async fn get_status(data_dir: &Path) -> Result<UpdateState> {
    load_state(data_dir).await
}

/// Dismiss the available update notification.
pub async fn dismiss_update(data_dir: &Path) -> Result<()> {
    let mut state = load_state(data_dir).await?;
    state.available_update = None;
    save_state(data_dir, &state).await
}

/// Download update components to a staging directory.
/// Verifies SHA256 hash for each component.
pub async fn download_update(data_dir: &Path) -> Result<DownloadProgress> {
    let state = load_state(data_dir).await?;
    let manifest = state
        .available_update
        .as_ref()
        .ok_or_else(|| anyhow::anyhow!("No update available to download"))?;

    let staging_dir = data_dir.join("update-staging");
    fs::create_dir_all(&staging_dir)
        .await
        .context("Failed to create staging dir")?;

    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(300))
        .build()
        .context("Failed to create HTTP client")?;

    let mut downloaded = 0u64;
    let total_bytes: u64 = manifest.components.iter().map(|c| c.size_bytes).sum();

    for component in &manifest.components {
        info!(name = %component.name, url = %component.download_url, "Downloading component");

        let resp = client
            .get(&component.download_url)
            .send()
            .await
            .with_context(|| format!("Failed to download {}", component.name))?;

        if !resp.status().is_success() {
            anyhow::bail!(
                "Download failed for {}: HTTP {}",
                component.name,
                resp.status()
            );
        }

        let bytes = resp
            .bytes()
            .await
            .with_context(|| format!("Failed to read {}", component.name))?;

        // Verify SHA256
        use sha2::{Digest, Sha256};
        let hash = hex::encode(Sha256::digest(&bytes));
        if hash != component.sha256 {
            anyhow::bail!(
                "SHA256 mismatch for {}: expected {}, got {}",
                component.name,
                component.sha256,
                hash
            );
        }

        let dest = staging_dir.join(&component.name);
        fs::write(&dest, &bytes)
            .await
            .with_context(|| format!("Failed to write {}", component.name))?;

        downloaded += component.size_bytes;
        info!(
            name = %component.name,
            bytes = bytes.len(),
            "Component downloaded and verified"
        );
    }

    // Mark update as downloaded
    let mut state = load_state(data_dir).await?;
    state.update_in_progress = true;
    save_state(data_dir, &state).await?;

    Ok(DownloadProgress {
        total_bytes,
        downloaded_bytes: downloaded,
        components_downloaded: manifest.components.len(),
        staging_dir: staging_dir.to_string_lossy().to_string(),
    })
}

/// Apply a downloaded update. Backs up current binaries, replaces with staged versions.
pub async fn apply_update(data_dir: &Path) -> Result<()> {
    let staging_dir = data_dir.join("update-staging");
    if !staging_dir.exists() {
        anyhow::bail!("No staged update found. Download first.");
    }

    let backup_dir = data_dir.join("update-backup");
    fs::create_dir_all(&backup_dir)
        .await
        .context("Failed to create backup dir")?;

    // Back up current backend binary
    let current_binary = Path::new("/usr/local/bin/archipelago");
    if current_binary.exists() {
        let backup_path = backup_dir.join("archipelago");
        fs::copy(current_binary, &backup_path)
            .await
            .context("Failed to backup current binary")?;
        info!("Current binary backed up");
    }

    // Apply staged components
    let mut entries = fs::read_dir(&staging_dir)
        .await
        .context("Failed to read staging dir")?;

    while let Some(entry) = entries.next_entry().await? {
        let name = entry.file_name().to_string_lossy().to_string();
        let src = entry.path();

        // Map component names to destinations
        let dest = match name.as_str() {
            "archipelago" => Some(Path::new("/usr/local/bin/archipelago").to_path_buf()),
            _ => {
                // For frontend or config files, determine destination
                if name.ends_with(".tar.gz") || name.ends_with(".zip") {
                    // Archive — extract to appropriate location
                    debug!(name = %name, "Skipping archive (manual extraction needed)");
                    None
                } else {
                    debug!(name = %name, "Unknown component, skipping");
                    None
                }
            }
        };

        if let Some(dest_path) = dest {
            fs::copy(&src, &dest_path)
                .await
                .with_context(|| format!("Failed to apply {}", name))?;
            info!(name = %name, dest = %dest_path.display(), "Component applied");
        }
    }

    // Update state
    let mut state = load_state(data_dir).await?;
    if let Some(manifest) = &state.available_update {
        state.current_version = manifest.version.clone();
    }
    state.available_update = None;
    state.update_in_progress = false;
    state.rollback_available = true;
    save_state(data_dir, &state).await?;

    // Clean staging
    let _ = fs::remove_dir_all(&staging_dir).await;

    info!("Update applied. Restart service to take effect.");
    Ok(())
}

/// Rolling container restart — restarts containers one at a time with health checks.
/// This enables zero-downtime updates for containerized apps.
pub async fn rolling_container_restart() -> Result<RollingRestartReport> {
    use std::process::Command;

    let output = Command::new("sudo")
        .args(["podman", "ps", "--format", "{{.Names}}"])
        .output()
        .context("Failed to list containers")?;
    let names: Vec<String> = String::from_utf8_lossy(&output.stdout)
        .lines()
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect();

    let total = names.len();
    let mut restarted = 0;
    let mut failed = Vec::new();

    info!(total = total, "Starting rolling container restart");

    for name in &names {
        debug!(container = %name, "Restarting container");

        let restart = Command::new("sudo")
            .args(["podman", "restart", "--time", "30", name])
            .output();

        match restart {
            Ok(out) if out.status.success() => {
                // Wait for container to be healthy
                let mut healthy = false;
                for _ in 0..12 {
                    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                    let check = Command::new("sudo")
                        .args(["podman", "inspect", name, "--format", "{{.State.Status}}"])
                        .output();
                    if let Ok(out) = check {
                        let status = String::from_utf8_lossy(&out.stdout).trim().to_string();
                        if status == "running" {
                            healthy = true;
                            break;
                        }
                    }
                }
                if healthy {
                    restarted += 1;
                    debug!(container = %name, "Container restarted successfully");
                } else {
                    failed.push(name.clone());
                    warn!(container = %name, "Container not healthy after restart");
                }
            }
            _ => {
                failed.push(name.clone());
                warn!(container = %name, "Container restart command failed");
            }
        }
    }

    info!(restarted = restarted, failed = failed.len(), "Rolling restart complete");

    Ok(RollingRestartReport {
        total,
        restarted,
        failed,
    })
}

/// Report from a rolling container restart.
#[derive(Debug, Serialize, Deserialize)]
pub struct RollingRestartReport {
    pub total: usize,
    pub restarted: usize,
    pub failed: Vec<String>,
}

/// Rollback to the previous version from backup.
pub async fn rollback_update(data_dir: &Path) -> Result<()> {
    let backup_dir = data_dir.join("update-backup");
    if !backup_dir.exists() {
        anyhow::bail!("No rollback backup available");
    }

    let backup_binary = backup_dir.join("archipelago");
    if backup_binary.exists() {
        fs::copy(&backup_binary, "/usr/local/bin/archipelago")
            .await
            .context("Failed to restore backup binary")?;
        info!("Binary rolled back to previous version");
    }

    let mut state = load_state(data_dir).await?;
    state.rollback_available = false;
    save_state(data_dir, &state).await?;

    let _ = fs::remove_dir_all(&backup_dir).await;

    info!("Rollback complete. Restart service to take effect.");
    Ok(())
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DownloadProgress {
    pub total_bytes: u64,
    pub downloaded_bytes: u64,
    pub components_downloaded: usize,
    pub staging_dir: String,
}

/// Set the update schedule preference.
pub async fn set_schedule(data_dir: &Path, schedule: UpdateSchedule) -> Result<()> {
    let mut state = load_state(data_dir).await?;
    state.schedule = schedule;
    save_state(data_dir, &state).await?;
    info!(schedule = ?schedule, "Update schedule changed");
    Ok(())
}

/// Get the current schedule.
pub async fn get_schedule(data_dir: &Path) -> Result<UpdateSchedule> {
    let state = load_state(data_dir).await?;
    Ok(state.schedule)
}

/// Background update scheduler. Runs in a loop, checking/applying based on schedule.
/// Call this once at startup via `tokio::spawn`.
pub async fn run_update_scheduler(data_dir: std::path::PathBuf) {
    use tokio::time::{interval, Duration};

    // Check every hour; act based on schedule setting
    let mut tick = interval(Duration::from_secs(3600));

    loop {
        tick.tick().await;

        let state = match load_state(&data_dir).await {
            Ok(s) => s,
            Err(e) => {
                debug!("Update scheduler: failed to load state: {}", e);
                continue;
            }
        };

        match state.schedule {
            UpdateSchedule::Manual => {
                debug!("Update scheduler: manual mode, skipping");
                continue;
            }
            UpdateSchedule::DailyCheck => {
                // Only check once per day
                if let Some(ref last) = state.last_check {
                    if let Ok(last_time) = chrono::DateTime::parse_from_rfc3339(last) {
                        let elapsed = chrono::Utc::now() - last_time.with_timezone(&chrono::Utc);
                        if elapsed.num_hours() < 24 {
                            debug!("Update scheduler: checked recently, skipping");
                            continue;
                        }
                    }
                }
                info!("Update scheduler: running daily check");
                if let Err(e) = check_for_updates(&data_dir).await {
                    debug!("Update scheduler: check failed: {}", e);
                }
            }
            UpdateSchedule::AutoApply => {
                // Auto-apply: check, download, and apply during 3 AM window
                let hour = chrono::Local::now().hour();
                if hour != 3 {
                    // Still do daily check outside the window
                    if let Some(ref last) = state.last_check {
                        if let Ok(last_time) = chrono::DateTime::parse_from_rfc3339(last) {
                            let elapsed =
                                chrono::Utc::now() - last_time.with_timezone(&chrono::Utc);
                            if elapsed.num_hours() < 24 {
                                continue;
                            }
                        }
                    }
                    info!("Update scheduler: auto-apply check (outside window)");
                    if let Err(e) = check_for_updates(&data_dir).await {
                        debug!("Update scheduler: check failed: {}", e);
                    }
                    continue;
                }

                // 3 AM — check, download, and apply
                info!("Update scheduler: 3 AM auto-apply window");
                match check_for_updates(&data_dir).await {
                    Ok(s) if s.available_update.is_some() => {
                        info!("Update scheduler: downloading update");
                        if let Err(e) = download_update(&data_dir).await {
                            debug!("Update scheduler: download failed: {}", e);
                            continue;
                        }
                        info!("Update scheduler: applying update");
                        if let Err(e) = apply_update(&data_dir).await {
                            debug!("Update scheduler: apply failed: {}", e);
                            continue;
                        }
                        info!("Update scheduler: update applied, restart needed");
                        // Signal for service restart (systemd will handle via exit code)
                        std::process::exit(0);
                    }
                    Ok(_) => {
                        debug!("Update scheduler: no update available");
                    }
                    Err(e) => {
                        debug!("Update scheduler: check failed: {}", e);
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_update_schedule_default_is_daily_check() {
        let schedule = UpdateSchedule::default();
        assert_eq!(schedule, UpdateSchedule::DailyCheck);
    }

    #[test]
    fn test_update_state_default_values() {
        let state = UpdateState::default();
        assert_eq!(state.current_version, env!("CARGO_PKG_VERSION"));
        assert!(state.last_check.is_none());
        assert!(state.available_update.is_none());
        assert!(!state.update_in_progress);
        assert!(!state.rollback_available);
        assert_eq!(state.schedule, UpdateSchedule::DailyCheck);
    }

    #[test]
    fn test_update_state_serialization_roundtrip() {
        let state = UpdateState {
            current_version: "0.2.0".to_string(),
            last_check: Some("2025-01-01T00:00:00Z".to_string()),
            available_update: None,
            update_in_progress: false,
            rollback_available: true,
            schedule: UpdateSchedule::AutoApply,
        };
        let json = serde_json::to_string(&state).unwrap();
        let deserialized: UpdateState = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.current_version, "0.2.0");
        assert!(deserialized.rollback_available);
        assert_eq!(deserialized.schedule, UpdateSchedule::AutoApply);
    }

    #[test]
    fn test_update_schedule_serde_rename() {
        let json = serde_json::to_string(&UpdateSchedule::DailyCheck).unwrap();
        assert_eq!(json, "\"daily_check\"");
        let json = serde_json::to_string(&UpdateSchedule::Manual).unwrap();
        assert_eq!(json, "\"manual\"");
        let json = serde_json::to_string(&UpdateSchedule::AutoApply).unwrap();
        assert_eq!(json, "\"auto_apply\"");
    }

    #[test]
    fn test_update_state_schedule_defaults_on_missing_field() {
        // When schedule field is missing from JSON, it should default to DailyCheck
        let json = r#"{
            "current_version": "0.1.0",
            "last_check": null,
            "available_update": null,
            "update_in_progress": false,
            "rollback_available": false
        }"#;
        let state: UpdateState = serde_json::from_str(json).unwrap();
        assert_eq!(state.schedule, UpdateSchedule::DailyCheck);
    }

    #[tokio::test]
    async fn test_load_state_creates_default_when_missing() {
        let dir = tempfile::tempdir().unwrap();
        let state = load_state(dir.path()).await.unwrap();
        assert_eq!(state.current_version, env!("CARGO_PKG_VERSION"));
        assert!(!state.update_in_progress);
        // File should now exist after load created the default
        assert!(dir.path().join(UPDATE_STATE_FILE).exists());
    }

    #[tokio::test]
    async fn test_save_and_load_state_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let state = UpdateState {
            current_version: "1.0.0".to_string(),
            last_check: Some("2025-06-15T12:00:00Z".to_string()),
            available_update: Some(UpdateManifest {
                version: "1.1.0".to_string(),
                release_date: "2025-06-20".to_string(),
                changelog: vec!["Fix bugs".to_string(), "New feature".to_string()],
                components: vec![ComponentUpdate {
                    name: "archipelago".to_string(),
                    current_version: "1.0.0".to_string(),
                    new_version: "1.1.0".to_string(),
                    download_url: "https://example.com/binary".to_string(),
                    sha256: "abc123".to_string(),
                    size_bytes: 5000,
                }],
            }),
            update_in_progress: true,
            rollback_available: false,
            schedule: UpdateSchedule::Manual,
        };
        save_state(dir.path(), &state).await.unwrap();
        let loaded = load_state(dir.path()).await.unwrap();
        assert_eq!(loaded.current_version, "1.0.0");
        assert!(loaded.update_in_progress);
        assert_eq!(loaded.schedule, UpdateSchedule::Manual);
        let manifest = loaded.available_update.unwrap();
        assert_eq!(manifest.version, "1.1.0");
        assert_eq!(manifest.components.len(), 1);
        assert_eq!(manifest.components[0].size_bytes, 5000);
    }

    #[tokio::test]
    async fn test_dismiss_update_clears_available() {
        let dir = tempfile::tempdir().unwrap();
        let state = UpdateState {
            available_update: Some(UpdateManifest {
                version: "2.0.0".to_string(),
                release_date: "2025-07-01".to_string(),
                changelog: vec![],
                components: vec![],
            }),
            ..UpdateState::default()
        };
        save_state(dir.path(), &state).await.unwrap();
        dismiss_update(dir.path()).await.unwrap();
        let loaded = load_state(dir.path()).await.unwrap();
        assert!(loaded.available_update.is_none());
    }

    #[tokio::test]
    async fn test_set_and_get_schedule() {
        let dir = tempfile::tempdir().unwrap();
        // Initialize state
        let _ = load_state(dir.path()).await.unwrap();

        set_schedule(dir.path(), UpdateSchedule::AutoApply).await.unwrap();
        let schedule = get_schedule(dir.path()).await.unwrap();
        assert_eq!(schedule, UpdateSchedule::AutoApply);

        set_schedule(dir.path(), UpdateSchedule::Manual).await.unwrap();
        let schedule = get_schedule(dir.path()).await.unwrap();
        assert_eq!(schedule, UpdateSchedule::Manual);
    }

    #[tokio::test]
    async fn test_get_status_returns_current_state() {
        let dir = tempfile::tempdir().unwrap();
        let state = UpdateState {
            current_version: "3.0.0".to_string(),
            rollback_available: true,
            ..UpdateState::default()
        };
        save_state(dir.path(), &state).await.unwrap();
        let status = get_status(dir.path()).await.unwrap();
        assert_eq!(status.current_version, "3.0.0");
        assert!(status.rollback_available);
    }
}