pub mod build;
pub mod command;
pub mod compose;
pub mod deploy;
pub mod health;
pub(crate) mod launch_node;
pub mod logs;
pub mod mirror;
pub mod mods;
pub(crate) mod path_ref;
pub mod push;
pub use logs::show_logs;
pub mod build_dir;
pub mod stale;
pub mod stop;
pub mod sync;
pub mod tail;
use crate::{
    config::{Cluster, Config, NcclConfig, Node, Recipe, load_recipe},
    doctor::try_repaired_image,
    ssh::run_ssh,
};
use anyhow::{Context, Result};
use tokio::time::{Duration, sleep};
use tracing::{info, warn};
/// Options for the serve operation.
#[derive(Clone, Debug)]
pub struct ServeOpts {
    pub keep: bool,
    pub skip_before_serve: bool,
    pub dry_run: bool,
    /// Skip the engine-readiness wait after launch (`--no-wait`).
    pub no_wait_ready: bool,
    pub auto_power_on: bool,
    pub sync_image: bool,
    pub sync_image_source: Option<String>,
    /// Skip the post-build registry push even when `[registry]` is set.
    pub no_push: bool,
    pub build_node: Option<String>,
    pub force_build: bool,
    /// Ignore occupied nodes during cluster selection (`--force`).
    pub force: bool,
    /// CLI override for recipe `solo`: `Some(true)` (--solo), `Some(false)`
    /// (--no-solo), `None` (unset → use recipe `solo` field).
    pub solo: Option<bool>,
    /// Mirror the model into a per-node local folder and mount that instead
    /// of the shared HF cache. `Some("")` = bare `--mirror` → default
    /// `~/.cache/hf-models`; `Some(path)` = that exact path; `None` = off.
    pub mirror: Option<String>,
    /// CLI override for the host-side HF cache dir (`--hf-cache`); wins over
    /// the config `hf_cache_dir` default. Mutually exclusive with `mirror`
    /// (both pick the HF mount source).
    pub hf_cache: Option<String>,
    /// CLI override for recipe scaling: `Some("")` = bare `--fit` → scale to
    /// the whole selected cluster; `Some("N")` = scale to N nodes;
    /// `None` = unset. See `docker::serve`.
    pub fit: Option<String>,
    pub clear: bool,
}
/// Shell-quote for remote `sh -c` (single-quote wrapper, escapes quotes).
/// Re-exported so docker submodules can `use super::shq;`; the impl lives
/// in [`crate::ssh`].
pub(crate) use crate::ssh::shq;
pub async fn serve(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    node_name: Option<&str>,
    opts: ServeOpts,
) -> Result<Cluster> {
    let mut recipe = load_recipe(recipe_name, &config.recipe_sources).await?;
    let rn = recipe.name.clone();
    let re = recipe.engine.clone();
    // CLI `--cluster auto` (the default) falls back to the recipe's `cluster`
    // field — same behavior as the tmux path.
    let cluster_name = if cluster_name == "auto" {
        recipe.cluster.clone().unwrap_or_else(|| "auto".to_string())
    } else {
        cluster_name.to_string()
    };
    info!(
        "{}",
        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
    );
    // `--fit`: override required_nodes + tensor_parallel_size so TP tracks
    // the node count (the recipe's `required_nodes == Some(tp_size)` path in
    // `start_cluster` makes tp_size follow nnodes). `pipeline_parallel_size`
    // is left untouched. Bare `--fit` (whole cluster) is handled after
    // cluster selection, when nnodes is known; for now just clear the cap so
    // selection/trim keeps every cluster node.
    //
    // Mutually exclusive with `--solo` (solo = 1 node, no dist flags).
    let fit: Option<u32> = match opts.fit.as_deref() {
        None => None,
        Some("") => Some(0),
        Some(s) => Some(
            s.parse()
                .context(format!("--fit expects a node count, got '{}'", s))?,
        ),
    };
    if fit.is_some() && opts.solo.is_some() {
        anyhow::bail!("--fit and --solo are mutually exclusive");
    }
    match fit {
        Some(0) => recipe.required_nodes = None,
        Some(n) => {
            recipe.required_nodes = Some(n);
            recipe.tensor_parallel_size = n;
            info!("--fit: scaling to {} nodes (tp = {})", n, n);
        }
        None => {}
    }
    // Prepare cluster: probe, auto-power-on if requested, select.
    use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
    use crate::orchestrate::{BootTiming, PrepareRequest, prepare_cluster};

    let power_ctrl = opts
        .auto_power_on
        .then(|| ha_controller_from_config(config))
        .flatten();
    // Effective solo: CLI flag overrides recipe field. When solo and no
    // explicit --node, launch on exactly one node (the first free node in
    // the smallest cluster that has one) with no distributed flags. When
    // solo and --node is given, the explicit node wins but we still run
    // single-node (start_cluster sees one node → no distributed flags).
    let effective_solo = opts.solo.unwrap_or(recipe.solo.unwrap_or(false));
    // For solo, ignore the recipe's required_nodes — we always want one node.
    // Build a shallow clone with required_nodes overridden so cluster
    // selection picks the smallest cluster with >= 1 free node.
    let recipe_for_select = if effective_solo {
        let mut r = recipe.clone();
        r.required_nodes = Some(1);
        r
    } else {
        recipe.clone()
    };
    let (mut cluster, nccl) = if opts.dry_run {
        // Dry-run: skip all SSH probing/power/pinging so it works even when the
        // cluster is off or unreachable (treat every configured node as free).
        crate::orchestrate::prepare_cluster_dry_run(
            config,
            &recipe_for_select,
            &cluster_name,
            node_name,
        )?
    } else {
        prepare_cluster(
            config,
            &recipe_for_select,
            &cluster_name,
            node_name,
            &PrepareRequest {
                prober: &SshProber,
                power: power_ctrl.as_ref(),
                pinger: &SshNodePinger,
                auto_power_on: opts.auto_power_on,
                boot_timing: BootTiming::default(),
                force_occupied: opts.force,
            },
        )
        .await?
    };
    if effective_solo && node_name.is_none() && !cluster.nodes.is_empty() {
        // Solo: keep only the first node. `prepare_cluster` already selected
        // the smallest cluster with >= required_nodes free (required_nodes is
        // ignored for solo — we always want one node), so the first free node
        // is the solo target.
        let solo_node = cluster.nodes[0].clone();
        info!(
            "{}",
            fluent_i18n::t!("solo-node-selected", { "name" => &solo_node.name, "cluster" => &cluster.name })
        );
        cluster.nodes = vec![solo_node];
    }
    // Bare `--fit` (whole cluster): now that `cluster.nodes` is resolved, fix
    // tp to the node count so `start_cluster`'s `tp_size` auto-track fires
    // (required_nodes == Some(tp_size) → tp_size = nnodes). required_nodes is
    // already None (no trim), so every node is kept.
    if fit == Some(0) {
        let nn = cluster.nodes.len() as u32;
        recipe.required_nodes = Some(nn);
        recipe.tensor_parallel_size = nn;
        info!(
            "--fit: scaling to whole cluster '{}' ({} nodes, tp = {})",
            cluster.name, nn, nn
        );
    }
    crate::before_serve::run(
        config,
        &cluster,
        &recipe,
        opts.skip_before_serve,
        opts.dry_run,
    )
    .await?;

    // `built_on` records the node that (re)built the image this serve: after
    // a fresh build that image is authoritative — the sync step below must
    // pull FROM it, never let a peer's older image (by creation timestamp)
    // overwrite the new build.
    let mut built_on: Option<String> = None;
    if let Some(ref build_node_name) = opts.build_node {
        if !opts.dry_run && !cluster.nodes.iter().any(|n| n.name == *build_node_name) {
            anyhow::bail!("build-node '{}' not in cluster", build_node_name);
        }
        if !opts.dry_run {
            let did_build = crate::docker::build::handle_build(
                config,
                &cluster,
                recipe_name,
                &recipe,
                &opts,
                build_node_name,
            )
            .await?;
            if did_build {
                built_on = Some(build_node_name.clone());
            }
        }
    }

    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?;
    if image != recipe.docker_image {
        recipe.docker_image = image.clone();
    }
    // Surface image staleness (bumped Dockerfile vs image checksum label)
    // when not building. Non-blocking; runs on dry-run too (read-only inspect).
    if opts.build_node.is_none() {
        stale::warn_if_stale(config, &cluster, &recipe).await?;
    }
    // Sync Docker image across nodes if needed.
    let should_sync =
        opts.sync_image && recipe.sync_image && !opts.dry_run && cluster.nodes.len() > 1;
    if should_sync {
        // After a fresh build the build node holds the authoritative image:
        // force it as the sync source so a peer's older image (by creation
        // timestamp, or from a no-op `docker load` round-trip) can't overwrite
        // the new build. The "most recently created" heuristic only applies
        // when nothing was built here and no explicit --sync-image-source was
        // given — i.e. the nodes merely disagree and we pick the newest.
        let source = built_on.as_deref().or(opts.sync_image_source.as_deref());
        self::sync::sync_image_to_nodes(config, &cluster, &recipe.docker_image, source).await?;
    }

    // Drop page cache on all nodes right before starting containers.
    // Must run after image sync (docker load fills page cache).
    if config.drop_caches && !opts.dry_run {
        info!("Dropping page cache on all nodes...");
        for node in &cluster.nodes {
            let cmd = "echo 3 | sudo tee /proc/sys/vm/drop_caches".to_string();
            if let Err(e) = run_ssh(config, node, cmd, false).await {
                warn!("drop_caches on {} failed: {}", node.name, e);
            }
        }
    }

    // `--mirror`: mount a per-node local mirror of the HF cache instead of
    // the shared one (first serve copies the model per node; later serves
    // load it locally). See mirror.rs.
    if opts.mirror.is_some() && opts.hf_cache.is_some() {
        anyhow::bail!(
            "--hf-cache and --mirror are mutually exclusive (both pick the HF mount source)"
        );
    }
    let (mirror, cache_dir) =
        mirror::resolve(config, opts.mirror.as_deref(), opts.hf_cache.as_deref());
    if let Some(ref m) = mirror {
        info!(
            "serving from local HF mirror {} (models copied per node on first use)",
            m.display()
        );
    }
    start_cluster(
        config,
        &cluster,
        nccl,
        recipe_name,
        &recipe,
        StartClusterOpts {
            cache_dir,
            mirror,
            keep: opts.keep,
            node_name: node_name.map(String::from),
            mods_remote_path,
            dry_run: opts.dry_run,
            clear: opts.clear,
        },
    )
    .await?;

    // Launch ledger + readiness: the run is recorded at launch time; then
    // the engine must actually come up (bench-style /v1/models probe). A
    // run that never reaches readiness is closed with its uptime and the
    // serve fails loudly instead of returning a half-started serve.
    if !opts.dry_run {
        crate::runlock::mark_start(recipe_name, &config.recipe_sources, Some(&cluster.name))?;
        if !opts.no_wait_ready
            && let Err(e) = crate::readiness::wait_engine(
                &cluster,
                recipe.port,
                crate::readiness::SERVE_WAIT_SECS,
            )
            .await
        {
            let _ = crate::runlock::mark_stop(recipe_name, &config.recipe_sources);
            return Err(e);
        }
    }
    Ok(cluster)
}
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 {} \
             --format '{{{{.Label \"sm121.cluster\"}}}}' | head -1",
            shq(&format!("label=sm121.recipe={}", 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,
    /// When set, the per-node mirror dir mounted at /root/.cache/huggingface
    /// (repos are copied into it from the shared cache before launch).
    mirror: Option<std::path::PathBuf>,
    keep: bool,
    node_name: Option<String>,
    mods_remote_path: String,
    dry_run: bool,
    clear: bool,
}
async fn start_cluster(
    config: &Config,
    cluster: &Cluster,
    nccl: Option<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 {
        // Trim to `required_nodes` when set: a recipe that needs N nodes
        // must launch on exactly N nodes, not on every free node in the
        // cluster. Previously a `required_nodes = 1` recipe on a 4-node
        // cluster launched 4 containers, each told `--nnodes 4 --tp 4` —
        // a broken rendezvous that never coordinated. `prepare_cluster`
        // guarantees >= required_nodes are free, so taking the first N is
        // safe and deterministic.
        let all = cluster.nodes.iter().collect::<Vec<&Node>>();
        match recipe.required_nodes {
            Some(n) if (n as usize) < all.len() => all.into_iter().take(n as usize).collect(),
            _ => all,
        }
    };
    let nnodes = nodes.len() as u32;
    // `--mirror`: ensure every repo the recipe needs exists in the per-node
    // mirror before launch (copies from the shared HF cache when missing).
    if let Some(ref mirror) = opts.mirror {
        mirror::ensure_mirrors(
            config,
            &nodes,
            recipe,
            &config.resolved_hf_cache_dir(),
            mirror,
            opts.dry_run,
        )
        .await?;
    }
    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
    };
    // Resolve static mounts once (alias/host + tilde). Model path references
    // (`repo`+`subfolder`) are kept as specs and resolved per node inside the
    // spawn loop, since the HF snapshot hash can differ across nodes.
    let resolved = command::resolve_extra_mounts(config, recipe)?;
    command::clear_extra_mounts(config, &nodes, &resolved, opts.clear, opts.dry_run).await?;
    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 {} 2>/dev/null; true",
            shq(&format!("sm121_{}_{}", recipe_name, node.rank))
        );
        let dm_ctx = command::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(),
            extra_mounts: resolved.static_mounts.clone(),
        };
        let hf_cache = opts.cache_dir.clone();
        let path_refs = resolved.path_refs.clone();
        handles.push(crate::locale::spawn_with_locale(async move {
            launch_node::launch_node(
                &config,
                launch_node::LaunchNodeCtx {
                    node: &node,
                    rn: &rn,
                    dry,
                    engine: &engine,
                    is_head,
                    cleanup,
                    dm_ctx: &dm_ctx,
                    hf_cache: &hf_cache,
                    path_refs: &path_refs,
                },
            )
            .await
        }));
    }
    // NOTE: all node launches are awaited (join_all) before reporting — a
    // failure on one node does not cancel the others, so a partial cluster
    // may be left running. Intentional: reporting every node's result is
    // more useful than aborting mid-launch and leaving the operator to
    // guess which nodes actually started.
    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.effective_master_ip(), "port" => recipe.port })
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests;