launcher/src/config.rs

use anyhow::{Context, Result};
use std::collections::HashSet;
use tracing::info;

mod types;
pub use types::*;

pub fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
    if let Some(c) = clusters.iter().find(|c| c.name == name) {
        return Ok(c);
    }
    let matches: Vec<&Cluster> = clusters
        .iter()
        .filter(|c| c.name.starts_with(name))
        .collect();
    match matches.len() {
        0 => anyhow::bail!(fluent_i18n::t!("cluster-not-found", { "name" => name.to_string() })),
        1 => Ok(matches[0]),
        n => anyhow::bail!(
            "Pref '{}' matches {} clusters: {}",
            name,
            n,
            matches
                .iter()
                .map(|c| c.name.as_str())
                .collect::<Vec<_>>()
                .join(", ")
        ),
    }
}

pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
    required_nodes: Option<u32>,
    occupied: &HashSet<String>,
) -> Result<(Cluster, &'c NcclConfig)> {
    if cluster_name != "auto" {
        let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
        let nccl = config
            .nccl
            .get(&cluster.name)
            .context(fluent_i18n::t!("nccl-config-missing"))?;
        return Ok((cluster.clone(), nccl));
    }
    let free = |c: &Cluster| -> u32 {
        c.nodes
            .iter()
            .filter(|n| !occupied.contains(&n.name))
            .count() as u32
    };
    let mut candidates: Vec<(&Cluster, &NcclConfig)> = config
        .clusters
        .iter()
        .filter_map(|c| Some((c, config.nccl.get(&c.name)?)))
        .filter(|(c, _)| required_nodes.is_none_or(|r| free(c) >= r))
        .collect();
    candidates.sort_by_key(|(c, _)| c.nodes.len());
    if let Some((cluster, nccl)) = candidates.into_iter().next() {
        let avail = free(cluster);
        info!(
            "{}",
            fluent_i18n::t!("cluster-selected", { "name" => cluster.name.to_string(), "nodes" => avail })
        );
        return Ok((cluster.clone(), nccl));
    }
    anyhow::bail!(fluent_i18n::t!("cluster-auto-fail"));
}

pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
    Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
}

pub fn resolve_cluster<'c>(
    config: &'c Config,
    recipe: &Recipe,
    cluster_name: &str,
    occupied: &HashSet<String>,
) -> Result<(Cluster, &'c NcclConfig)> {
    let choice = if cluster_name != "auto" {
        cluster_name.to_string()
    } else if let Some(ref rc) = recipe.cluster {
        rc.clone()
    } else {
        "auto".to_string()
    };
    select_cluster(config, &choice, recipe.required_nodes, occupied)
}

pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
    let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods"));
    if cwd.is_dir() {
        return Some(cwd);
    }
    let cfg = default_config_dir()?
        .join("recipes")
        .join(name)
        .join("mods");
    cfg.is_dir().then_some(cfg)
}

pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
    let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
    if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
        return Ok(toml::from_str(&content)?);
    }
    if let Some(dir) = default_config_dir() {
        let in_default = dir.join("recipes").join(recipe_name).join("recipe.toml");
        if let Ok(content) = tokio::fs::read_to_string(&in_default).await {
            return Ok(toml::from_str(&content)?);
        }
    }
    let content = tokio::fs::read_to_string(&cwd_path).await?;
    Ok(toml::from_str(&content)?)
}

pub fn default_config_dir() -> Option<std::path::PathBuf> {
    dirs::config_dir().map(|p| p.join("sm121"))
}

#[cfg(test)]
mod tests;