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` (supports http(s):// and file://) and optionally decompress
/// according to `decompress`.
pub async fn ensure_images(cfg: &OrchestratorConfig) -> Result<()> {
    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()?;
        }

        let source = image.source.as_str();
        let is_file = source.starts_with("file://");
        let tmp_path = image.local_path.with_extension("part");

        if is_file {
            // Local file source: copy or decompress from local path
            let src_path = PathBuf::from(&source[7..]); // naive parse; paths should be absolute
            tracing::info!(label = %label, src = ?src_path, local = ?image.local_path, "preparing base image from file:// source");

            match image.decompress.unwrap_or(Decompress::None) {
                Decompress::None => {
                    // Copy to temporary then atomically move into place
                    tokio::fs::copy(&src_path, &tmp_path)
                        .await
                        .into_diagnostic()?;
                    tokio::fs::rename(&tmp_path, &image.local_path)
                        .await
                        .into_diagnostic()?;
                }
                Decompress::Zstd => {
                    let src = src_path.clone();
                    let tmp_out = tmp_path.clone();
                    task::spawn_blocking(move || -> miette::Result<()> {
                        let infile = fs::File::open(&src).into_diagnostic()?;
                        let mut decoder =
                            zstd::stream::read::Decoder::new(infile).into_diagnostic()?;
                        let mut outfile = fs::File::create(&tmp_out).into_diagnostic()?;
                        std::io::copy(&mut decoder, &mut outfile).into_diagnostic()?;
                        Ok(())
                    })
                    .await
                    .into_diagnostic()??;
                    tokio::fs::rename(&tmp_path, &image.local_path)
                        .await
                        .into_diagnostic()?;
                }
            }
        } else {
            // Remote URL (HTTP/HTTPS): download to temporary file first
            tracing::info!(label = %label, url = %image.source, local = ?image.local_path, "downloading base image");
            let resp = reqwest::get(&image.source).await.into_diagnostic()?;
            let status = resp.status();
            if !status.is_success() {
                miette::bail!(
                    "failed to download {url}: {status}",
                    url = image.source,
                    status = status
                );
            }
            let bytes = resp.bytes().await.into_diagnostic()?;
            tokio::fs::write(&tmp_path, &bytes)
                .await
                .into_diagnostic()?;

            // Decompress or move into place
            match image.decompress.unwrap_or(Decompress::None) {
                Decompress::None => {
                    tokio::fs::rename(&tmp_path, &image.local_path)
                        .await
                        .into_diagnostic()?;
                }
                Decompress::Zstd => {
                    let src = tmp_path.clone();
                    let dst = image.local_path.clone();
                    task::spawn_blocking(move || -> miette::Result<()> {
                        let infile = fs::File::open(&src).into_diagnostic()?;
                        let mut decoder =
                            zstd::stream::read::Decoder::new(infile).into_diagnostic()?;
                        let mut outfile = fs::File::create(&dst).into_diagnostic()?;
                        std::io::copy(&mut decoder, &mut outfile).into_diagnostic()?;
                        // remove compressed temp
                        std::fs::remove_file(&src).ok();
                        Ok(())
                    })
                    .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));
    }
}