//! Assembly of the `docker run` command per node — flags, env vars, mounts,
//! and the vLLM/SGLang tail (engine flags + distributed rendezvous).
//!
//! [`build_docker_command`] is the single entry point; everything else here is
//! private to this module.

use super::{mods, shq, tail};
use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
use crate::ssh::run_ssh;
use anyhow::Result;
use std::borrow::Cow;
use tracing::info;

/// Shared context for [`build_docker_command`] — constructed once, cloned per node.
#[derive(Clone)]
pub(crate) struct DockerCommandCtx {
    pub cluster: Cluster,
    /// NCCL config — `None` for single-node recipes (no NCCL env vars
    /// are emitted). For multi-node, always `Some`.
    pub nccl: Option<NcclConfig>,
    pub recipe: Recipe,
    pub recipe_name: String,
    pub pp_size: u32,
    pub tp_size: u32,
    pub nnodes: u32,
    pub cache_dir: String,
    pub keep: bool,
    pub init_port: u16,
    pub mods_remote_path: String,
    /// Static `extra_mounts` (alias/host) with host paths already resolved
    /// once against `Config.mount_aliases` / `ssh_user`. Path references
    /// (`repo`/`subfolder`) are NOT here — `start_cluster` resolves those
    /// per-node and substitutes `{{<name>}}` into the recipe args/env before
    /// calling [`build_docker_command`].
    pub extra_mounts: Vec<ResolvedMount>,
}

/// A recipe `ExtraMount` after host-path resolution (alias + tilde), ready to
/// turn into a `-v` flag / `-e` env pair.
#[derive(Clone)]
pub(crate) struct ResolvedMount {
    pub host: String,
    pub container: String,
    pub read_only: bool,
    pub env: Option<String>,
}

/// A model path reference (`repo`+`subfolder`) before per-node resolution.
/// At serve time, the launcher resolves the path on each node (reading
/// `refs/main`) and substitutes `{{<name>}}` in the recipe args/env with the
/// resolved full HF cache path. The HF cache is already mounted at
/// `/root/.cache/huggingface`, so the resolved path's symlinks resolve
/// correctly (a separate bind-mount would break them).
#[derive(Clone)]
pub(crate) struct PathRef {
    pub repo: String,
    pub subfolder: Option<String>,
    pub name: String,
}

/// Partition a recipe's `extra_mounts` into static (resolved now) and
/// path references (deferred to per-node resolution + substitution).
/// Validates that a mount is one or the other — `repo` set means
/// `alias`/`host` must be empty, and `subfolder` without `repo` is an error.
pub(crate) fn resolve_extra_mounts(
    config: &Config,
    recipe: &Recipe,
) -> Result<ResolvedExtraMounts> {
    let mut static_mounts = Vec::new();
    let mut path_refs = Vec::new();
    for m in &recipe.extra_mounts {
        if m.repo.as_deref().is_some_and(|r| !r.trim().is_empty()) {
            if m.alias.is_some() || m.host.is_some() {
                anyhow::bail!(
                    "extra_mount '{}': `repo` is mutually exclusive with `alias`/`host`",
                    m.container
                );
            }
            path_refs.push(PathRef {
                repo: m.repo.as_ref().unwrap().trim().to_string(),
                subfolder: m
                    .subfolder
                    .as_deref()
                    .filter(|s| !s.trim().is_empty())
                    .map(|s| s.trim().to_string()),
                name: m.container.clone(),
            });
        } else if m.subfolder.as_deref().is_some_and(|s| !s.trim().is_empty()) {
            anyhow::bail!("extra_mount '{}': `subfolder` requires `repo`", m.container);
        } else {
            let host = crate::config::resolve_mount_host(
                m.alias.as_deref(),
                m.host.as_deref(),
                &config.mount_aliases,
                &config.ssh_user,
                &m.container,
            )?;
            static_mounts.push(ResolvedMount {
                host,
                container: m.container.clone(),
                read_only: m.read_only,
                env: m.env.clone(),
            });
        }
    }
    Ok(ResolvedExtraMounts {
        static_mounts,
        path_refs,
    })
}

/// Result of [`resolve_extra_mounts`]: static mounts ready to clone per node,
/// path references deferred to per-node resolution + substitution.
pub(crate) struct ResolvedExtraMounts {
    pub static_mounts: Vec<ResolvedMount>,
    pub path_refs: Vec<PathRef>,
}

/// `--clear`: delete the host-side static mount dirs on the launch nodes
/// before containers start. Containers run as root, so caches there are
/// usually root-owned → `sudo rm -rf`. Dry run prints the exact command
/// (via [`run_ssh`]) and runs nothing. A failure on any node aborts the
/// serve — never start against a stale cache. Repo-based path refs are
/// never cleared (they point into the HF cache).
pub(crate) async fn clear_extra_mounts(
    config: &Config,
    nodes: &[&Node],
    resolved: &ResolvedExtraMounts,
    clear: bool,
    dry_run: bool,
) -> Result<()> {
    if !clear {
        return Ok(());
    }
    let paths = clear_paths(&resolved.static_mounts);
    if paths.is_empty() {
        info!("--clear: no host-side extra_mounts to clear");
        return Ok(());
    }
    info!(
        "--clear: clearing {} host extra_mount dir(s) on {} launch node(s)",
        paths.len(),
        nodes.len()
    );
    for node in nodes {
        for p in &paths {
            let cmd = format!("sudo rm -rf {}", shq(p));
            run_ssh(config, node, cmd, dry_run).await?;
        }
    }
    Ok(())
}

/// Dedup + sort the static mount host paths (the set `--clear` deletes).
pub(crate) fn clear_paths(mounts: &[ResolvedMount]) -> Vec<String> {
    let mut paths: Vec<String> = mounts.iter().map(|m| m.host.clone()).collect();
    paths.sort();
    paths.dedup();
    paths
}

/// Substitute `{{<name>}}` placeholders in `args` and `env` with resolved paths.
/// `path_refs` maps placeholder name → resolved path (per-node). Done before
/// building the docker command so vLLM args reference the correct snapshot.
pub(crate) fn substitute_path_refs(recipe: &Recipe, path_refs: &[(String, String)]) -> Recipe {
    let mut r = recipe.clone();
    for (name, path) in path_refs {
        let placeholder = format!("{{{{{}}}}}", name);
        r.args = r.args.substitute(&placeholder, path);
        for v in r.env.values_mut() {
            // Ref entries reference recipe fields — nothing to interpolate
            // (they resolve at serve time via `resolved_env`).
            if let crate::config::EnvValue::Plain(s) = v {
                *s = s.replace(&placeholder, path);
            }
        }
    }
    r
}

/// Dry-run placeholder for a model path reference: nodes may be off, so no
/// HF cache to read. Emits a descriptive path so the printed `docker run`
/// command shows *what* would be substituted, not a broken lookup.
pub(crate) fn path_ref_placeholder(spec: &PathRef) -> String {
    let sub = spec.subfolder.as_deref().unwrap_or("");
    format!(
        "<hf_cache>/hub/{}/snapshots/<refs/main>/{}",
        crate::config::model_cache_dir_name(&spec.repo),
        sub
    )
}

/// Collect NCCL cluster env vars for multi-node deployments.
///
/// `ctx.nccl` is `None` for single-node recipes — but this function is
/// only called when `ctx.nnodes > 1`, so `None` indicates a configuration
/// error (multi-node without NCCL config). In that case, only the worker
/// IP env var is emitted (NCCL env vars are skipped), which will cause
/// NCCL to fail at runtime with a visible error.
fn collect_nccl_env<'a>(ctx: &'a DockerCommandCtx, node: &'a Node) -> Vec<(&'a str, Cow<'a, str>)> {
    let is_sglang = ctx.recipe.engine == "sglang";
    let mut env: Vec<(&str, Cow<str>)> = Vec::new();
    if let Some(nccl) = ctx.nccl.as_ref() {
        let sock = nccl.socket_ifname.as_str();
        let gloo_if = nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
        let ucx_dev = nccl.ucx_net_devices.as_deref().unwrap_or(sock);
        env.extend(vec![
            ("NCCL_IB_DISABLE", format!("{}", nccl.ib_disable).into()),
            ("NCCL_IB_HCA", nccl.ib_hca.as_str().into()),
            ("NCCL_SOCKET_IFNAME", sock.into()),
            ("GLOO_SOCKET_IFNAME", gloo_if.into()),
            ("TP_SOCKET_IFNAME", gloo_if.into()),
            (
                "NCCL_IB_GID_INDEX",
                nccl.ib_gid_index
                    .map_or_else(|| Cow::Borrowed("auto"), |g| g.to_string().into()),
            ),
            ("MN_IF_NAME", gloo_if.into()),
            ("UCX_NET_DEVICES", ucx_dev.into()),
            ("OMPI_MCA_btl_tcp_if_include", gloo_if.into()),
            ("OMP_NUM_THREADS", "4".into()),
            ("NCCL_IGNORE_CPU_AFFINITY", "1".into()),
            ("NCCL_NET", "IB".into()),
            ("NCCL_DEBUG", "WARN".into()),
        ]);
        if let Some(v) = nccl.nccl_cross_nic {
            env.push(("NCCL_CROSS_NIC", v.to_string().into()));
        }
        if let Some(ref v) = nccl.nccl_proto {
            env.push(("NCCL_PROTO", v.as_str().into()));
        }
        if let Some(ref v) = nccl.nccl_p2p_level {
            env.push(("NCCL_P2P_LEVEL", v.as_str().into()));
        }
    }
    // Worker IP is always set for multi-node (collect_nccl_env is only
    // called when nnodes > 1).
    let worker_ip = Cow::Borrowed(node.effective_ib_ip());
    if is_sglang {
        env.push(("SGLANG_WORKER_IP", worker_ip));
    } else {
        env.push(("VLLM_HOST_IP", worker_ip));
    }
    // Per-node rendezvous env: entrypoint-driven images (e.g. aiden's spark
    // launchers) read NODE_RANK/MASTER_ADDR from the environment instead of
    // CLI flags. Inert for stock vllm/sglang (they use the CLI flags).
    env.push(("NODE_RANK", node.rank.to_string().into()));
    env.push(("MASTER_ADDR", ctx.cluster.effective_master_ip().into()));
    env
}

/// Build the list of `-e KEY=VALUE` Docker env var pairs for a node.
fn collect_docker_env(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Result<Vec<String>> {
    let mut pairs: Vec<(&str, Cow<str>)> =
        vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
    if ctx.nnodes > 1 {
        pairs.extend(collect_nccl_env(ctx, node));
    }
    if is_head {
        pairs.push(("VLLM_LOGGING_COLOR", "1".into()));
        pairs.push(("TERM", "xterm-256color".into()));
    }
    if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
        pairs.push((
            "RAY_ADDRESS",
            format!("{}:6379", ctx.cluster.effective_master_ip()).into(),
        ));
        pairs.push(("RAY_PORT", "6379".into()));
        pairs.push(("RAY_NODE_IP_ADDRESS", node.effective_ib_ip().into()));
        pairs.push((
            "RAY_OVERRIDE_NODE_IP_ADDRESS",
            node.effective_ib_ip().into(),
        ));
    }
    // extra_mount `env` defaults (point at the container path) are injected
    // before recipe [env] so the recipe can still override them (last wins).
    for m in &ctx.extra_mounts {
        if let Some(name) = &m.env {
            pairs.push((name.as_str(), m.container.as_str().into()));
        }
    }
    let recipe_env = ctx.recipe.resolved_env()?;
    for (k, v) in &recipe_env {
        pairs.push((k.as_str(), v.as_str().into()));
    }
    // Quote the whole `K=V` so a recipe `[env]` value or name containing
    // shell metacharacters cannot inject into the remote `docker run` shell.
    Ok(pairs
        .into_iter()
        .map(|(k, v)| format!("-e {}", shq(&format!("{k}={v}"))))
        .collect())
}

/// Build the list of Docker run flags (--gpus, --ipc, --ulimit, etc.).
fn build_docker_flags(ctx: &DockerCommandCtx) -> Vec<String> {
    let mut flags: Vec<String> = vec![
        "--gpus all".to_string(),
        "--ipc=host".to_string(),
        "--privileged".to_string(),
        "--network=host".to_string(),
        "--init".to_string(),
        "--ulimit memlock=-1".to_string(),
        "--ulimit stack=67108864".to_string(),
        "--ulimit nofile=1048576".to_string(),
    ];
    if !ctx.recipe.shm_size.is_empty() {
        flags.push(format!("--shm-size {}", ctx.recipe.shm_size.as_str()));
    }
    if !ctx.keep {
        flags.push("--rm".to_string());
    }
    flags
}

/// Build the full Docker command line for a node.
pub(crate) fn build_docker_command(
    ctx: &DockerCommandCtx,
    node: &Node,
    is_head: bool,
) -> Result<Vec<String>> {
    let flags = build_docker_flags(ctx);
    let mut cmd = vec![
        format!("docker run -d {}", flags.join(" ")),
        format!("-v {}:/root/.cache/huggingface", shq(&ctx.cache_dir)),
    ];
    if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
        cmd.push(format!("-v {}:/mods:ro", shq(&ctx.mods_remote_path)));
    }
    for m in &ctx.extra_mounts {
        // Adjacent quoted words concatenate into one shell arg, so
        // `-v 'host':\'container':ro` is a single `-v host:container:ro`.
        let mut spec = format!("-v {}:{}", shq(&m.host), shq(&m.container));
        if m.read_only {
            spec.push_str(":ro");
        }
        cmd.push(spec);
    }
    cmd.extend(collect_docker_env(ctx, node, is_head)?);
    cmd.push(format!(
        "--name {}",
        shq(&format!("sm121_{}_{}", ctx.recipe_name, node.rank))
    ));
    cmd.push(format!(
        "--label {}",
        shq(&format!("sm121.recipe={}", ctx.recipe_name))
    ));
    cmd.push(format!(
        "--label {}",
        shq(&format!("sm121.rank={}", node.rank))
    ));
    cmd.push(format!(
        "--label {}",
        shq(&format!("sm121.cluster={}", ctx.cluster.name))
    ));
    // Add solo-node label for single-node deployments
    if ctx.cluster.nodes.len() == 1 && ctx.cluster.extra_entities.is_empty() {
        cmd.push(format!(
            "--label {}",
            shq(&format!("sm121.solo-node={}", node.name))
        ));
    }
    let master_ip = ctx.cluster.effective_master_ip();
    let tctx = tail::TailCtx {
        recipe: &ctx.recipe,
        tp_size: ctx.tp_size,
        pp_size: ctx.pp_size,
        nnodes: ctx.nnodes,
        init_port: ctx.init_port,
        master_ib_ip: &master_ip,
    };
    let tail = tail::build_tail(&tctx, node);
    let img = shq(&ctx.recipe.docker_image);
    if ctx.recipe.mods.commands.is_empty() {
        cmd.push(img);
        // Shell-quote every tail word (entrypoint, model, args, flags) so a
        // crafted recipe arg/env cannot break out of the remote `docker run`.
        cmd.extend(tail.iter().map(|t| shq(t)));
    } else {
        let script = mods::build_mod_script(&ctx.recipe.mods.commands, &tail);
        // `--entrypoint sh` and `-c` stay literal; `img` and `script` are
        // shell-quoted so they survive the remote shell as one argv each.
        cmd.extend(["--entrypoint sh".into(), img, "-c".into(), script]);
    }
    Ok(cmd)
}

#[cfg(test)]
#[path = "command_tests.rs"]
mod tests;