launcher/src/main.rs

#![deny(clippy::wildcard_imports)]

use anyhow::Result;
use clap::Parser;

fluent_i18n::i18n!("locales", fallback = "en");

mod after_stop;
mod before_serve;
mod builtin_scripts;
mod cli;
mod cluster_state;
mod config;
mod docker;
mod doctor;
mod ends;
mod exec;
mod health;
mod home_assistant;
mod host_keys;
mod orchestrate;
mod power;
mod pre_serve;
mod ssh;
mod tmux;

use crate::cli::{Cli, Commands, PowerAction};
use crate::config::{
    Cluster, Config, default_config_dir, load_recipe, resolve_cluster, select_cluster_by_name,
};
use crate::ssh::fetch_host_key;
use std::collections::HashSet;

/// Resolve the path to the config file.
///
/// Search order:
/// 1. `--config <path>` from CLI (always wins; error if missing)
/// 2. `./config.toml` in cwd
/// 3. `<default_config_dir>/config.toml`
///
/// Returns an error if nothing is found, with a hint pointing the user
/// to copy the example config to one of the two paths.
fn resolve_config_path(cli_config: &str) -> std::path::PathBuf {
    use std::path::PathBuf;
    let explicit = PathBuf::from(cli_config);
    // User passed a non-default path: use it as-is.
    if cli_config != "config.toml" {
        return explicit;
    }
    // Default value (or no flag): try cwd first, then default folder.
    if explicit.exists() {
        return explicit;
    }
    if let Some(dir) = default_config_dir() {
        let candidate = dir.join("config.toml");
        if candidate.exists() {
            return candidate;
        }
    }
    // Fall through; the loader will report a clean error.
    explicit
}

#[tokio::main]
async fn main() -> Result<()> {
    // Locale: check LANG env var first (macOS sys-locale ignores it).
    // Normalize "ru_RU.UTF-8" → "ru-RU" for unic-langid.
    let locale = std::env::var("LANG")
        .ok()
        .map(|l| l.split('.').next().unwrap_or(&l).replace('_', "-"));
    let _ = fluent_i18n::set_locale(locale.as_deref());
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();
    let cli = Cli::parse();

    let config_path = resolve_config_path(&cli.config);
    let raw_config = tokio::fs::read_to_string(&config_path).await?;
    let config: Config = toml::from_str(&raw_config)?;
    let dry_run = cli.dry_run;

    match cli.command {
        Commands::Serve {
            recipe,
            cluster,
            node,
            keep,
            skip_before_serve,
            follow,
            auto_power_on,
            sync_image,
            sync_image_source,
            build_node,
            force_build,
            solo,
        } => {
            handle_serve(
                &config,
                &recipe,
                &cluster,
                node,
                docker::ServeOpts {
                    keep,
                    skip_before_serve,
                    dry_run,
                    auto_power_on,
                    sync_image,
                    sync_image_source,
                    build_node,
                    force_build,
                    solo,
                },
                follow,
                dry_run,
            )
            .await?;
        }
        Commands::Build {
            recipe,
            build_node,
            cluster,
            force_build,
        } => {
            handle_build_cmd(
                &config,
                &recipe,
                &build_node,
                &cluster,
                force_build,
                dry_run,
            )
            .await?;
        }
        Commands::Logs {
            recipe: recipe_name,
            cluster,
            follow,
            node,
        } => {
            handle_logs(&config, &recipe_name, &cluster, follow, node, dry_run).await?;
        }
        Commands::Stop {
            recipe: recipe_name,
            cluster,
            node,
            power_off,
        } => {
            handle_stop(&config, &recipe_name, &cluster, &node, power_off, dry_run).await?;
        }
        Commands::Health {
            cluster: cluster_name,
            fingerprint,
        } => {
            handle_health(&config, cluster_name, fingerprint).await?;
        }
        Commands::Doctor {
            recipe,
            cluster,
            repair,
        } => {
            doctor::doctor(&config, &recipe, &cluster, repair, dry_run).await?;
        }
        Commands::Tmux(tmux_cmd) => {
            crate::tmux::handle_tmux(&config, tmux_cmd, dry_run).await?;
        }
        Commands::InitHostKeys { cluster } => {
            host_keys::init_host_keys(
                &config,
                cluster.as_deref(),
                config_path.to_str().unwrap_or("config.toml"),
            )
            .await?;
        }
        Commands::Exec {
            cluster,
            recipe,
            node,
            alias,
            command,
        } => {
            handle_exec(&config, cluster, recipe, node, alias, command, dry_run).await?;
        }
        Commands::Power { action } => {
            handle_power(&config, action, dry_run).await?;
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Per-command handler functions
// ---------------------------------------------------------------------------

/// Handle `serve` command: launch inference containers on a cluster.
async fn handle_serve(
    config: &Config,
    recipe: &str,
    cluster: &str,
    node: Option<String>,
    opts: docker::ServeOpts,
    follow: bool,
    dry_run: bool,
) -> Result<()> {
    let cluster = docker::serve(config, recipe, cluster, node.as_deref(), opts).await?;
    if follow && !dry_run {
        tracing::info!("{}", fluent_i18n::t!("container-log-header"));
        docker::show_logs(config, &cluster, recipe, true, node.as_deref(), dry_run).await?;
    }
    Ok(())
}

/// Handle `build` command: build a Docker image on a remote node.
async fn handle_build_cmd(
    config: &Config,
    recipe: &str,
    build_node: &str,
    cluster: &str,
    force_build: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe_obj = crate::config::load_recipe(recipe).await?;
    let cluster = resolve_cluster_for_op(config, recipe, cluster).await?;
    let opts = docker::ServeOpts {
        keep: false,
        skip_before_serve: true,
        dry_run,
        auto_power_on: false,
        sync_image: false,
        sync_image_source: None,
        build_node: Some(build_node.to_string()),
        force_build,
        solo: None,
    };
    docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
    Ok(())
}

/// Handle `logs` command: show container logs.
async fn handle_logs(
    config: &Config,
    recipe_name: &str,
    cluster: &str,
    follow: bool,
    node: Option<String>,
    dry_run: bool,
) -> Result<()> {
    let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
    docker::show_logs(
        config,
        &cluster,
        recipe_name,
        follow,
        node.as_deref(),
        dry_run,
    )
    .await?;
    Ok(())
}

/// Handle `stop` command: stop containers and optionally power off.
async fn handle_stop(
    config: &Config,
    recipe_name: &str,
    cluster: &str,
    node: &[String],
    power_off: bool,
    dry_run: bool,
) -> Result<()> {
    let (stopped_cluster, stopped_nodes): (Cluster, Vec<String>) = if !node.is_empty() {
        let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
        let nodes = node.to_vec();
        docker::stop_cluster(config, &cluster, recipe_name, &nodes, dry_run).await?;
        (cluster, nodes)
    } else {
        let variants = docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
        if variants.is_empty() {
            anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
        }
        if variants.len() == 1 {
            let v = &variants[0];
            let cluster = select_cluster_by_name(config, &v.cluster)?;
            let nodes = v.nodes.clone();
            docker::stop_cluster(config, &cluster, recipe_name, &nodes, dry_run).await?;
            (cluster, nodes)
        } else {
            let mut msg = format!(
                "Multiple deployments of '{}' found — use --node to select:\n",
                recipe_name
            );
            for (i, v) in variants.iter().enumerate() {
                let nodes = v
                    .nodes
                    .iter()
                    .map(|n| format!("--node {}", n))
                    .collect::<Vec<_>>()
                    .join(" ");
                msg.push_str(&format!(
                    "  {}: cluster '{}' with {}\n",
                    i + 1,
                    v.cluster,
                    nodes
                ));
            }
            anyhow::bail!(msg);
        }
    };

    // Run recipe `after_stop` hooks (e.g. restore_dgx_services) on the
    // nodes that were actually stopped. Best-effort against the loaded
    // recipe; if the recipe file is unavailable, skip silently.
    if let Ok(recipe) = crate::config::load_recipe(recipe_name).await {
        crate::after_stop::run(config, &stopped_cluster, &recipe, dry_run).await?;
    }

    if power_off {
        crate::power::power_off(config, &stopped_cluster, &stopped_nodes, dry_run).await?;
    }
    Ok(())
}

/// Handle `health` command: check cluster health or show host key fingerprints.
async fn handle_health(
    config: &Config,
    cluster_name: Option<String>,
    fingerprint: bool,
) -> Result<()> {
    if fingerprint {
        let clusters_to_check: Vec<Cluster> = match &cluster_name {
            Some(name) => vec![select_cluster_by_name(config, name)?],
            None => config.clusters.clone(),
        };
        for cluster in &clusters_to_check {
            println!(
                "\n{}",
                fluent_i18n::t!("cluster-section", { "name" => &cluster.name })
            );
            for node in &cluster.nodes {
                match fetch_host_key(config, node).await {
                    Ok(fp) => println!(
                        "{}",
                        fluent_i18n::t!("host-key-line", { "name" => &node.name, "ip" => &node.ssh_ip, "key" => &fp })
                    ),
                    Err(e) => println!(
                        "{}",
                        fluent_i18n::t!("host-key-error", { "name" => &node.name, "ip" => &node.ssh_ip, "error" => &e.to_string() })
                    ),
                }
            }
        }
    } else {
        match cluster_name {
            Some(ref name) => {
                let cluster = select_cluster_by_name(config, name)?;
                docker::health::health_check(config, &cluster, false).await?;
            }
            None => {
                for cluster in &config.clusters {
                    docker::health::health_check(config, cluster, false).await?;
                }
                docker::health::health_summary(config).await?;
            }
        }
    }
    Ok(())
}

/// Handle `exec` command: run a command on cluster nodes via SSH.
async fn handle_exec(
    config: &Config,
    cluster: Option<String>,
    recipe: Option<String>,
    node: Option<String>,
    alias: Option<String>,
    command: Vec<String>,
    dry_run: bool,
) -> Result<()> {
    let (cluster, _recipe) =
        exec::resolve_target(config, cluster.as_deref(), recipe.as_deref()).await?;
    let cmd_str = exec::resolve_command(config, alias.as_deref(), &command)?;
    exec::run(config, &cluster, node.as_deref(), &cmd_str, dry_run).await?;
    Ok(())
}

/// Handle `power` command: power on/off/switch nodes.
async fn handle_power(config: &Config, action: PowerAction, dry_run: bool) -> Result<()> {
    match action {
        PowerAction::Off { cluster, node } => {
            let cluster = select_cluster_by_name(config, &cluster)?;
            let nodes = node.into_iter().collect::<Vec<_>>();
            power::power_off(config, &cluster, &nodes, dry_run).await?;
        }
        PowerAction::On { cluster, node } => {
            let cluster = select_cluster_by_name(config, &cluster)?;
            let nodes = node.into_iter().collect::<Vec<_>>();
            power::power_on(config, &cluster, &nodes, dry_run).await?;
        }
        PowerAction::Switch { from, to } => {
            let from = select_cluster_by_name(config, &from)?;
            let to = select_cluster_by_name(config, &to)?;
            power::power_switch(config, &from, &to, dry_run).await?;
        }
    }
    Ok(())
}

/// Resolve a cluster for container operations (stop/logs).
/// When `auto`, tries live container label lookup first, then falls back
/// to config-based auto-selection (for dry-run / offline scenarios).
async fn resolve_cluster_for_op(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
) -> Result<Cluster> {
    if cluster_name != "auto" {
        let recipe = load_recipe(recipe_name).await?;
        let (cluster, _) = resolve_cluster(config, &recipe, cluster_name, &HashSet::new())?;
        return Ok(cluster);
    }
    // Try live container label lookup first
    if let Ok(c) = docker::resolve_cluster_by_container_label(config, recipe_name).await {
        return Ok(c.clone());
    }
    // Fall back to config-based auto (dry-run / containers not running)
    let recipe = load_recipe(recipe_name).await?;
    let (cluster, _) = resolve_cluster(config, &recipe, "auto", &HashSet::new())?;
    Ok(cluster)
}