//! 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, tail};
use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
use anyhow::Result;
use std::borrow::Cow;

/// 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>,
}

/// 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() {
            *v = v.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.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.ib_ip.as_str());
    if is_sglang {
        env.push(("SGLANG_WORKER_IP", worker_ip));
    } else {
        env.push(("VLLM_HOST_IP", worker_ip));
    }
    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) -> 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.master_ib_ip).into(),
        ));
        pairs.push(("RAY_PORT", "6379".into()));
        pairs.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
        pairs.push(("RAY_OVERRIDE_NODE_IP_ADDRESS", node.ib_ip.as_str().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()));
        }
    }
    for (k, v) in &ctx.recipe.env {
        pairs.push((k.as_str(), v.as_str().into()));
    }
    pairs
        .into_iter()
        .map(|(k, v)| format!("-e {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,
) -> Vec<String> {
    let flags = build_docker_flags(ctx);
    let mut cmd = vec![
        format!("docker run -d {}", flags.join(" ")),
        format!("-v {}:/root/.cache/huggingface", ctx.cache_dir),
    ];
    if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
        cmd.push(format!("-v {}:/mods:ro", ctx.mods_remote_path));
    }
    for m in &ctx.extra_mounts {
        let mut spec = format!("-v {}:{}", m.host, 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 sm121_{}_{}", ctx.recipe_name, node.rank));
    cmd.push(format!("--label sm121.recipe={}", ctx.recipe_name));
    cmd.push(format!("--label sm121.rank={}", node.rank));
    cmd.push(format!("--label 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 sm121.solo-node={}", node.name));
    }
    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: &ctx.cluster.master_ib_ip,
    };
    let tail = tail::build_tail(&tctx, node);
    if ctx.recipe.mods.commands.is_empty() {
        cmd.push(ctx.recipe.docker_image.clone());
        cmd.extend(tail);
    } else {
        let script = mods::build_mod_script(&ctx.recipe.mods.commands, &tail);
        cmd.extend([
            "--entrypoint sh".into(),
            ctx.recipe.docker_image.clone(),
            "-c".into(),
            script,
        ]);
    }
    cmd
}

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