use super::*;
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,
            },
            Node {
                name: "n1".into(),
                ssh_ip: "10.0.0.2".into(),
                ib_ip: "2".into(),
                rank: 1,
                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,
    }
}

fn trec() -> 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: vec!["--x".into(), "4096".into()],
        env: HashMap::from([("K".into(), "V".into())]),
        cluster: None,
        entrypoint: vec![],
        distributed_backend: "mp".into(),
        before_serve: vec![],
        shm_size: String::new(),
    }
}

fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
        2
    } else {
        recipe.tensor_parallel_size
    };
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "r",
        node,
        head,
        2,
        tp,
        2,
        "/home/user/.cache/huggingface",
        false,
        25000,
    )
    .join(" ")
}

fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
        2
    } else {
        recipe.tensor_parallel_size
    };
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "r",
        node,
        head,
        2,
        tp,
        2,
        "/home/user/.cache/huggingface",
        true,
        25000,
    )
    .join(" ")
}

#[test]
fn mp() {
    let cmd = run(&trec(), &tcls().nodes[0], true);
    assert!(cmd.contains("docker run -d --gpus all"));
    assert!(cmd.contains("-e NCCL_IB_HCA=r"));
    assert!(cmd.contains("-e K=V"));
    assert!(cmd.contains("--name sm121_r_0"));
    assert!(
        !cmd.contains("-p "),
        "host networking: no port mapping expected"
    );
    assert!(cmd.contains("--label sm121.recipe=r"));
    assert!(cmd.contains("--label sm121.rank=0"));
    assert!(cmd.contains("img"));
    assert!(cmd.contains("mdl"));
    assert!(cmd.contains("--tensor-parallel-size 1"));
    assert!(cmd.contains("--pipeline-parallel-size 2"));
    assert!(cmd.contains("--nnodes 2"));
    assert!(cmd.contains("--node-rank 0"));
    assert!(cmd.contains("--master-addr 1"));
    assert!(cmd.contains("--master-port 25000"));
    // vLLM auto-detects mp; only explicitly pass for ray
    assert!(
        !cmd.contains("--distributed-executor-backend"),
        "mp backend should be auto-detected"
    );
    assert!(!cmd.contains("RAY_ADDRESS"));
    assert!(cmd.contains("--rm"), "default --rm should be present");
    // New: defaults & cluster env
    assert!(
        !cmd.contains("--shm-size"),
        "shm-size should not be set by default"
    );
    assert!(cmd.contains("--ulimit memlock=-1"));
    assert!(cmd.contains("--ulimit stack=67108864"));
    assert!(
        cmd.contains("-e GLOO_SOCKET_IFNAME=e"),
        "default GLOO_SOCKET_IFNAME missing"
    );
    assert!(
        cmd.contains("-e TP_SOCKET_IFNAME=e"),
        "default TP_SOCKET_IFNAME missing"
    );
    assert!(
        cmd.contains("-e OMP_NUM_THREADS=4"),
        "default OMP_NUM_THREADS missing"
    );
    assert!(
        cmd.contains("-e NCCL_IGNORE_CPU_AFFINITY=1"),
        "default NCCL_IGNORE_CPU_AFFINITY missing"
    );
    assert!(cmd.contains("-e NCCL_NET=IB"), "default NCCL_NET missing");
    assert!(
        cmd.contains("-e NCCL_DEBUG=WARN"),
        "default NCCL_DEBUG missing"
    );
    assert!(
        cmd.contains("-e MN_IF_NAME=e"),
        "cluster MN_IF_NAME missing"
    );
    assert!(
        cmd.contains("-e UCX_NET_DEVICES=e"),
        "cluster UCX_NET_DEVICES missing"
    );
    assert!(
        cmd.contains("-e OMPI_MCA_btl_tcp_if_include=e"),
        "cluster OMPI_MCA_btl_tcp_if_include missing"
    );
    // Head node: no --headless
    assert!(!cmd.contains("--headless"));
}

#[test]
fn worker_no_port() {
    let cmd = run(&trec(), &tcls().nodes[1], false);
    assert!(!cmd.contains("-p 8000:8000"));
    assert!(cmd.contains("--node-rank 1"));
    assert!(cmd.contains("--label sm121.rank=1"));
    assert!(cmd.contains("--name sm121_r_1"));
    // New: vLLM workers run --headless
    assert!(cmd.contains("--headless"), "vLLM worker should be headless");
}

#[test]
fn ray_backend() {
    let mut r = trec();
    r.distributed_backend = "ray".into();
    let cmd = run(&r, &tcls().nodes[0], true);
    assert!(cmd.contains("-e RAY_ADDRESS=1:6379"));
    assert!(cmd.contains("-e RAY_PORT=6379"));
    assert!(cmd.contains("--node-rank 0"));
    assert!(cmd.contains("--distributed-executor-backend ray"));
}

#[test]
fn custom_tensor_parallel_size() {
    // Per-node TP=1 is the default. Custom TP (e.g. multi-node TP)
    // can be set via recipe.tensor_parallel_size.
    let mut r = trec();
    r.tensor_parallel_size = 2;
    let cmd = run(&r, &tcls().nodes[0], true);
    assert!(cmd.contains("--tensor-parallel-size 2"), "custom TP");
}

#[test]
fn auto_pp_size() {
    let cmd = run(&trec(), &tcls().nodes[0], true);
    assert!(cmd.contains("--pipeline-parallel-size 2"));
}

#[test]
fn entrypoint() {
    let mut r = trec();
    r.entrypoint = vec!["py".into(), "-m".into(), "mod".into()];
    let cmd = run(&r, &tcls().nodes[0], true);
    let ip = cmd.find("img").unwrap();
    let mp = cmd.find("mdl").unwrap();
    let ep = cmd.find("py -m mod").unwrap();
    assert!(
        ip < ep && ep < mp,
        "entrypoint should be between image and model"
    );
}

// ── SGLang engine tests ──────────────────────────────────────

fn sglang_recipe() -> Recipe {
    Recipe {
        engine: "sglang".into(),
        name: "S".into(),
        docker_image: "sglang-img".into(),
        model: "org/sgl-model".into(),
        port: 8000,
        required_nodes: Some(4),
        pipeline_parallel_size: None,
        tensor_parallel_size: 4,
        args: vec!["--mem-fraction-static".into(), "0.8".into()],
        env: HashMap::from([("TORCH_NCCL_ASYNC_ERROR_HANDLING".into(), "1".into())]),
        cluster: None,
        entrypoint: vec!["sglang".into(), "serve".into()],
        distributed_backend: "mp".into(),
        before_serve: vec![],
        shm_size: String::new(),
    }
}

fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
    let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
        4
    } else {
        recipe.tensor_parallel_size
    };
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "sg",
        node,
        head,
        4,
        tp,
        4,
        "/home/user/.cache/huggingface",
        false,
        25000,
    )
    .join(" ")
}

#[test]
fn sglang_host_ip() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    assert!(cmd.contains("-e SGLANG_WORKER_IP=1"));
    assert!(!cmd.contains("VLLM_HOST_IP"));
}

#[test]
fn sglang_head() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    assert!(cmd.contains("--model-path"));
    assert!(cmd.contains("org/sgl-model"));
    assert!(cmd.contains("--tp 4"));
    assert!(cmd.contains("--nnodes 4"));
    assert!(cmd.contains("--node-rank 0"));
    assert!(cmd.contains("--dist-init-addr 1:25000"));
    assert!(!cmd.contains("--master-addr"));
    assert!(!cmd.contains("--tensor-parallel-size"));
    assert!(!cmd.contains("--pipeline-parallel-size"));
    assert!(!cmd.contains("--distributed-executor-backend"));
    assert!(cmd.contains("-e TORCH_NCCL_ASYNC_ERROR_HANDLING=1"));
    assert!(
        !cmd.contains("-p "),
        "host networking: no port mapping expected"
    );
    assert!(cmd.contains("--port 8000"));
    assert!(cmd.contains("-e HF_HOME=/root/.cache/huggingface"));
    assert!(cmd.contains("-v /home/user/.cache/huggingface:/root/.cache/huggingface"));
    assert!(
        cmd.contains("--rm"),
        "SGLang default --rm should be present"
    );
}

#[test]
fn sglang_worker() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[1], false);
    assert!(!cmd.contains("-p 8000:8000"));
    assert!(cmd.contains("--node-rank 1"));
    assert!(cmd.contains("--port 8000"));
}

#[test]
fn sglang_entrypoint_model_path_order() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    let ep = cmd.find("sglang serve").unwrap();
    let mp = cmd.find("--model-path").unwrap();
    let mdl = cmd.find("org/sgl-model").unwrap();
    let tp = cmd.find("--tp 4").unwrap();
    assert!(
        ep < mp && mp < mdl,
        "entrypoint before model-path before model"
    );
    assert!(mdl < tp, "model before --tp");
    let port = cmd.find("--port 8000").unwrap();
    assert!(tp < port, "--tp before --port");
}

#[test]
fn keep_disables_rm() {
    let cmd = run_keep(&trec(), &tcls().nodes[0], true);
    assert!(!cmd.contains("--rm"), "--keep should suppress --rm");
    assert!(cmd.contains("docker run -d"));
    assert!(cmd.contains("--gpus all"));
    assert!(cmd.contains("-e HF_HOME=/root/.cache/huggingface"));
    assert!(cmd.contains("-v /home/user/.cache/huggingface:/root/.cache/huggingface"));
}

#[test]
fn recipe_env_overrides_cluster_defaults() {
    // Recipe-level env should win over the defaults we inject.
    let mut r = trec();
    r.env.insert("NCCL_CUMEM_ENABLE".into(), "1".into());
    r.env.insert("OMP_NUM_THREADS".into(), "16".into());
    r.env.insert("NCCL_IGNORE_CPU_AFFINITY".into(), "0".into());
    r.env.insert("NCCL_NET".into(), "SOCKET".into());
    r.env.insert("NCCL_CROSS_NIC".into(), "0".into());
    r.env.insert("NCCL_DEBUG".into(), "INFO".into());
    let cmd = run(&r, &tcls().nodes[0], true);
    // Docker uses the last -e for duplicate keys, so recipe values come after.
    // rfind returns the last occurrence, which should be the recipe override.
    let cumem_pos = cmd.rfind("NCCL_CUMEM_ENABLE=").unwrap();
    assert!(
        cmd[cumem_pos..].starts_with("NCCL_CUMEM_ENABLE=1"),
        "recipe override should be last: slice starts at {}",
        cumem_pos
    );
    let omp_pos = cmd.rfind("OMP_NUM_THREADS=").unwrap();
    assert!(
        cmd[omp_pos..].starts_with("OMP_NUM_THREADS=16"),
        "recipe override should be last"
    );
    // New NCCL defaults can also be overridden.
    let ignore_pos = cmd.rfind("NCCL_IGNORE_CPU_AFFINITY=").unwrap();
    assert!(
        cmd[ignore_pos..].starts_with("NCCL_IGNORE_CPU_AFFINITY=0"),
        "NCCL_IGNORE_CPU_AFFINITY should be overridable"
    );
    let net_pos = cmd.rfind("NCCL_NET=").unwrap();
    assert!(
        cmd[net_pos..].starts_with("NCCL_NET=SOCKET"),
        "NCCL_NET should be overridable"
    );
    let cross_pos = cmd.rfind("NCCL_CROSS_NIC=").unwrap();
    assert!(
        cmd[cross_pos..].starts_with("NCCL_CROSS_NIC=0"),
        "NCCL_CROSS_NIC should be overridable"
    );
    let debug_pos = cmd.rfind("NCCL_DEBUG=").unwrap();
    assert!(
        cmd[debug_pos..].starts_with("NCCL_DEBUG=INFO"),
        "NCCL_DEBUG should be overridable"
    );
}

#[test]
fn custom_init_port() {
    // Non-default init_port should propagate to vLLM master-port and
    // SGLang dist-init-addr.
    let r = trec();
    let tp = if r.required_nodes == Some(r.tensor_parallel_size) {
        2
    } else {
        r.tensor_parallel_size
    };
    let cmd = build_docker_command(
        &tcls(),
        &tnccl(),
        &r,
        "r",
        &tcls().nodes[0],
        true,
        2,
        tp,
        2,
        "/home/user/.cache/huggingface",
        false,
        29500,
    )
    .join(" ");
    assert!(cmd.contains("--master-port 29500"));
    assert!(cmd.contains("--master-addr 1"));
}

// ── Logs command format ───────────────────────────────────────

#[test]
fn logs_tail_cmd_format() {
    let cmd = super::logs_tail_cmd("test-sleep", 0, 100);
    assert!(cmd.contains("docker ps -q"));
    assert!(cmd.contains("--filter \"label=sm121.recipe=test-sleep\""));
    assert!(cmd.contains("--filter \"label=sm121.rank=0\""));
    assert!(cmd.contains("xargs -r docker logs"));
    assert!(cmd.contains("--timestamps"));
    assert!(cmd.contains("--tail=100"));
    assert!(cmd.ends_with(" 2>&1"));
}

#[test]
fn logs_tail_cmd_per_rank() {
    let cmd0 = super::logs_tail_cmd("glm51", 0, 100);
    let cmd2 = super::logs_tail_cmd("glm51", 2, 100);
    assert!(cmd0.contains("label=sm121.rank=0"));
    assert!(cmd2.contains("label=sm121.rank=2"));
    assert!(!cmd0.contains("label=sm121.rank=2"));
}

#[test]
fn logs_follow_cmd_format() {
    let cmd = super::logs_follow_cmd("test-sleep");
    assert!(cmd.contains("xargs -r docker logs -f"));
    assert!(cmd.contains("--timestamps"));
    // Follow targets the head node only (rank=0).
    assert!(cmd.contains("--filter \"label=sm121.rank=0\""));
    // Follow should not include --tail (we want a live stream from now).
    assert!(!cmd.contains("--tail"));
    // Follow must redirect stderr so docker's own errors are captured
    // (mirrors the behaviour of logs_tail_cmd).
    assert!(
        cmd.contains("2>&1"),
        "follow cmd must capture stderr, got: {cmd}"
    );
}

// ── Shell quoting ─────────────────────────────────────────────

#[test]
fn shq_wraps_simple_strings() {
    assert_eq!(shq("foo"), "'foo'");
    assert_eq!(shq(""), "''");
    assert_eq!(shq("foo bar"), "'foo bar'");
}

#[test]
fn shq_escapes_inner_single_quotes() {
    // "don't" → 'don'\''t'
    assert_eq!(shq("don't"), "'don'\\''t'");
}

#[test]
fn entrypoint_script_is_shell_quoted() {
    // sh -c SCRIPT should wrap the script in single quotes so it survives
    // bash word-splitting on the SSH server.
    let mut r = trec();
    r.entrypoint = vec!["sh".into(), "-c".into(), "echo $i; sleep 1".into()];
    let cmd = run(&r, &tcls().nodes[0], true);
    // The script element should appear wrapped in single quotes
    assert!(
        cmd.contains("'echo $i; sleep 1'"),
        "script should be shell-quoted, got: {cmd}"
    );
}

#[test]
fn entrypoint_script_non_sh_passes_through() {
    // For non `sh -c` entrypoints (e.g. ["sh", "-c", "...", "extra"]) or
    // direct binaries, we leave the args alone — no false positive shq.
    let mut r = trec();
    r.entrypoint = vec!["sh".into(), "-c".into()];
    let cmd = run(&r, &tcls().nodes[0], true);
    assert!(cmd.contains("docker run"));
}