a3s-box-core 3.1.0

//! VM Snapshot Types — Configuration-based snapshot metadata.
//!
//! Snapshots capture the full VM configuration (not memory state) so a box
//! can be reconstructed from the saved spec. Combined with rootfs caching,
//! restore achieves sub-500ms cold start.

use crate::error::{BoxError, Result};
use crate::traits::ExecutionHealthCheck;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

use crate::config::DEFAULT_VCPUS;

/// Resolved OCI image defaults required to reproduce a captured rootfs.
///
/// These values are distinct from [`SnapshotMetadata`] command, environment,
/// and working-directory fields: those fields are user overrides, while this
/// structure records the immutable defaults resolved from the source image.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SnapshotImageConfig {
    /// Image entrypoint.
    #[serde(default)]
    pub entrypoint: Option<Vec<String>>,
    /// Image command arguments.
    #[serde(default)]
    pub cmd: Option<Vec<String>>,
    /// Image environment in declaration order.
    #[serde(default)]
    pub env: Vec<(String, String)>,
    /// Image working directory.
    #[serde(default)]
    pub working_dir: Option<String>,
    /// Image user.
    #[serde(default)]
    pub user: Option<String>,
    /// Ports declared by the image.
    #[serde(default)]
    pub exposed_ports: Vec<String>,
    /// Image labels.
    #[serde(default)]
    pub labels: HashMap<String, String>,
    /// Volumes declared by the image.
    #[serde(default)]
    pub volumes: Vec<String>,
    /// Image stop signal.
    #[serde(default)]
    pub stop_signal: Option<String>,
    /// Image health check.
    #[serde(default)]
    pub health_check: Option<SnapshotImageHealthCheck>,
    /// Image ONBUILD triggers.
    #[serde(default)]
    pub onbuild: Vec<String>,
}

/// Health-check defaults embedded in a resolved OCI image configuration.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SnapshotImageHealthCheck {
    /// Health-check command.
    #[serde(default)]
    pub test: Vec<String>,
    /// Interval in seconds.
    #[serde(default)]
    pub interval: Option<u64>,
    /// Timeout in seconds.
    #[serde(default)]
    pub timeout: Option<u64>,
    /// Retry count.
    #[serde(default)]
    pub retries: Option<u32>,
    /// Start period in seconds.
    #[serde(default)]
    pub start_period: Option<u64>,
}

impl SnapshotImageHealthCheck {
    /// Whether the OCI health-check declaration contains an executable command.
    pub fn is_enabled(&self) -> bool {
        let Some(marker) = self.test.first() else {
            return false;
        };
        if marker.eq_ignore_ascii_case("NONE") {
            return false;
        }
        if marker.eq_ignore_ascii_case("CMD") || marker.eq_ignore_ascii_case("CMD-SHELL") {
            return self
                .test
                .get(1..)
                .is_some_and(|command| command.iter().any(|part| !part.trim().is_empty()));
        }
        self.test.iter().any(|part| !part.trim().is_empty())
    }
}

/// Metadata for a saved VM snapshot.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotMetadata {
    /// Unique snapshot identifier
    pub id: String,
    /// User-assigned name (or auto-generated)
    pub name: String,
    /// Box ID this snapshot was taken from
    pub source_box_id: String,
    /// OCI image reference used by the source box
    pub image: String,
    /// Number of vCPUs
    pub vcpus: u32,
    /// Memory in MB
    pub memory_mb: u32,
    /// Volume mounts (host:guest pairs)
    pub volumes: Vec<String>,
    /// Environment variables
    pub env: HashMap<String, String>,
    /// Command override
    pub cmd: Vec<String>,
    /// Entrypoint override
    #[serde(default)]
    pub entrypoint: Option<Vec<String>>,
    /// Working directory inside the box
    #[serde(default)]
    pub workdir: Option<String>,
    /// Resolved defaults from the source OCI image.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub image_config: Option<SnapshotImageConfig>,
    /// Explicit health check selected for the source box.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub health_check: Option<ExecutionHealthCheck>,
    /// Whether health checks were explicitly disabled for the source box.
    #[serde(default)]
    pub healthcheck_disabled: bool,
    /// Port mappings
    #[serde(default)]
    pub port_map: Vec<String>,
    /// User-defined labels
    #[serde(default)]
    pub labels: HashMap<String, String>,
    /// Network mode
    #[serde(default)]
    pub network_mode: Option<String>,
    /// Rootfs cache key (for fast restore via cached rootfs)
    #[serde(default)]
    pub rootfs_cache_key: Option<String>,
    /// Size of the snapshot on disk in bytes
    #[serde(default)]
    pub size_bytes: u64,
    /// Creation timestamp
    #[serde(default = "epoch")]
    pub created_at: DateTime<Utc>,
    /// User-provided description
    #[serde(default)]
    pub description: String,
}

/// Default `created_at` for metadata written by an older/partial build that
/// omitted the field. Without a default a single missing field would make
/// serde reject the whole record, dropping the snapshot from `list()`/`prune()`
/// (invisible, un-prunable, leaking disk) — see SnapshotStore::list.
fn epoch() -> DateTime<Utc> {
    DateTime::<Utc>::UNIX_EPOCH
}

impl SnapshotMetadata {
    /// Create a new snapshot metadata with required fields.
    pub fn new(id: String, name: String, source_box_id: String, image: String) -> Self {
        Self {
            id,
            name,
            source_box_id,
            image,
            vcpus: DEFAULT_VCPUS,
            memory_mb: 512,
            volumes: Vec::new(),
            env: HashMap::new(),
            cmd: Vec::new(),
            entrypoint: None,
            workdir: None,
            image_config: None,
            health_check: None,
            healthcheck_disabled: false,
            port_map: Vec::new(),
            labels: HashMap::new(),
            network_mode: None,
            rootfs_cache_key: None,
            size_bytes: 0,
            created_at: Utc::now(),
            description: String::new(),
        }
    }

    /// Set description.
    pub fn with_description(mut self, desc: &str) -> Self {
        self.description = desc.to_string();
        self
    }

    /// Set resources.
    pub fn with_resources(mut self, vcpus: u32, memory_mb: u32) -> Self {
        self.vcpus = vcpus;
        self.memory_mb = memory_mb;
        self
    }

    /// Return the captured OCI image defaults required for a safe restore.
    ///
    /// Snapshot records created before this field existed remain readable so
    /// operators can inspect and delete them, but restoring one would lose
    /// image entrypoint, environment, user, and working-directory semantics.
    pub fn require_image_config(&self) -> Result<&SnapshotImageConfig> {
        self.image_config.as_ref().ok_or_else(|| {
            BoxError::ConfigError(format!(
                "Snapshot '{}' does not contain resolved OCI image configuration and cannot be restored safely; recreate it with the current A3S Box version",
                self.id
            ))
        })
    }

    /// Whether restoring this snapshot would enable an explicit or image health check.
    pub fn has_effective_health_check(&self) -> bool {
        !self.healthcheck_disabled
            && (self.health_check.is_some()
                || self
                    .image_config
                    .as_ref()
                    .and_then(|config| config.health_check.as_ref())
                    .is_some_and(SnapshotImageHealthCheck::is_enabled))
    }
}

/// Configuration for snapshot operations.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotConfig {
    /// Whether snapshots are enabled
    pub enabled: bool,
    /// Directory to store snapshots (default: ~/.a3s/snapshots)
    pub snapshot_dir: Option<PathBuf>,
    /// Maximum number of snapshots to keep (0 = unlimited)
    pub max_snapshots: usize,
    /// Maximum total size in bytes (0 = unlimited)
    pub max_total_bytes: u64,
}

impl Default for SnapshotConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            snapshot_dir: None,
            max_snapshots: 0,
            max_total_bytes: 0,
        }
    }
}

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

    #[test]
    fn test_snapshot_metadata_new() {
        let meta = SnapshotMetadata::new(
            "snap-001".to_string(),
            "my-snapshot".to_string(),
            "box-abc".to_string(),
            "alpine:latest".to_string(),
        );
        assert_eq!(meta.id, "snap-001");
        assert_eq!(meta.name, "my-snapshot");
        assert_eq!(meta.source_box_id, "box-abc");
        assert_eq!(meta.image, "alpine:latest");
        assert_eq!(meta.vcpus, DEFAULT_VCPUS);
        assert_eq!(meta.memory_mb, 512);
        assert!(meta.volumes.is_empty());
        assert!(meta.env.is_empty());
        assert!(meta.description.is_empty());
        assert!(meta.image_config.is_none());
        assert!(meta.health_check.is_none());
        assert!(!meta.healthcheck_disabled);
    }

    #[test]
    fn test_snapshot_metadata_with_description() {
        let meta = SnapshotMetadata::new(
            "snap-002".to_string(),
            "test".to_string(),
            "box-xyz".to_string(),
            "ubuntu:22.04".to_string(),
        )
        .with_description("Before migration");
        assert_eq!(meta.description, "Before migration");
    }

    #[test]
    fn test_snapshot_metadata_with_resources() {
        let meta = SnapshotMetadata::new(
            "snap-003".to_string(),
            "test".to_string(),
            "box-xyz".to_string(),
            "python:3.12".to_string(),
        )
        .with_resources(4, 2048);
        assert_eq!(meta.vcpus, 4);
        assert_eq!(meta.memory_mb, 2048);
    }

    #[test]
    fn test_snapshot_metadata_tolerates_missing_size_and_created_at() {
        // Metadata written by an older/partial build (or lightly corrupted) that
        // omits size_bytes/created_at must still deserialize, so the snapshot
        // stays visible to list()/count()/prune() instead of being silently
        // dropped and leaking disk. Only the truly-required identity fields are
        // present here.
        let json = r#"{
            "id": "snap-old",
            "name": "old",
            "source_box_id": "box-1",
            "image": "alpine:latest",
            "vcpus": 2,
            "memory_mb": 512,
            "volumes": [],
            "env": {},
            "cmd": []
        }"#;
        let parsed: SnapshotMetadata =
            serde_json::from_str(json).expect("must tolerate missing size_bytes/created_at");
        assert_eq!(parsed.id, "snap-old");
        assert_eq!(parsed.size_bytes, 0);
        assert_eq!(parsed.created_at, DateTime::<Utc>::UNIX_EPOCH);
        assert!(parsed.health_check.is_none());
        assert!(!parsed.healthcheck_disabled);
        assert!(parsed.require_image_config().is_err());
    }

    #[test]
    fn test_snapshot_metadata_serde_roundtrip() {
        let mut meta = SnapshotMetadata::new(
            "snap-rt".to_string(),
            "roundtrip".to_string(),
            "box-123".to_string(),
            "nginx:latest".to_string(),
        );
        meta.vcpus = 4;
        meta.memory_mb = 1024;
        meta.volumes = vec!["/data:/data".to_string()];
        meta.env.insert("FOO".to_string(), "bar".to_string());
        meta.cmd = vec!["nginx".to_string(), "-g".to_string()];
        meta.entrypoint = Some(vec!["/docker-entrypoint.sh".to_string()]);
        meta.workdir = Some("/app".to_string());
        meta.image_config = Some(SnapshotImageConfig {
            entrypoint: Some(vec!["/usr/bin/envd".to_string()]),
            cmd: Some(vec!["--listen".to_string(), "0.0.0.0:49983".to_string()]),
            env: vec![
                ("PATH".to_string(), "/usr/local/bin:/usr/bin".to_string()),
                ("HOME".to_string(), "/home/user".to_string()),
            ],
            working_dir: Some("/home/user".to_string()),
            user: Some("1000:1000".to_string()),
            exposed_ports: vec!["49983/tcp".to_string()],
            labels: HashMap::from([("runtime".to_string(), "envd".to_string())]),
            volumes: vec!["/home/user".to_string()],
            stop_signal: Some("SIGTERM".to_string()),
            health_check: Some(SnapshotImageHealthCheck {
                test: vec!["CMD".to_string(), "envd-health".to_string()],
                interval: Some(10),
                timeout: Some(2),
                retries: Some(3),
                start_period: Some(5),
            }),
            onbuild: vec!["RUN prepare-runtime".to_string()],
        });
        meta.health_check = Some(ExecutionHealthCheck {
            cmd: vec!["test".to_string(), "-f".to_string(), "/ready".to_string()],
            interval_secs: 11,
            timeout_secs: 3,
            retries: 7,
            start_period_secs: 2,
        });
        meta.healthcheck_disabled = true;
        meta.port_map = vec!["8080:80".to_string()];
        meta.labels.insert("env".to_string(), "prod".to_string());
        meta.network_mode = Some("bridge".to_string());
        meta.rootfs_cache_key = Some("abc123".to_string());
        meta.size_bytes = 1024 * 1024;
        meta.description = "test snapshot".to_string();

        let json = serde_json::to_string(&meta).unwrap();
        let parsed: SnapshotMetadata = serde_json::from_str(&json).unwrap();

        assert_eq!(parsed.id, "snap-rt");
        assert_eq!(parsed.name, "roundtrip");
        assert_eq!(parsed.source_box_id, "box-123");
        assert_eq!(parsed.image, "nginx:latest");
        assert_eq!(parsed.vcpus, 4);
        assert_eq!(parsed.memory_mb, 1024);
        assert_eq!(parsed.volumes, vec!["/data:/data"]);
        assert_eq!(parsed.env.get("FOO").unwrap(), "bar");
        assert_eq!(parsed.cmd, vec!["nginx", "-g"]);
        assert_eq!(
            parsed.entrypoint,
            Some(vec!["/docker-entrypoint.sh".to_string()])
        );
        assert_eq!(parsed.workdir, Some("/app".to_string()));
        assert_eq!(parsed.image_config, meta.image_config);
        assert_eq!(parsed.health_check, meta.health_check);
        assert!(parsed.healthcheck_disabled);
        assert!(!parsed.has_effective_health_check());
        assert_eq!(parsed.port_map, vec!["8080:80"]);
        assert_eq!(parsed.labels.get("env").unwrap(), "prod");
        assert_eq!(parsed.network_mode, Some("bridge".to_string()));
        assert_eq!(parsed.rootfs_cache_key, Some("abc123".to_string()));
        assert_eq!(parsed.size_bytes, 1024 * 1024);
        assert_eq!(parsed.description, "test snapshot");
    }

    #[test]
    fn test_snapshot_metadata_deserialize_minimal() {
        let json = r#"{
            "id": "snap-min",
            "name": "minimal",
            "source_box_id": "box-1",
            "image": "alpine:latest",
            "vcpus": 1,
            "memory_mb": 256,
            "volumes": [],
            "env": {},
            "cmd": [],
            "size_bytes": 0,
            "created_at": "2024-01-01T00:00:00Z",
            "description": ""
        }"#;
        let meta: SnapshotMetadata = serde_json::from_str(json).unwrap();
        assert_eq!(meta.id, "snap-min");
        assert!(meta.entrypoint.is_none());
        assert!(meta.workdir.is_none());
        assert!(meta.image_config.is_none());
        assert!(meta.health_check.is_none());
        assert!(!meta.healthcheck_disabled);
        assert!(meta.port_map.is_empty());
        assert!(meta.labels.is_empty());
        assert!(meta.network_mode.is_none());
        assert!(meta.rootfs_cache_key.is_none());
    }

    #[test]
    fn test_snapshot_config_default() {
        let config = SnapshotConfig::default();
        assert!(config.enabled);
        assert!(config.snapshot_dir.is_none());
        assert_eq!(config.max_snapshots, 0);
        assert_eq!(config.max_total_bytes, 0);
    }

    #[test]
    fn test_snapshot_config_serde_roundtrip() {
        let config = SnapshotConfig {
            enabled: true,
            snapshot_dir: Some(PathBuf::from("/custom/snapshots")),
            max_snapshots: 10,
            max_total_bytes: 5 * 1024 * 1024 * 1024,
        };
        let json = serde_json::to_string(&config).unwrap();
        let parsed: SnapshotConfig = serde_json::from_str(&json).unwrap();
        assert!(parsed.enabled);
        assert_eq!(
            parsed.snapshot_dir,
            Some(PathBuf::from("/custom/snapshots"))
        );
        assert_eq!(parsed.max_snapshots, 10);
        assert_eq!(parsed.max_total_bytes, 5 * 1024 * 1024 * 1024);
    }

    #[test]
    fn test_snapshot_metadata_clone() {
        let meta = SnapshotMetadata::new(
            "snap-clone".to_string(),
            "clone-test".to_string(),
            "box-c".to_string(),
            "redis:7".to_string(),
        );
        let cloned = meta.clone();
        assert_eq!(cloned.id, meta.id);
        assert_eq!(cloned.name, meta.name);
        assert_eq!(cloned.source_box_id, meta.source_box_id);
    }

    #[test]
    fn effective_health_check_respects_none_and_disabled_declarations() {
        let mut meta = SnapshotMetadata::new(
            "snap-health".to_string(),
            "health".to_string(),
            "box-health".to_string(),
            "image:latest".to_string(),
        );
        meta.image_config = Some(SnapshotImageConfig {
            health_check: Some(SnapshotImageHealthCheck {
                test: vec!["NONE".to_string()],
                ..Default::default()
            }),
            ..Default::default()
        });
        assert!(!meta.has_effective_health_check());

        meta.image_config.as_mut().unwrap().health_check = Some(SnapshotImageHealthCheck {
            test: vec!["CMD".to_string(), "true".to_string()],
            ..Default::default()
        });
        assert!(meta.has_effective_health_check());

        meta.healthcheck_disabled = true;
        assert!(!meta.has_effective_health_check());
    }
}