launcher/src/tmux/serve.rs

use anyhow::{Context, Result};
use std::collections::HashSet;
use std::time::Duration;
use tokio::time::sleep;
use tracing::{info, warn};

use crate::config::{Config, Node, Recipe};
use crate::tmux::TmuxManager;

/// Launches a cluster via tmux sessions.
/// For Ray backend: starts head node first, waits for readiness, then workers.
pub async fn tmux_serve(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    skip_before_serve: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
    let recipe: Recipe = toml::from_str(&tokio::fs::read_to_string(&recipe_path).await?)?;

    // Feature gaps: tmux deployment doesn't support Docker-only features.
    if !recipe.mods.commands.is_empty() {
        anyhow::bail!(
            "{}",
            fluent_i18n::t!("tmux-mods-unsupported", { "name" => &recipe.name })
        );
    }
    if recipe.composed_model.is_some() {
        anyhow::bail!(
            "{}",
            fluent_i18n::t!("tmux-composed-model-unsupported", { "name" => &recipe.name })
        );
    }

    info!(
        "🚀 TMUX launch: {} ({}) [backend: {}] [engine: {}]",
        recipe.name,
        recipe_name,
        recipe
            .distributed_backend
            .map_or("none".into(), |b| b.to_string()),
        recipe.engine
    );

    let choice = if cluster_name != "auto" {
        cluster_name.to_string()
    } else if let Some(ref rc) = recipe.cluster {
        rc.clone()
    } else {
        "auto".to_string()
    };
    let (cluster, nccl) =
        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;

    if let Some(required) = recipe.required_nodes {
        let actual = cluster.nodes.len() as u32;
        if actual != required {
            anyhow::bail!(
                "Recipe '{}' requires {} nodes, but cluster '{}' has {} nodes",
                recipe.name,
                required,
                cluster.name,
                actual
            );
        }
    }

    crate::before_serve::run(config, &cluster, &recipe, skip_before_serve, dry_run).await?;

    // Prepare the Python venv on each node if it doesn't exist yet.
    prepare_venv(config, &recipe, &cluster, recipe_name, dry_run).await?;

    let manager = TmuxManager;

    let is_ray = recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray);

    if is_ray {
        info!("{}", fluent_i18n::t!("tmux-ray-backend"));

        let head_node = cluster
            .nodes
            .iter()
            .find(|n| n.rank == 0)
            .context("No head node (rank 0) in cluster")?;

        let head_cmds = manager.build_session_commands(head_node, &recipe, nccl, &cluster);
        manager
            .create_session(config, head_node, &head_cmds, dry_run)
            .await?;

        if !dry_run {
            info!("{}", fluent_i18n::t!("tmux-waiting-head"));
            let sess = manager.session_name(&recipe, 0);
            let head_ready = wait_for_head_ready(config, head_node, &sess, &recipe).await?;
            if head_ready {
                info!("{}", fluent_i18n::t!("tmux-head-ready"));
            } else {
                warn!("⚠️ Head node may not be fully ready — starting workers anyway.");
            }
        }

        for node in &cluster.nodes {
            if node.rank == 0 {
                continue;
            }
            let worker_cmds = manager.build_session_commands(node, &recipe, nccl, &cluster);
            manager
                .create_session(config, node, &worker_cmds, dry_run)
                .await?;
        }
    } else {
        info!("{}", fluent_i18n::t!("tmux-mp-backend"));

        let mut handles = Vec::new();
        for node in &cluster.nodes {
            let config = config.clone();
            let node = node.clone();
            let recipe = recipe.clone();
            let nccl = nccl.clone();
            let cluster = cluster.clone();
            let manager = manager.clone();
            let dry = dry_run;

            handles.push(crate::locale::spawn_with_locale(async move {
                let cmds = manager.build_session_commands(&node, &recipe, &nccl, &cluster);
                manager
                    .create_session(&config, &node, &cmds, dry)
                    .await
                    .map(|_| node.name.clone())
            }));
        }

        let results = futures::future::join_all(handles).await;
        let mut failures = Vec::new();
        for result in results {
            match result {
                Ok(Ok(name)) => info!(
                    "{}",
                    fluent_i18n::t!("tmux-session-started", { "name" => &name })
                ),
                Ok(Err(e)) => failures.push(e),
                Err(e) => failures.push(anyhow::anyhow!("Task join: {}", e)),
            }
        }

        if !failures.is_empty() {
            for e in &failures {
                warn!("  ✗ {}", e);
            }
            anyhow::bail!("Failed to start sessions on {} node(s)", failures.len());
        }
    }

    if !dry_run {
        sleep(Duration::from_secs(5)).await;
        info!(
            "🎉 Cluster launched via tmux! API → http://{}:{}",
            cluster.master_ib_ip, recipe.port
        );
        info!(
            "📌 To view sessions: `cargo run tmux status --recipe {}`",
            recipe_name
        );
        info!(
            "📌 To attach: `cargo run tmux attach --recipe {}`",
            recipe_name
        );
    }

    Ok(())
}

/// Polls the head node until the API endpoint responds or timeout.
async fn wait_for_head_ready(
    config: &Config,
    head_node: &crate::config::Node,
    session_name: &str,
    recipe: &Recipe,
) -> Result<bool> {
    let max_attempts = 30;
    for i in 0..max_attempts {
        let pid_cmd = format!(
            "tmux list-panes -t {}:0 -F '#{{pane_pid}}' | head -1",
            session_name
        );
        match crate::ssh::run_ssh(config, head_node, pid_cmd, false).await {
            Ok(pid_str) => {
                let pid = pid_str.trim();
                if pid.is_empty() || pid == "0" {
                    sleep(Duration::from_secs(5)).await;
                    continue;
                }
            }
            Err(_) => {
                sleep(Duration::from_secs(5)).await;
                continue;
            }
        }

        let api_check = format!(
            "curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{}/v1/models 2>/dev/null || echo 'not_ready'",
            recipe.port
        );
        if let Ok(out) = crate::ssh::run_ssh(config, head_node, api_check, false).await {
            let code = out.trim();
            if code == "200" {
                info!(
                    "{}",
                    fluent_i18n::t!("tmux-head-api-ready", { "attempt" => i + 1 })
                );
                return Ok(true);
            }
        }

        if (i + 1) % 5 == 0 {
            info!(
                "  Still waiting for head node... ({}/{})",
                i + 1,
                max_attempts
            );
        }
        sleep(Duration::from_secs(5)).await;
    }

    warn!("  Head node did not become ready within timeout");
    Ok(false)
}

/// Prepares the Python venv on each node if it doesn't exist.
///
/// Checks `test -f <venv>/bin/activate` on each node; if absent, runs
/// the setup commands in order, then re-checks. The setup commands come
/// from either `recipe.venv_setup_script` (a script file read from the
/// recipe directory) or `recipe.venv_setup` (inline list or multi-line
/// string). A failure on any command aborts the serve. Nodes are
/// prepared in parallel.
async fn prepare_venv(
    config: &Config,
    recipe: &Recipe,
    cluster: &crate::config::Cluster,
    recipe_name: &str,
    dry_run: bool,
) -> Result<()> {
    let venv = match &recipe.venv {
        Some(v) => v,
        None => return Ok(()),
    };

    // Resolve setup commands: script file takes precedence over inline.
    let setup: Vec<String> = if let Some(script_path) = &recipe.venv_setup_script {
        let full_path = format!("recipes/{}/{}", recipe_name, script_path);
        if dry_run {
            info!("  read venv_setup_script: {} (dry-run)", full_path);
        } else {
            info!("  read venv_setup_script: {}", full_path);
        }
        let body = tokio::fs::read_to_string(&full_path)
            .await
            .map_err(|e| anyhow::anyhow!("Cannot read venv_setup_script '{}': {}", full_path, e))?;
        body.lines()
            .map(|l| l.trim().to_string())
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
            .collect()
    } else {
        recipe.venv_setup.to_vec()
    };

    if setup.is_empty() {
        return Ok(());
    }

    info!("{}", fluent_i18n::t!("venv-preparing", { "venv" => venv }));

    let mut handles = Vec::with_capacity(cluster.nodes.len());
    for node in &cluster.nodes {
        let config = config.clone();
        let node = node.clone();
        let venv = venv.clone();
        let setup = setup.clone();
        let dry = dry_run;

        handles.push(crate::locale::spawn_with_locale(async move {
            (
                node.name.clone(),
                prepare_venv_for_node(&config, &node, &venv, &setup, dry).await,
            )
        }));
    }

    let mut any_failed = false;
    for h in handles {
        let (name, result) = h.await?;
        if let Err(e) = result {
            any_failed = true;
            warn!(
                "{}",
                fluent_i18n::t!("venv-prepare-failed", { "name" => &name, "error" => e.to_string() })
            );
        }
    }

    if any_failed {
        anyhow::bail!("Venv preparation failed on at least one node (see warnings above)");
    }

    Ok(())
}

/// Checks if the venv exists on a single node and creates it if not.
async fn prepare_venv_for_node(
    config: &Config,
    node: &Node,
    venv: &str,
    setup: &[String],
    dry_run: bool,
) -> Result<()> {
    let check_cmd = format!("test -f {}/bin/activate", venv);

    if dry_run {
        info!(
            "  ssh {}@{}: {} (dry-run)",
            config.ssh_user, node.ssh_ip, check_cmd
        );
        for cmd in setup {
            info!("    {} (dry-run)", cmd);
        }
        return Ok(());
    }

    // Check if venv already exists.
    match crate::ssh::run_ssh(config, node, check_cmd.clone(), false).await {
        Ok(_) => {
            info!(
                "{}",
                fluent_i18n::t!("venv-exists", { "name" => &node.name })
            );
            return Ok(());
        }
        Err(_) => {
            info!(
                "{}",
                fluent_i18n::t!("venv-creating", { "name" => &node.name })
            );
        }
    }

    // Run setup commands in order.
    for cmd in setup {
        info!("  ssh {}@{}: {}", config.ssh_user, node.ssh_ip, cmd);
        crate::ssh::run_ssh(config, node, cmd.clone(), false).await?;
    }

    // Re-check that the venv was created.
    match crate::ssh::run_ssh(config, node, check_cmd, false).await {
        Ok(_) => {
            info!(
                "{}",
                fluent_i18n::t!("venv-created", { "name" => &node.name })
            );
            Ok(())
        }
        Err(_) => {
            anyhow::bail!(
                "{}",
                fluent_i18n::t!("venv-not-created", { "name" => &node.name })
            );
        }
    }
}