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

use super::{mods, tail};

/// Shared context for [`build_docker_command`] — constructed once, cloned per node.
#[derive(Clone)]
pub(crate) struct DockerCommandCtx {
    pub cluster: Cluster,
    pub nccl: 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`. Model-subfolder mounts
    /// are NOT here — `start_cluster` resolves those per node and merges the
    /// results into this `Vec` (as `ResolvedMount`s) 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. Also the final form a model-subfolder
/// mount takes after per-node resolution.
#[derive(Clone)]
pub(crate) struct ResolvedMount {
    pub host: String,
    pub container: String,
    pub read_only: bool,
    pub env: Option<String>,
}

/// A model-subfolder mount spec (`repo`+`subfolder`) before per-node
/// resolution. The host path is resolved at serve time from each node's HF
/// cache (so it can differ across nodes); `container` is the fixed path the
/// recipe references in args.
#[derive(Clone)]
pub(crate) struct ModelMountSpec {
    pub repo: String,
    pub subfolder: Option<String>,
    pub container: String,
    pub read_only: bool,
    pub env: Option<String>,
}

/// Partition a recipe's `extra_mounts` into static (resolved now) and
/// model-subfolder (deferred to per-node resolution). 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 model_mounts = 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
                );
            }
            model_mounts.push(ModelMountSpec {
                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()),
                container: m.container.clone(),
                read_only: m.read_only,
                env: m.env.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,
        model_mounts,
    })
}

/// Result of [`resolve_extra_mounts`]: static mounts ready to clone per node,
/// model-subfolder specs deferred to per-node resolution.
pub(crate) struct ResolvedExtraMounts {
    pub static_mounts: Vec<ResolvedMount>,
    pub model_mounts: Vec<ModelMountSpec>,
}

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

/// Collect NCCL cluster env vars for multi-node deployments.
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 sock = ctx.nccl.socket_ifname.as_str();
    let gloo_if = ctx.nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
    let ucx_dev = ctx.nccl.ucx_net_devices.as_deref().unwrap_or(sock);
    let mut env: Vec<(&str, Cow<str>)> = vec![
        ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
        ("NCCL_IB_HCA", ctx.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",
            ctx.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) = ctx.nccl.nccl_cross_nic {
        env.push(("NCCL_CROSS_NIC", v.to_string().into()));
    }
    if let Some(ref v) = ctx.nccl.nccl_proto {
        env.push(("NCCL_PROTO", v.as_str().into()));
    }
    if let Some(ref v) = ctx.nccl.nccl_p2p_level {
        env.push(("NCCL_P2P_LEVEL", v.as_str().into()));
    }
    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)]
mod tests {
    use super::*;
    use crate::config::{
        Cluster, DistributedBackend, ExtraMount, NcclConfig, Node, Recipe, RecipeArgs,
    };
    use std::collections::HashMap;

    fn tcls() -> Cluster {
        Cluster {
            name: "t".into(),
            master_ib_ip: "1".into(),
            nodes: vec![Node {
                name: "n0".into(),
                ssh_ip: "10.0.0.1".into(),
                ib_ip: "1".into(),
                rank: 0,
                host_key: None,
                entity: None,
            }],
            extra_entities: vec![],
        }
    }

    fn tnccl() -> NcclConfig {
        NcclConfig {
            ib_hca: "r".into(),
            socket_ifname: "e".into(),
            ib_gid_index: 3,
            ib_disable: 0,
            gloo_socket_ifname: None,
            ucx_net_devices: None,
            nccl_cross_nic: None,
            nccl_proto: None,
            nccl_p2p_level: None,
        }
    }

    fn minimal_config() -> Config {
        Config {
            ssh_user: "u".into(),
            ssh_key_path: "/tmp/key".into(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            global_mods_dir: None,
            drop_caches: true,
            mount_aliases: HashMap::from([("models".to_string(), "/mnt/nas/models".to_string())]),
        }
    }

    fn run_with(recipe: &Recipe) -> String {
        let resolved = resolve_extra_mounts(&minimal_config(), recipe).unwrap();
        let ctx = DockerCommandCtx {
            cluster: tcls(),
            nccl: tnccl(),
            recipe: recipe.clone(),
            recipe_name: "r".to_string(),
            pp_size: 2,
            tp_size: 1,
            nnodes: 2,
            cache_dir: "/home/user/.cache/huggingface".to_string(),
            keep: false,
            init_port: 25000,
            mods_remote_path: String::new(),
            extra_mounts: resolved.static_mounts,
        };
        build_docker_command(&ctx, &tcls().nodes[0], true).join(" ")
    }

    fn base_recipe() -> Recipe {
        Recipe {
            engine: "vllm".into(),
            name: "T".into(),
            docker_image: "img".into(),
            model: "mdl".into(),
            port: 8000,
            required_nodes: Some(2),
            pipeline_parallel_size: Some(2),
            tensor_parallel_size: 1,
            args: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
            env: HashMap::new(),
            cluster: None,
            entrypoint: vec![],
            distributed_backend: Some(DistributedBackend::Mp),
            before_serve: vec![],
            mods: Default::default(),
            shm_size: String::new(),
            sync_image: true,
            docker_image_dir: None,
            extra_mounts: vec![],
        }
    }

    #[test]
    fn extra_mounts_are_mounted() {
        let mut r = base_recipe();
        r.extra_mounts.push(ExtraMount {
            alias: Some("models".into()),
            host: Some("~/.cache/models".into()),
            repo: None,
            subfolder: None,
            container: "/models:ro".into(),
            read_only: true,
            env: Some("MODEL_DIR".into()),
        });
        let cmd = run_with(&r);
        // alias present in minimal_config() → alias host wins
        assert!(cmd.contains("-v /mnt/nas/models:/models:ro:ro"));
        assert!(cmd.contains("-v /home/user/.cache/huggingface:/root/.cache/huggingface"));
        // env injected pointing at the container path, before recipe env
        assert!(cmd.contains("-e MODEL_DIR=/models:ro"));
    }

    #[test]
    fn extra_mounts_fallback_to_host_when_alias_absent() {
        let mut r = base_recipe();
        r.extra_mounts.push(ExtraMount {
            alias: Some("no_such_alias".into()),
            host: Some("~/.cache/jit".into()),
            repo: None,
            subfolder: None,
            container: "/root/.cache/jit".into(),
            read_only: false,
            env: None,
        });
        let cmd = run_with(&r);
        // alias not in config → host default used with ~ expansion (ssh_user=u)
        assert!(cmd.contains("-v /home/u/.cache/jit:/root/.cache/jit"));
        assert!(!cmd.contains(":ro"));
    }

    #[test]
    fn model_subfolder_mount_uses_placeholder_in_dry_run() {
        // In build_docker_command (no per-node SSH resolution here), a model
        // mount shows the descriptive placeholder host so the printed command
        // documents what would be mounted. start_cluster replaces this with
        // the real resolved path at serve time (or keeps it under dry-run).
        let mut r = base_recipe();
        r.extra_mounts.push(ExtraMount {
            alias: None,
            host: None,
            repo: Some("aidendle94/GLM-5.2-Hybrid-FP8-MXFP4".into()),
            subfolder: Some("mtp-draft".into()),
            container: "/models/mtp-draft".into(),
            read_only: true,
            env: Some("MTP_DRAFT_PATH".into()),
        });
        let cmd = run_with(&r);
        // model_mounts are NOT in extra_mounts here (run_with doesn't resolve
        // them per-node), so the -v flag for the model mount is absent — only
        // start_cluster emits it. Verify the spec is parsed + the env is NOT
        // injected yet (env injection follows extra_mounts, which is empty of
        // the model mount here). This test guards the partition.
        assert!(!cmd.contains("-v /models/mtp-draft"));
        assert!(!cmd.contains("MTP_DRAFT_PATH"));
        // But the recipe parsed into a model_mount spec:
        let resolved = resolve_extra_mounts(&minimal_config(), &r).unwrap();
        assert_eq!(resolved.model_mounts.len(), 1);
        assert_eq!(resolved.static_mounts.len(), 0);
        let ph = model_mount_placeholder(&resolved.model_mounts[0]);
        assert_eq!(ph.container, "/models/mtp-draft");
        assert!(
            ph.host
                .contains("models--aidendle94--GLM-5.2-Hybrid-FP8-MXFP4")
        );
        assert!(ph.host.contains("mtp-draft"));
        assert!(ph.read_only);
    }
}