//! Per-node RoCE GID index auto-detection.
//!
//! `NCCL_IB_GID_INDEX` re-numbers across reboots on DGX Spark (the RoCE-v2
//! IPv4 GID's index is not stable), so pinning it in config can go stale.
//! When the cluster's `[nccl]` block omits `ib_gid_index`, the launcher
//! resolves it per node at serve time by scanning the node's sysfs IB GID
//! table for the first index whose type is "RoCE v2" over an IPv4-mapped GID
//! (mirrors the Aiden/0rand serving stack's `compose.head.yaml` gid auto-detect:
//! https://forums.developer.nvidia.com/t/372268/708). Explicit
//! `ib_gid_index = <n>` in config opts out (pins the value).

use crate::config::{Config, NcclConfig, Node};
use crate::ssh::{run_ssh, shq};
use anyhow::{Context, Result};

/// Build the shell that scans a node's IB GID table for the RoCE-v2 IPv4 GID
/// index. Mirrors the reference compose head/worker launch script: for each
/// HCA in `NCCL_IB_HCA`, walk sysfs `gid_attrs/types` + `gids` up to index 15
/// and pick the first index whose type is "RoCE v2" and whose GID is the
/// IPv4-mapped form (`0000:0000:0000:0000:0000:ffff:*`).
pub(crate) fn detect_gid_cmd(hcas: &str) -> String {
    // `hcas` is a comma-separated HCA list from config (trusted, but shell
    // quote it anyway so a crafted value cannot inject).
    let hcas_q = shq(hcas);
    format!(
        "for HCA in $(echo {hcas_q} | tr ',' ' '); do \
           for i in $(seq 0 15); do \
             t=$(cat /sys/class/infiniband/$HCA/ports/1/gid_attrs/types/$i 2>/dev/null || true); \
             g=$(cat /sys/class/infiniband/$HCA/ports/1/gids/$i 2>/dev/null || true); \
             case \"$t\" in *\"RoCE v2\"*) case \"$g\" in *0000:0000:0000:0000:0000:ffff:*) \
               echo \"$i\"; exit 0;; esac;; esac; \
           done; \
         done; \
         exit 1"
    )
}

/// Auto-detect the RoCE GID index on `node` by SSHing in and reading its
/// sysfs GID table. On failure returns an error naming the HCA list and
/// pointing at the explicit-`ib_gid_index` opt-out.
pub(crate) async fn detect_gid_index_on_node(
    config: &Config,
    node: &Node,
    hcas: &str,
) -> Result<u32> {
    let out = run_ssh(config, node, detect_gid_cmd(hcas), false)
        .await
        .context("NCCL_IB_GID_INDEX auto-detect failed")?;
    let idx = out.trim().parse::<u32>().with_context(|| {
        format!(
            "Could not determine a RoCE-v2 IPv4 GID index on '{}' (HCA(s) {}). \
             Set `ib_gid_index = <n>` explicitly in the cluster's [nccl] section \
             to opt out of auto-detection.",
            node.name, hcas
        )
    })?;
    if idx > 15 {
        anyhow::bail!(
            "Auto-detected NCCL_IB_GID_INDEX={idx} on '{}' is out of range",
            node.name
        );
    }
    Ok(idx)
}

/// Resolve a node's NCCL config for launch: if `ib_gid_index` is pinned in
/// config, return it unchanged; if `None` and `dry_run`, print the detect
/// command and return a config keeps `ib_gid_index = None` (the docker env
/// renders a placeholder); if `None` and a real run, SSH-detect per node and
/// pin the resolved value.
pub(crate) async fn resolve_nccl_for_node(
    config: &Config,
    node: &Node,
    nccl: Option<&NcclConfig>,
    dry_run: bool,
) -> Result<Option<NcclConfig>> {
    let Some(nccl) = nccl else {
        // Single-node recipe — no NCCL config at all.
        return Ok(None);
    };
    if nccl.ib_gid_index.is_some() {
        // Pinned (opt-out) — use as-is.
        return Ok(Some(nccl.clone()));
    }
    if dry_run {
        // No SSH in dry-run; show the detect command and leave the value
        // unresolved so the printed `-e NCCL_IB_GID_INDEX=...` is a
        // self-explanatory placeholder.
        println!(
            "  ssh {}@{}: {}",
            config.ssh_user,
            node.ssh_ip,
            detect_gid_cmd(&nccl.ib_hca)
        );
        return Ok(Some(nccl.clone()));
    }
    let idx = detect_gid_index_on_node(config, node, &nccl.ib_hca).await?;
    let mut resolved = nccl.clone();
    resolved.ib_gid_index = Some(idx);
    Ok(Some(resolved))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::NcclConfig;

    #[test]
    fn detect_cmd_scans_roce_v2_ipv4() {
        let cmd = detect_gid_cmd("rocep1s0f1,roceP2p1s0f1");
        // HCA list is shell-quoted; both HCAs are iterated in order.
        assert!(cmd.contains("'rocep1s0f1,roceP2p1s0f1'"));
        assert!(cmd.contains("gid_attrs/types"));
        assert!(cmd.contains("0000:0000:0000:0000:0000:ffff:"));
        assert!(cmd.contains("RoCE v2"));
        assert!(cmd.contains("seq 0 15"));
        // Must fail loudly when nothing matches.
        assert!(cmd.ends_with("exit 1"));
    }

    #[tokio::test]
    async fn resolve_returns_pinned_config_unchanged() {
        // A pinned (opt-out) `ib_gid_index` must short-circuit with no SSH —
        // `detect_gid_index_on_node` is never reached, so this works even
        // with an unreachable node.
        let nccl = NcclConfig {
            ib_hca: "r".into(),
            socket_ifname: "e".into(),
            ib_gid_index: Some(3),
            ib_disable: 0,
            gloo_socket_ifname: None,
            ucx_net_devices: None,
            nccl_cross_nic: None,
            nccl_proto: None,
            nccl_p2p_level: None,
        };
        let node = crate::config::Node {
            name: "n1".into(),
            ssh_ip: "10.0.0.1".into(),
            ib_ip: Some("10.0.1.1".into()),
            rank: 0,
            host_key: None,
            entity: None,
        };
        let config = crate::config::Config {
            ssh_user: "u".into(),
            ssh_key_path: std::path::PathBuf::from("/p"),
            clusters: vec![],
            nccl: Default::default(),
            hf_cache_dir: None,
            init_port: 30000,
            before_serve: Default::default(),
            exec_aliases: Default::default(),
            home_assistant: None,
            mikrotik: None,
            global_mods_dir: None,
            drop_caches: false,
            mount_aliases: Default::default(),
            recipe_sources: Vec::new(),
            docker_image_sources: Vec::new(),
            cluster_sources: Vec::new(),
            registry: None,
        };
        for dry in [false, true] {
            let r = resolve_nccl_for_node(&config, &node, Some(&nccl), dry)
                .await
                .unwrap()
                .unwrap();
            assert_eq!(r.ib_gid_index, Some(3));
        }
    }
}