launcher/src/cli.rs

use clap::{Args, Parser, Subcommand, ValueEnum};

/// Output format for human-readable reports (`--format`).
#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
pub enum OutputFormat {
    Plain,
    Md,
    Html,
}

#[derive(Parser)]
#[command(author, version, about = "DGX Spark Launcher for vLLM and SGLang")]
pub struct Cli {
    #[arg(long, short, default_value = "config.toml")]
    pub config: String,

    #[arg(long, short, help = "Show commands without executing")]
    pub dry_run: bool,

    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Launch a cluster
    Serve {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,

        /// Keep containers after stop (default: auto-remove on stop)
        #[arg(long)]
        keep: bool,

        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,

        /// Don't follow logs after serving (following is the default).
        #[arg(long, action = clap::ArgAction::SetTrue)]
        no_follow: bool,

        /// Power on cluster if switch is off (requires HA config).
        #[arg(short = 'O', long = "on")]
        auto_power_on: bool,

        /// Sync Docker image across nodes before starting.
        #[arg(long, default_value_t = true)]
        sync_image: bool,

        /// Explicit source node for image sync (default: auto-detect).
        #[arg(long)]
        sync_image_source: Option<String>,

        /// Build image on this node if not found (requires docker-images/<recipe>/).
        #[arg(long)]
        build_node: Option<String>,

        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,

        /// Run even if nodes are occupied (another container running).
        /// Skips the occupancy check during cluster selection — use to run
        /// a lightweight test recipe alongside a live serve, or to force
        /// onto a busy node. Dangerous for GPU recipes (OOM, port
        /// conflicts); the operator is responsible for coexistence.
        #[arg(long)]
        force: bool,

        /// Serve single-node on the first free node found (overrides recipe
        /// `solo`). Use `--no-solo` to force multi-node even if the recipe
        /// sets `solo = true`.
        #[arg(long)]
        solo: Option<bool>,

        /// Mirror the model into a per-node local folder and mount that
        /// instead of the shared HF cache (first serve copies the model from
        /// the shared cache; later serves load it locally). Without a value
        /// uses ~/.cache/hf-models; `--mirror <path>` uses that exact path.
        #[arg(long, num_args = 0..=1, default_missing_value = "")]
        mirror: Option<String>,

        /// Scale the recipe to a node count and let tensor-parallel track it
        /// (overrides `required_nodes` and `tensor_parallel_size`).
        /// `--fit=N` scales to N nodes; bare `--fit` uses the whole selected
        /// cluster. Incompatible with `--solo`. `pipeline_parallel_size` is
        /// left untouched.
        #[arg(long, num_args = 0..=1, default_missing_value = "")]
        fit: Option<String>,

        /// Clear the host-side `extra_mounts` folders (alias/host) on the
        /// launch nodes with `sudo rm -rf` before containers start — the
        /// containers run as root, so caches there are usually root-owned.
        /// Repo-based mounts (HF cache path refs) are never cleared. `-d`
        /// prints the exact command instead of running it.
        #[arg(long)]
        clear: bool,
    },

    /// Show container logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Target a specific node within the cluster (default: head node if following, all nodes if not).
        #[arg(long)]
        node: Option<String>,
    },

    /// Stop a cluster
    Stop {
        #[arg(long, required = true)]
        recipe: String,

        /// Target specific nodes (repeatable, e.g. --node dgx1 --node dgx2).
        #[arg(long)]
        node: Vec<String>,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Power off cluster after stopping containers (requires HA config).
        #[arg(long = "off")]
        power_off: bool,
    },

    /// Check cluster health
    Health {
        #[arg(
            long,
            help = "Cluster name (e.g. 4nodes-switch). If omitted — checks all."
        )]
        cluster: Option<String>,
        #[arg(long, help = "Show node host keys (for adding to config.toml)")]
        fingerprint: bool,
        /// Output format (cluster-health mode only; ignored with --fingerprint).
        #[arg(long, value_enum, default_value_t = OutputFormat::Plain)]
        format: OutputFormat,
    },

    /// Check and repair a recipe (doctor mode)
    Doctor {
        #[arg(long, required = true, help = "Recipe name to check")]
        recipe: String,
        #[arg(
            long,
            default_value = "auto",
            help = "Cluster (if not auto — uses recipe.cluster)"
        )]
        cluster: String,
        #[arg(long, help = "Auto-fix issues (create Dockerfile + build image)")]
        repair: bool,
    },

    /// Publish a recipe (+ image overlay) as a standalone reproducible bundle
    Publish(PublishArgs),

    /// Browse and install recipes from the public registry
    #[command(subcommand)]
    Registry(RegistryCommands),

    /// Manage tmux sessions (Docker alternative)
    #[command(subcommand)]
    Tmux(TmuxCommands),

    /// Auto-discover host keys and write them into config.toml
    InitHostKeys {
        #[arg(long, help = "Cluster name (e.g. 4nodes-switch). If omitted — all.")]
        cluster: Option<String>,
    },

    /// Run an arbitrary shell command on a cluster or single node
    Exec {
        /// Cluster name (mutually exclusive with --recipe).
        #[arg(long, conflicts_with = "recipe")]
        cluster: Option<String>,

        /// Recipe name; resolves to the recipe's pinned cluster.
        #[arg(long, conflicts_with = "cluster")]
        recipe: Option<String>,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,

        /// Name of an alias from `[exec_aliases]` in config.toml.
        /// Mutually exclusive with a trailing command.
        #[arg(long)]
        alias: Option<String>,

        /// Command to run (everything after `--`).
        #[arg(trailing_var_arg = true)]
        command: Vec<String>,
    },

    /// Manage configuration (global config path, init)
    #[command(subcommand)]
    Config(ConfigCommands),

    /// Launch the web dashboard (axum + SSE) on a local port
    Web {
        #[arg(long, default_value_t = 3100)]
        port: u16,

        #[arg(long, default_value = "127.0.0.1")]
        bind: String,

        /// Bearer token for data endpoints (required when --bind is not loopback)
        #[arg(long)]
        token: Option<String>,

        /// Open the browser once the server is up
        #[arg(long)]
        open: bool,
    },

    /// Build Docker image for a recipe on a remote node
    Build {
        #[arg(long, required = true, help = "Recipe name")]
        recipe: String,

        #[arg(long, required = true, help = "Node to build on (e.g. dgx1.lan)")]
        build_node: String,

        #[arg(long, default_value = "auto", help = "Cluster name")]
        cluster: String,

        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,
    },

    /// Power on/off nodes via Home Assistant (smart plugs / Wake-on-LAN).
    Power {
        #[command(subcommand)]
        action: PowerAction,
    },
}

#[derive(Subcommand)]
pub enum PowerAction {
    /// Power off: SSH `sudo shutdown now` + wait + HA `turn_off`.
    Off {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Power on: HA `turn_on` only (no SSH; the smart plug wakes the machine).
    On {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Switch clusters: disable nodes from source, enable nodes in target.
    Switch {
        /// Current cluster (nodes + entities to shut down).
        #[arg(long)]
        from: String,
        /// Target cluster (nodes + entities to wake up).
        #[arg(long)]
        to: String,
    },
}

/// `sm121 config` subcommands — manage the global config file.
#[derive(Subcommand)]
pub enum ConfigCommands {
    /// Show the active config path and where the global config lives
    Path,
    /// Create the global config file (template) if it does not exist
    InitGlobal,
}

#[derive(Args)]
pub struct PublishArgs {
    #[arg(long, required = true)]
    pub recipe: String,

    /// Cluster whose fabric env + node count fill the generated run script
    #[arg(long)]
    pub cluster: Option<String>,

    /// Output dir (default: ./publish/<recipe>)
    #[arg(long)]
    pub out: Option<std::path::PathBuf>,

    /// Also create a public GitHub gist (requires `gh` CLI, authenticated)
    #[arg(long)]
    pub gist: bool,

    /// Upload the bundle to a recipe registry (see `sm121 registry --help`)
    #[arg(long, value_name = "URL")]
    pub registry: Option<String>,

    /// Registry publish token (or $SM121RECIPES_TOKEN)
    #[arg(long)]
    pub token: Option<String>,
}

#[derive(Subcommand)]
pub enum RegistryCommands {
    /// List recipes on the registry
    List {
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
    },
    /// Print a registry recipe (recipe.toml)
    Show {
        name: String,
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
    },
    /// Download a recipe into ./recipes/<name>/ (+ ./docker-images/<name>/ if it has an image)
    Install {
        name: String,
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
        /// Overwrite an existing ./recipes/<name>/
        #[arg(long)]
        force: bool,
    },
}

#[derive(Subcommand)]
pub enum TmuxCommands {
    /// Launch cluster via tmux
    Serve {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,
    },
    /// Stop tmux sessions
    Stop {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Detailed tmux session status
    Status {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Show tmux session logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Attach to head node tmux session
    Attach {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Parse `sm121 health <args>` and return the parsed `--format`.
    fn health_format(args: &[&str]) -> Result<(OutputFormat, bool), clap::Error> {
        let cli = Cli::try_parse_from(
            std::iter::once("sm121")
                .chain(std::iter::once("health"))
                .chain(args.iter().copied()),
        )?;
        let Commands::Health {
            fingerprint,
            format,
            ..
        } = cli.command
        else {
            panic!("expected Health")
        };
        Ok((format, fingerprint))
    }

    #[test]
    fn health_format_defaults_to_plain() {
        assert_eq!(health_format(&[]).unwrap(), (OutputFormat::Plain, false));
    }

    #[test]
    fn health_format_parses_each_variant() {
        assert_eq!(
            health_format(&["--format", "plain"]).unwrap().0,
            OutputFormat::Plain
        );
        assert_eq!(
            health_format(&["--format", "md"]).unwrap().0,
            OutputFormat::Md
        );
        assert_eq!(
            health_format(&["--format", "html"]).unwrap().0,
            OutputFormat::Html
        );
    }

    #[test]
    fn health_format_combines_with_cluster_and_fingerprint() {
        let (fmt, fp) = health_format(&["--cluster", "c1", "--format", "md"]).unwrap();
        assert_eq!((fmt, fp), (OutputFormat::Md, false));
        let (fmt, fp) = health_format(&["--fingerprint", "--format", "html"]).unwrap();
        assert_eq!((fmt, fp), (OutputFormat::Html, true));
    }

    #[test]
    fn health_format_rejects_invalid_value() {
        assert!(health_format(&["--format", "bogus"]).is_err());
    }
}