pandora_lib_promethion/src/slurm_helpers.rs

use std::collections::HashMap;
use std::process::Command;

use anyhow::Context;

#[derive(Debug)]
pub struct NodeInfo {
    pub partition: String,
    pub node: String,
    pub cpu_total: i32,
    pub cpu_free: i32,
    pub gpu_total: i32,
    pub gpu_free: i32,
    pub gpu_names: Vec<String>, // <- models/types for GRES
}

fn parse_kv_line(line: &str) -> HashMap<String, String> {
    let mut map = HashMap::new();
    for token in line.split_whitespace() {
        if let Some(eq) = token.find('=') {
            let (k, v) = token.split_at(eq);
            map.insert(k.to_string(), v[1..].to_string());
        }
    }
    map
}

fn parse_gpu_from_tres(s: &str) -> u32 {
    // handles "gres/gpu=4" and "gres/gpu:V100=4"
    if let Some(idx) = s.find("gres/gpu") {
        let rest = &s[idx + "gres/gpu".len()..];
        let rest = rest.strip_prefix(':').unwrap_or(rest); // optional ":TYPE"
        let rest = rest.split('=').nth(1).unwrap_or("");
        let num: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect();
        return num.parse().unwrap_or(0);
    }
    0
}

fn parse_gpu_from_gres(s: &str) -> u32 {
    // handles "gpu:4" or "gpu:V100:4"
    if let Some(idx) = s.find("gpu:") {
        let rest = &s[idx + "gpu:".len()..];
        // could be "V100:4" or "4"
        let parts: Vec<&str> = rest.split([':', '(']).collect();
        // last numeric field is count
        for p in parts.iter().rev() {
            if let Ok(v) = p.parse() {
                return v;
            }
        }
    }
    0
}

fn extract_total_gpus(fields: &HashMap<String, String>) -> u32 {
    if let Some(cfg) = fields.get("CfgTRES") {
        let n = parse_gpu_from_tres(cfg);
        if n > 0 {
            return n;
        }
    }
    if let Some(gres) = fields.get("Gres") {
        return parse_gpu_from_gres(gres);
    }
    0
}

fn extract_alloc_gpus(fields: &HashMap<String, String>) -> u32 {
    if let Some(alloc) = fields.get("AllocTRES") {
        let n = parse_gpu_from_tres(alloc);
        if n > 0 {
            return n;
        }
    }
    if let Some(gres_used) = fields.get("GresUsed") {
        return parse_gpu_from_gres(gres_used);
    }
    0
}

fn extract_gpu_names_from_gres(s: &str) -> Vec<String> {
    // Examples:
    //   "gpu:V100:4(S:0-3)"
    //   "gpu:A100:8(S:0-7)"
    //   "gpu:4(S:0-3)"           -> no type
    let mut names = Vec::new();

    for item in s.split(',') {
        if let Some(idx) = item.find("gpu:") {
            let rest = &item[idx + "gpu:".len()..];
            // strip trailing "(S:...)" etc.
            let rest = rest.split('(').next().unwrap_or(rest);
            let parts: Vec<&str> = rest.split(':').collect();

            if parts.len() == 1 {
                // "4" -> no explicit type
                continue;
            } else {
                // "V100:4" -> type = parts[0]
                let name = parts[0].trim();
                if !name.is_empty() && !name.chars().all(|c| c.is_ascii_digit()) {
                    names.push(name.to_string());
                }
            }
        }
    }
    names.sort();
    names.dedup();
    names
}

fn extract_gpu_names_from_cfg_tres(s: &str) -> Vec<String> {
    // Example:
    //   "cpu=64,gres/gpu:V100=4,gres/gpu:A100=2"
    let mut names = Vec::new();

    for item in s.split(',') {
        if let Some(idx) = item.find("gres/gpu") {
            let rest = &item[idx + "gres/gpu".len()..];
            if let Some(rest) = rest.strip_prefix(':') {
                // rest = "V100=4"
                let name = rest.split('=').next().unwrap_or("").trim();
                if !name.is_empty() {
                    names.push(name.to_string());
                }
            }
        }
    }
    names.sort();
    names.dedup();
    names
}

fn extract_gpu_names(fields: &HashMap<String, String>) -> Vec<String> {
    // Prefer CfgTRES if it has types, else Gres
    if let Some(cfg) = fields.get("CfgTRES") {
        let names = extract_gpu_names_from_cfg_tres(cfg);
        if !names.is_empty() {
            return names;
        }
    }
    if let Some(gres) = fields.get("Gres") {
        return extract_gpu_names_from_gres(gres);
    }
    Vec::new()
}

fn parse_node(fields: HashMap<String, String>) -> NodeInfo {
    let partition = fields.get("Partitions").cloned().unwrap_or_default();
    let node = fields.get("NodeName").cloned().unwrap_or_default();

    let cpu_total: i32 = fields
        .get("CPUTot")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);
    let cpu_alloc: i32 = fields
        .get("CPUAlloc")
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);
    let cpu_free = cpu_total - cpu_alloc;

    let gpu_total = extract_total_gpus(&fields) as i32;
    let gpu_alloc = extract_alloc_gpus(&fields) as i32;
    let gpu_free = gpu_total - gpu_alloc;

    let gpu_names = extract_gpu_names(&fields);

    NodeInfo {
        partition,
        node,
        cpu_total,
        cpu_free,
        gpu_total,
        gpu_free,
        gpu_names,
    }
}

pub fn slurm_availables() -> anyhow::Result<Vec<NodeInfo>> {
    let output = Command::new("scontrol")
        .args(["show", "nodes", "-o"])
        .output()
        .context("failed to run scontrol")?;

    let stdout = String::from_utf8_lossy(&output.stdout);

    let mut results = Vec::new();
    for line in stdout.lines() {
        let kv = parse_kv_line(line);
        let info = parse_node(kv);
        results.push(info);
    }

    Ok(results)
}

pub fn max_gpu_per_node(gpu_name: &str) -> Option<i32> {
    let nodes = slurm_availables().ok()?;
    nodes
        .iter()
        .filter(|n| n.gpu_names.iter().any(|g| g.eq_ignore_ascii_case(gpu_name)))
        .map(|n| n.gpu_free)
        .max()
}

/// Return the **maximum idle CPUs on a single non-drained node** in `shortq`.
pub fn max_cpus_per_task_shortq() -> anyhow::Result<u32> {
    let out = Command::new("sinfo")
        .args(["-p", "shortq", "-N", "-h", "-o", "%C %t"])
        .output()?;

    let stdout = String::from_utf8_lossy(&out.stdout);
    let mut max_idle = 0;

    for line in stdout.lines() {
        let mut parts = line.split_whitespace();
        let cpus = parts.next().unwrap_or("");
        let state = parts.next().unwrap_or("").to_lowercase();
        if state.contains("drain") {
            continue;
        }

        // CPUS(A/I/O/T) → take idle `I` at index 1
        if let Some(i) = cpus.split('/').nth(1).and_then(|v| v.parse::<u32>().ok()) {
            max_idle = max_idle.max(i);
        }
    }

    Ok(max_idle)
}

// pub fn available_cpus_shortq() -> Result<u32> {
//     use std::process::Command;
//
//     let out = Command::new("sinfo")
//         .args(["-p", "shortq", "-N", "-h", "-o", "%C %t"])
//         .output()?;
//
//     let stdout = String::from_utf8_lossy(&out.stdout);
//     let mut idle = 0;
//
//     for line in stdout.lines() {
//         let mut parts = line.split_whitespace();
//         let cpus = parts.next().unwrap_or("");
//         let state = parts.next().unwrap_or("").to_lowercase();
//         if state.contains("drain") {
//             continue;
//         }
//
//         // CPUS(A/I/O/T) → idle is index 1
//         if let Some(i) = cpus.split('/').nth(1).and_then(|v| v.parse::<u32>().ok()) {
//             idle = idle.saturating_add(i);
//         }
//     }
//
//     Ok(idle)
// }

#[cfg(test)]
mod tests {
    use crate::slurm_helpers::{max_cpus_per_task_shortq, max_gpu_per_node, slurm_availables};

    #[test]
    fn slurm_info() -> anyhow::Result<()> {
        println!(
            "{:15} {:15} {:>7} {:>8} {:>7} {:>8}  GPU_NAMES",
            "Partition", "Node", "CPU_TOT", "CPU_FREE", "GPU_TOT", "GPU_FREE"
        );

        for info in slurm_availables()? {
            let names = if info.gpu_names.is_empty() {
                "-".to_string()
            } else {
                info.gpu_names.join(",")
            };

            println!(
                "{:15} {:15} {:7} {:8} {:7} {:8}  {}",
                info.partition,
                info.node,
                info.cpu_total,
                info.cpu_free,
                info.gpu_total,
                info.gpu_free,
                names
            );
        }

        let u = max_gpu_per_node("h100");
        println!("H100 max available: {}", u.unwrap_or(0));
        let u = max_gpu_per_node("a100");
        println!("A100 max available: {}", u.unwrap_or(0));
        let u = max_cpus_per_task_shortq();
        println!("CPU max available: {}", u.unwrap_or(0));

        Ok(())
    }
}