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


#[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, download
/// from `source` 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()?;
        }
        // Download to temporary file
        let tmp_path = image.local_path.with_extension("part");
        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(())
}