solstice-ci/crates/orchestrator/src/config.rs

use std::{
    collections::BTreeMap,
    fs,
    path::{Path, PathBuf},
};

use miette::{IntoDiagnostic as _, Result};
use serde::Deserialize;
use tokio::task;

#[derive(Debug, Clone, Deserialize)]
pub struct OrchestratorConfig {
    pub default_label: String,
    #[serde(default)]
    pub aliases: BTreeMap<String, String>,
    #[serde(default)]
    pub sizes: BTreeMap<String, SizePreset>,
    #[serde(default)]
    pub images: BTreeMap<String, ImageEntry>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct SizePreset {
    pub cpu: u16,
    pub ram_mb: u32,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ImageEntry {
    /// Remote source URL. If local_path does not exist, we will download it.
    pub source: String,
    /// Target local path for the prepared base image (raw .img or qcow2)
    pub local_path: PathBuf,
    /// Decompression method for downloaded artifact ("zstd" or "none"/missing)
    #[serde(default)]
    pub decompress: Option<Decompress>,
    /// Images must support NoCloud for metadata injection
    #[serde(default)]
    pub nocloud: bool,
    /// Default VM resource overrides for this label
    #[serde(default)]
    pub defaults: Option<ImageDefaults>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ImageDefaults {
    pub cpu: Option<u16>,
    pub ram_mb: Option<u32>,
    pub disk_gb: Option<u32>,
    /// Default SSH username for this image (e.g., ubuntu, root)
    pub ssh_user: Option<String>,
}

#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Decompress {
    Zstd,
    None,
}

impl Default for Decompress {
    fn default() -> Self {
        Decompress::None
    }
}

impl OrchestratorConfig {
    pub async fn load(path: Option<&Path>) -> Result<Self> {
        let path = match path {
            Some(p) => p.to_path_buf(),
            None => default_example_path(),
        };
        // Use blocking read via spawn_blocking to avoid blocking Tokio
        let cfg: OrchestratorConfig = task::spawn_blocking(move || {
            let builder = config::Config::builder().add_source(config::File::from(path));
            let cfg = builder.build().into_diagnostic()?;
            cfg.try_deserialize().into_diagnostic()
        })
        .await
        .into_diagnostic()??;
        Ok(cfg)
    }

    /// Resolve an incoming label using aliases to the canonical key used in `images`.
    pub fn resolve_label<'a>(&'a self, label: Option<&'a str>) -> Option<&'a str> {
        let l = label.unwrap_or(&self.default_label);
        if let Some(canon) = self.aliases.get(l) {
            Some(canon.as_str())
        } else {
            Some(l)
        }
    }

    /// Get image entry for a resolved label (canonical key)
    pub fn image_for(&self, resolved_label: &str) -> Option<&ImageEntry> {
        self.images.get(resolved_label)
    }
}

fn default_example_path() -> PathBuf {
    // default to examples/orchestrator-image-map.yaml relative to cwd
    PathBuf::from("examples/orchestrator-image-map.yaml")
}

/// Ensure images referenced in config exist at local_path. If missing, fetch
/// from `source` using vm-manager's ImageManager (supports http(s), file://,
/// oci://, streaming downloads, and zstd decompression).
pub async fn ensure_images(cfg: &OrchestratorConfig) -> Result<()> {
    let cache_dir = cfg
        .images
        .values()
        .next()
        .and_then(|img| img.local_path.parent())
        .unwrap_or(Path::new("/var/lib/solstice/images"));
    let img_mgr = vm_manager::image::ImageManager::with_cache_dir(cache_dir.to_path_buf());

    for (label, image) in cfg.images.iter() {
        if image.local_path.exists() {
            continue;
        }
        // Create parent dirs
        if let Some(parent) = image.local_path.parent() {
            tokio::fs::create_dir_all(parent).await.into_diagnostic()?;
        }

        tracing::info!(label = %label, source = %image.source, local = ?image.local_path, "downloading base image via vm-manager");

        // Use vm-manager's ImageManager for streaming download + decompression
        img_mgr
            .download(&image.source, &image.local_path)
            .await
            .into_diagnostic()?;

        tracing::info!(label = %label, local = ?image.local_path, "image ready");
    }
    Ok(())
}

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

    #[test]
    fn test_alias_resolution_and_image_lookup() {
        let yaml = r#"
            default_label: illumos-latest
            aliases:
              illumos-latest: omnios-bloody
            images:
              omnios-bloody:
                source: "https://example.com/omnios.qcow2"
                local_path: "/tmp/omnios.qcow2"
                decompress: none
                nocloud: true
                defaults:
                  cpu: 2
                  ram_mb: 2048
                  disk_gb: 20
        "#;
        let cfg: OrchestratorConfig = serde_yaml::from_str(yaml).expect("parse yaml");
        // resolve default
        assert_eq!(cfg.resolve_label(None), Some("omnios-bloody"));
        // alias mapping
        assert_eq!(
            cfg.resolve_label(Some("illumos-latest")),
            Some("omnios-bloody")
        );
        // image for canonical key
        let img = cfg.image_for("omnios-bloody").expect("image exists");
        assert!(img.nocloud);
        assert_eq!(img.defaults.as_ref().and_then(|d| d.cpu), Some(2));
        assert_eq!(img.defaults.as_ref().and_then(|d| d.ram_mb), Some(2048));
        assert_eq!(img.defaults.as_ref().and_then(|d| d.disk_gb), Some(20));
    }
}