//! Builds the "tail" of the docker command — everything after the image
//! name (entrypoint, model, args, auto-generated flags).

use crate::config::{DistributedBackend, Node, Recipe};

use super::shq;

/// The subset of [`DockerCommandCtx`] that `build_tail` needs.
pub(crate) struct TailCtx<'a> {
    pub recipe: &'a Recipe,
    pub tp_size: u32,
    pub pp_size: u32,
    pub nnodes: u32,
    pub init_port: u16,
    pub master_ib_ip: &'a str,
}

/// Build the docker command tail — everything after the image name.
pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
    let is_sglang = ctx.recipe.engine == "sglang";
    let is_distributed = matches!(
        ctx.recipe.distributed_backend,
        Some(DistributedBackend::Mp | DistributedBackend::Ray)
    );

    let mut tail = Vec::new();

    if !ctx.recipe.entrypoint.is_empty() {
        let mut ep = ctx.recipe.entrypoint.clone();
        if ctx.recipe.mods.commands.is_empty() && ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
            let last = ep.len() - 1;
            ep[last] = shq(&ep[last]);
        }
        tail.extend(ep);
    }

    if is_sglang {
        tail.push("--model-path".to_string());
    }
    tail.push(ctx.recipe.model.clone());
    tail.extend(ctx.recipe.args.to_vec());

    // Single-node run (solo, required_nodes=1, or a single --node): emit
    // NO distributed flags — the engine defaults to single-node. Launching
    // one container with `--nnodes 1 --master-addr ...` is redundant at best
    // and rejected by some engines at worst. Multi-node runs get the full
    // rendezvous flag set below.
    let mut flags: Vec<String> = if ctx.nnodes > 1 {
        match ctx.recipe.distributed_backend {
            Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => vec![
                "--tp".into(),
                ctx.tp_size.to_string(),
                "--nnodes".into(),
                ctx.nnodes.to_string(),
                "--node-rank".into(),
                node.rank.to_string(),
                "--dist-init-addr".into(),
                format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
            ],
            Some(DistributedBackend::Ray | DistributedBackend::Mp) => {
                let mut fp = vec![
                    "--tensor-parallel-size".into(),
                    ctx.tp_size.to_string(),
                    "--pipeline-parallel-size".into(),
                    ctx.pp_size.to_string(),
                    "--nnodes".into(),
                    ctx.nnodes.to_string(),
                    "--node-rank".into(),
                    node.rank.to_string(),
                    "--master-addr".into(),
                    ctx.master_ib_ip.to_string(),
                    "--master-port".into(),
                    ctx.init_port.to_string(),
                ];
                if ctx.recipe.distributed_backend == Some(DistributedBackend::Ray) {
                    fp.push("--distributed-executor-backend".into());
                    fp.push("ray".into());
                }
                fp
            }
            _ => vec![],
        }
    } else {
        vec![]
    };
    if ctx.nnodes > 1 && !is_sglang && node.rank > 0 && is_distributed {
        flags.push("--headless".into());
    }
    tail.extend(flags);
    tail.push("--host".to_string());
    tail.push("0.0.0.0".to_string());
    tail.extend(["--port".into(), ctx.recipe.port.to_string()]);
    tail
}