pub mod deploy;
pub mod health;
pub mod logs;
pub mod mods;
pub use logs::show_logs;
pub mod tail;
use crate::{
    config::{Cluster, Config, NcclConfig, Node, Recipe, load_recipe},
    doctor::try_repaired_image,
    ssh::{run_ssh, run_ssh_with_retry},
};
use anyhow::{Context, Result};
use std::borrow::Cow;
use tokio::time::{Duration, sleep};
use tracing::info;
/// Options for the serve operation.
#[derive(Clone, Debug)]
pub struct ServeOpts {
    pub keep: bool,
    pub skip_before_serve: bool,
    pub dry_run: bool,
    pub auto_power_on: bool,
}
pub(crate) fn shq(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('\'');
    for c in s.chars() {
        if c == '\'' {
            out.push_str(r"'\''");
        } else {
            out.push(c);
        }
    }
    out.push('\'');
    out
}
pub async fn serve(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    node_name: Option<&str>,
    opts: ServeOpts,
) -> Result<Cluster> {
    let recipe = load_recipe(recipe_name).await?;
    let rn = recipe.name.clone();
    let re = recipe.engine.clone();
    info!(
        "{}",
        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
    );
    // Prepare cluster: probe, auto-power-on if requested, select.
    use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
    use crate::orchestrate::prepare_cluster;

    let power_ctrl = if opts.auto_power_on {
        ha_controller_from_config(config)
    } else {
        None
    };
    let (cluster, nccl) = prepare_cluster(
        config,
        &recipe,
        cluster_name,
        node_name,
        opts.auto_power_on,
        &SshProber,
        power_ctrl.as_ref(),
        &SshNodePinger,
    )
    .await?;
    crate::before_serve::run(
        config,
        &cluster,
        &recipe,
        opts.skip_before_serve,
        opts.dry_run,
    )
    .await?;
    let mods_remote_path = mods::upload_mods_to_cluster(
        config,
        &cluster,
        &recipe,
        recipe_name,
        node_name,
        opts.dry_run,
    )
    .await?;
    let image = try_repaired_image(config, &cluster, &recipe).await?;
    let mut recipe = recipe;
    if image != recipe.docker_image {
        recipe.docker_image = image.clone();
    }
    let cache_dir = config.resolved_hf_cache_dir();
    start_cluster(
        config,
        &cluster,
        &nccl,
        recipe_name,
        &recipe,
        StartClusterOpts {
            cache_dir,
            keep: opts.keep,
            node_name: node_name.map(String::from),
            mods_remote_path,
            dry_run: opts.dry_run,
        },
    )
    .await?;
    Ok(cluster)
}
pub async fn stop_cluster(
    config: &Config,
    cluster: &Cluster,
    recipe_name: &str,
    nodes: &[String],
    dry_run: bool,
) -> Result<()> {
    info!(
        "{}",
        fluent_i18n::t!("cluster-stopping", { "name" => &cluster.name })
    );
    let nodes: Vec<&Node> = if nodes.is_empty() {
        cluster.nodes.iter().collect()
    } else {
        let mut found_nodes: Vec<&Node> = Vec::new();
        let mut not_found: Vec<String> = Vec::new();
        for nn in nodes {
            if let Some(node) = cluster.nodes.iter().find(|n| n.name == *nn) {
                found_nodes.push(node);
            } else {
                not_found.push(nn.clone());
            }
        }
        if !not_found.is_empty() {
            let names = not_found.join(", ");
            anyhow::bail!(
                fluent_i18n::t!("exec-node-not-found", { "name" => &names, "cluster" => &cluster.name })
            );
        }
        found_nodes
    };
    for node in &nodes {
        run_ssh(
            config,
            node,
            format!(
                "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.cluster={}\" | xargs -r docker stop",
                recipe_name, cluster.name
            ),
            dry_run,
        )
        .await?;
    }
    info!("{}", fluent_i18n::t!("cluster-all-stopped"));
    Ok(())
}
pub async fn resolve_cluster_by_container_label<'c>(
    config: &'c Config,
    recipe_name: &str,
) -> Result<&'c Cluster> {
    // Scan all nodes for containers matching this recipe.
    for node in config.clusters.iter().flat_map(|c| &c.nodes) {
        let cmd = format!(
            "docker ps --filter \"label=sm121.recipe={}\" \
             --format '{{{{.Label \"sm121.cluster\"}}}}' | head -1",
            recipe_name
        );
        if let Ok(output) = run_ssh(config, node, cmd, false).await {
            let cluster_label = output.trim();
            if !cluster_label.is_empty() {
                // Found a container — match its cluster label against config.
                return config
                    .clusters
                    .iter()
                    .find(|c| c.name == cluster_label)
                    .context(format!(
                        "Container label says cluster '{}', but no such cluster in config",
                        cluster_label
                    ));
            }
        }
    }
    anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
}
/// Options for [start_cluster].
struct StartClusterOpts {
    cache_dir: std::path::PathBuf,
    keep: bool,
    node_name: Option<String>,
    mods_remote_path: String,
    dry_run: bool,
}
/// Shared context for [`build_docker_command`] — constructed once, cloned per node.
#[derive(Clone)]
struct DockerCommandCtx {
    cluster: Cluster,
    nccl: NcclConfig,
    recipe: Recipe,
    recipe_name: String,
    pp_size: u32,
    tp_size: u32,
    nnodes: u32,
    cache_dir: String,
    keep: bool,
    init_port: u16,
    mods_remote_path: String,
}
async fn start_cluster(
    config: &Config,
    cluster: &Cluster,
    nccl: &NcclConfig,
    recipe_name: &str,
    recipe: &Recipe,
    opts: StartClusterOpts,
) -> Result<()> {
    let nodes: Vec<&Node> = if let Some(ref nn) = opts.node_name {
        let v: Vec<&Node> = cluster.nodes.iter().filter(|n| n.name == *nn).collect();
        if v.is_empty() {
            anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
        }
        v
    } else {
        cluster.nodes.iter().collect()
    };
    let nnodes = nodes.len() as u32;
    let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0
        || recipe.required_nodes == recipe.pipeline_parallel_size
    {
        nnodes
    } else {
        recipe.pipeline_parallel_size.unwrap()
    };
    let tp_size = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
        nnodes
    } else {
        recipe.tensor_parallel_size
    };
    let mut handles = Vec::new();
    for node in &nodes {
        let node = (*node).clone();
        let config = config.clone();
        let rn = recipe_name.to_string();
        let dry = opts.dry_run;
        let engine = recipe.engine.clone();
        let is_head = node.rank == 0;
        let cleanup = format!(
            "docker rm -f sm121_{}_{} 2>/dev/null; true",
            recipe_name, node.rank
        );
        let dm_ctx = DockerCommandCtx {
            cluster: cluster.clone(),
            nccl: nccl.clone(),
            recipe: recipe.clone(),
            recipe_name: rn.clone(),
            pp_size,
            tp_size,
            nnodes,
            cache_dir: opts.cache_dir.to_string_lossy().into_owned(),
            keep: opts.keep,
            init_port: config.init_port,
            mods_remote_path: opts.mods_remote_path.clone(),
        };
        handles.push(tokio::spawn(async move {
            if !dry {
                let _ = run_ssh(&config, &node, cleanup, dry).await;
            }
            let cmd = build_docker_command(&dm_ctx, &node, is_head).join(" ");
            let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
            (node.name, node.rank, engine, result)
        }));
    }
    for result in futures::future::join_all(handles).await {
        match result {
            Ok((name, rank, engine, Ok(_))) => info!(
                "{}",
                fluent_i18n::t!("engine-started", { "engine" => engine.to_uppercase(), "name" => &name, "rank" => rank })
            ),
            Ok((name, rank, engine, Err(e))) => anyhow::bail!(
                "{}",
                fluent_i18n::t!("engine-start-failed", { "engine" => engine.to_uppercase(), "name" => &name, "rank" => rank, "error" => e.to_string() })
            ),
            Err(e) => anyhow::bail!("Task join error: {}", e),
        }
    }
    if !opts.dry_run {
        sleep(Duration::from_secs(10)).await;
        info!(
            "{}",
            fluent_i18n::t!("cluster-started", { "ip" => &cluster.master_ib_ip, "port" => recipe.port })
        );
    }
    Ok(())
}
fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
    let is_sglang = ctx.recipe.engine == "sglang";
    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());
    }
    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));
    }
    let mut env = vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
    if ctx.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));
        }
        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 ncclext: 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 {
            ncclext.push(("NCCL_CROSS_NIC", v.to_string().into()));
        }
        if let Some(ref v) = ctx.nccl.nccl_proto {
            ncclext.push(("NCCL_PROTO", v.as_str().into()));
        }
        if let Some(ref v) = ctx.nccl.nccl_p2p_level {
            ncclext.push(("NCCL_P2P_LEVEL", v.as_str().into()));
        }
        env.extend(ncclext);
    }
    if is_head {
        env.push(("VLLM_LOGGING_COLOR", "1".into()));
        env.push(("TERM", "xterm-256color".into()));
    }
    for (k, v) in &ctx.recipe.env {
        env.push((k.as_str(), v.as_str().into()));
    }
    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));
    }
    if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
        env.push((
            "RAY_ADDRESS",
            format!("{}:6379", ctx.cluster.master_ib_ip).into(),
        ));
        env.push(("RAY_PORT", "6379".into()));
        env.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
        env.push(("RAY_OVERRIDE_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
    }
    for (k, v) in env {
        cmd.push(format!("-e {k}={v}"));
    }
    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;