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};

const DEFAULT_UPDATE_MANIFEST_URL: &str =
    "https://git.tx1138.com/lfg2025/archy/raw/branch/main/releases/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();

        match name.as_str() {
            "archipelago" => {
                let dest = Path::new("/usr/local/bin/archipelago");
                fs::copy(&src, dest)
                    .await
                    .with_context(|| format!("Failed to apply {}", name))?;
                #[cfg(unix)]
                {
                    use std::os::unix::fs::PermissionsExt;
                    std::fs::set_permissions(dest, std::fs::Permissions::from_mode(0o755))
                        .context("Failed to set binary permissions")?;
                }
                info!(name = %name, "Backend binary applied");
            }
            _ if name.contains("frontend") && name.ends_with(".tar.gz") => {
                let web_ui_dir = Path::new("/opt/archipelago/web-ui");
                // Back up current frontend
                let frontend_backup = backup_dir.join("web-ui-backup.tar.gz");
                if web_ui_dir.exists() {
                    let status = tokio::process::Command::new("tar")
                        .args(["-czf", &frontend_backup.to_string_lossy(), "-C", "/opt/archipelago", "web-ui"])
                        .status()
                        .await
                        .context("Failed to backup frontend")?;
                    if status.success() {
                        info!("Frontend backed up");
                    }
                }
                // Extract new frontend
                let status = tokio::process::Command::new("tar")
                    .args(["-xzf", &src.to_string_lossy(), "-C", "/opt/archipelago"])
                    .status()
                    .await
                    .with_context(|| format!("Failed to extract {}", name))?;
                if !status.success() {
                    anyhow::bail!("tar extraction failed for {}", name);
                }
                info!(name = %name, "Frontend archive extracted to /opt/archipelago/web-ui");
            }
            _ => {
                debug!(name = %name, "Unknown component, skipping");
            }
        }
    }

    // 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(())
}

/// 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);
    }
}