launcher/src/main.rs

#![deny(clippy::wildcard_imports)]

use crate::docker::sync;
use anyhow::Result;
use clap::Parser;

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

mod after_stop;
mod ansi;
mod before_serve;
mod bench;
mod bench_artifacts;
mod bench_depth;
mod bench_report;
mod builtin_scripts;
mod cli;
mod cluster_state;
mod config;
mod detect;
mod docker;
mod doctor;
mod ends;
mod engine_flags;
mod exec;
mod gconfig;
mod health;
mod home_assistant;
mod host_keys;
mod locale;
mod mikrotik;
mod orchestrate;
mod power;
mod pre_serve;
mod publish;
mod publish_render;
mod readiness;
mod registry;
mod report;
mod runlock;
mod ssh;
mod staging;
mod stamp;
mod tmux;
mod web;
use crate::cli::{Cli, Commands, OutputFormat, PowerAction};
use crate::config::{
    Cluster, Config, DeployTarget, load_recipe, resolve_cluster, select_cluster_by_name,
};
use crate::ssh::fetch_host_key;
use std::collections::HashSet;

#[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());
    locale::set_global_locale(locale.clone()); // tokio workers miss the thread-local; spawn_with_locale re-applies
    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, raw_config) = crate::gconfig::read_config(&cli.config).await?;
    let mut config: Config = toml::from_str(&raw_config)?;
    crate::config::merge_cluster_sources(&mut config).await?;
    let dry_run = cli.dry_run;
    // `config` subcommands are pure path/file utilities: dispatch them
    // before cluster_sources merging so they work even with broken sources.
    if let Commands::Config(cmd) = &cli.command {
        return crate::gconfig::run(cmd, &cli.config, dry_run);
    }

    match cli.command {
        Commands::Serve {
            recipe,
            cluster,
            node,
            keep,
            skip_before_serve,
            no_follow,
            no_wait,
            auto_power_on,
            sync_image,
            sync_image_source,
            build_node,
            force_build,
            force,
            solo,
            mirror,
            hf_cache,
            stage,
            fit,
            clear,
            no_push,
        } => {
            handle_serve(
                &config,
                &recipe,
                &cluster,
                node,
                docker::ServeOpts {
                    keep,
                    skip_before_serve,
                    no_wait_ready: no_wait,
                    dry_run,
                    auto_power_on,
                    sync_image,
                    sync_image_source,
                    no_push,
                    build_node,
                    force_build,
                    force,
                    solo,
                    mirror,
                    hf_cache,
                    stage,
                    fit,
                    clear,
                },
                !no_follow,
                dry_run,
            )
            .await?;
        }
        Commands::SyncImage(a) => {
            sync::handle_sync_image(&config, &a.recipe, &a.cluster, a.source.as_deref(), dry_run)
                .await?;
        }
        Commands::Build {
            recipe,
            build_node,
            cluster,
            force_build,
            no_push,
        } => {
            docker::build::handle_build_cmd(
                &config,
                &recipe,
                &build_node,
                &cluster,
                force_build,
                no_push,
                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,
            no_wait,
        } => {
            crate::docker::stop::handle(
                &config,
                &recipe_name,
                &cluster,
                &node,
                !no_wait,
                power_off,
                dry_run,
            )
            .await?;
        }
        Commands::Bench {
            recipe,
            cluster,
            label,
            concurrency,
            seed,
            parallel,
            perf_only,
            tokenizer,
        } => {
            let args = crate::bench::Args {
                recipe,
                cluster,
                label,
                concurrency,
                seed,
                parallel,
                perf_only,
                tokenizer,
            };
            crate::bench::run(&config, &args, dry_run).await?;
        }
        Commands::Health {
            cluster: cluster_name,
            fingerprint,
            format,
        } => {
            handle_health(&config, cluster_name, fingerprint, format).await?;
        }
        Commands::Doctor {
            recipe,
            cluster,
            repair,
        } => {
            doctor::doctor(&config, &recipe, &cluster, repair, dry_run).await?;
        }
        Commands::Publish(a) => {
            publish::publish(
                &config,
                &a.recipe,
                a.cluster.as_deref(),
                a.out,
                a.gist,
                &publish::RecipeRegistryOpts {
                    url: a.registry.clone(),
                    token: a.token.clone(),
                },
                publish::PushOpts {
                    push: a.push,
                    node: a.node.clone(),
                    force: a.force,
                },
            )
            .await?;
        }
        Commands::Registry(cmd) => registry::handle(cmd).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::Config(_) => unreachable!("dispatched above, before merging"),
        Commands::Web {
            port,
            bind,
            token,
            open,
        } => {
            crate::web::run(config, port, bind, token, open).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?;
        }
        Commands::Report {
            serve,
            port,
            open,
            archive,
            unarchive,
            archived,
        } => {
            crate::report::run(
                &config,
                serve || open,
                port,
                open,
                archive,
                unarchive,
                archived,
            )
            .await?;
        }
        Commands::Runlock(cmd) => crate::runlock::run(&config, &cmd)?,
    }

    Ok(())
}

/// Handle `serve` command: launch inference containers on a cluster.
///
/// Dispatches to `docker::serve` or `tmux::serve` based on the recipe's
/// `deploy` field. Docker-only CLI opts (sync_image, build_node,
/// force_build, keep, auto_power_on) are ignored with a warning when
/// `deploy = "tmux"`.
pub(crate) async fn handle_serve(
    config: &Config,
    recipe: &str,
    cluster: &str,
    node: Option<String>,
    opts: docker::ServeOpts,
    follow: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe_obj = load_recipe(recipe, &config.recipe_sources).await?;
    match recipe_obj.deploy {
        DeployTarget::Docker => {
            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::logs::follow_with_liveness(
                    config,
                    &cluster,
                    recipe,
                    recipe_obj.port,
                    node.as_deref(),
                    dry_run,
                )
                .await?;
            }
        }
        DeployTarget::Tmux => {
            // Warn about docker-only CLI opts that are silently ignored.
            let ignored: Vec<&str> = [
                (opts.sync_image, "--sync-image"),
                (opts.sync_image_source.is_some(), "--sync-image-source"),
                (opts.build_node.is_some(), "--build-node"),
                (opts.force_build, "--force-build"),
                (opts.keep, "--keep"),
                (opts.solo.is_some(), "--solo/--no-solo"),
                (opts.mirror.is_some(), "--mirror"),
                (opts.fit.is_some(), "--fit"),
            ]
            .iter()
            .filter_map(|(set, flag)| if *set { Some(*flag) } else { None })
            .collect();
            if !ignored.is_empty() {
                tracing::warn!(
                    "{}",
                    fluent_i18n::t!("tmux-deploy-ignored-opts", { "opts" => &ignored.join(", ") })
                );
            }
            crate::tmux::serve::tmux_serve(
                config,
                recipe,
                cluster,
                opts.skip_before_serve,
                dry_run,
                opts.auto_power_on,
                opts.force,
            )
            .await?;
            if !dry_run {
                // no readiness probe for tmux; cluster stays None (unresolved here)
                crate::runlock::mark_start(recipe, &config.recipe_sources, None)?;
            }
            if follow && !dry_run {
                tracing::info!(
                    "{}",
                    fluent_i18n::t!("tmux-following-logs", { "name" => recipe.to_string() })
                );
                crate::tmux::ops::tmux_logs(config, recipe, cluster, true, dry_run).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 `health` command: check cluster health or show host key fingerprints.
async fn handle_health(
    config: &Config,
    cluster_name: Option<String>,
    fingerprint: bool,
    format: OutputFormat,
) -> 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, format).await?;
            }
            None => {
                let fleet = docker::health::collect(config, None, false).await?;
                docker::health::print_fleet(&fleet, format);
            }
        }
    }
    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).
pub(crate) 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, &config.recipe_sources).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, &config.recipe_sources).await?;
    let (cluster, _) = resolve_cluster(config, &recipe, "auto", &HashSet::new())?;
    Ok(cluster)
}