launcher/src/docker/health.rs

//! Cluster health — shows node status, running containers, and recipe
//! deployments across all clusters.

use crate::config::{Cluster, Config};
use crate::home_assistant::Client as HaClient;
use crate::ssh::run_ssh;
use anyhow::Result;
use std::collections::HashMap;

/// Per-node result from the parallel SSH sweep.
struct NodeHealth {
    name: String,
    ib_ip: String,
    online: bool,
    containers: Vec<String>,
    net_up: bool,
    net_details: Vec<String>,
}

/// Checks Docker health — shows node status and containers belonging to this cluster.
pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) -> Result<()> {
    let nodes = &cluster.nodes;
    let cluster_label = &cluster.name;
    let nccl = config.nccl.get(cluster_label);

    // Check HA cluster-level entities (switch, etc.) in parallel.
    let ha_handle = if let Some(ref ha_cfg) = config.home_assistant {
        let ha = HaClient::new(ha_cfg).ok();
        let entities = cluster.extra_entities.clone();
        Some(tokio::spawn(async move {
            let client = ha.as_ref()?;
            for e in &entities {
                match client.get_state(e).await {
                    Some(ref s) if s == "on" => continue,
                    _ => return Some(false),
                }
            }
            Some(true)
        }))
    } else {
        None
    };

    // Build the network check command snippet for this cluster's NCCL interfaces.
    let net_script = nccl.map(|n| {
        let socket = &n.socket_ifname;
        format!(
            "echo '---NET---'; echo -n 'socket:' && (ip -o link show dev {socket} 2>/dev/null | grep -qi 'state UP' && echo UP || echo DOWN); echo '---END---'"
        )
    });

    // SSH into each node in parallel.
    let mut handles = Vec::new();
    for node in nodes {
        let config = config.clone();
        let node = node.clone();
        let name = node.name.clone();
        let ib_ip = node.ib_ip.clone();
        let cl = cluster_label.clone();
        let ns = net_script.clone();
        let docker_cmd = format!(
            "docker ps --filter label=sm121.cluster={cl} \
             --format '{{{{.Names}}}}\t{{{{.Status}}}}\t{{{{.Image}}}}' 2>&1"
        );
        let cmd = if let Some(ref script) = ns {
            format!("{docker_cmd}; {script}")
        } else {
            docker_cmd
        };
        handles.push(tokio::spawn(async move {
            let result = run_ssh(&config, &node, cmd, dry_run).await;
            (name, ib_ip, result)
        }));
    }

    // Collect SSH node results.
    let mut node_health: Vec<NodeHealth> = Vec::new();
    for h in futures::future::join_all(handles).await {
        match h {
            Ok((name, ib_ip, Ok(output))) => {
                let mut net_up = true;
                let mut net_details = Vec::new();
                let mut lines: Vec<String> = Vec::new();
                let mut in_net = false;
                for line in output.lines() {
                    let line = line.trim();
                    if line.is_empty()
                        || line.contains("permission denied")
                        || line.starts_with("docker:")
                    {
                        continue;
                    }
                    if line == "---NET---" {
                        in_net = true;
                        continue;
                    }
                    if line == "---END---" {
                        in_net = false;
                        continue;
                    }
                    if in_net {
                        net_details.push(line.to_string());
                        if line.ends_with("DOWN")
                            || line.ends_with("missing")
                            || line.ends_with("error")
                        {
                            net_up = false;
                        }
                        continue;
                    }
                    lines.push(line.to_string());
                }
                node_health.push(NodeHealth {
                    name,
                    ib_ip,
                    online: true,
                    containers: lines,
                    net_up,
                    net_details,
                });
            }
            Ok((name, ib_ip, Err(_))) => {
                node_health.push(NodeHealth {
                    name,
                    ib_ip,
                    online: false,
                    containers: vec![],
                    net_up: false,
                    net_details: vec![],
                });
            }
            Err(join_err) => eprintln!("task join: {}", join_err),
        }
    }

    // Resolve HA result.
    let ha_ok = match ha_handle {
        Some(h) => h.await.unwrap_or(Some(true)).unwrap_or(true),
        None => true,
    };

    let total = nodes.len();
    let online = node_health.iter().filter(|n| n.online).count();
    let offline = total - online;
    let container_count: usize = node_health.iter().map(|n| n.containers.len()).sum();
    let net_ok = node_health.iter().all(|n| !n.online || n.net_up);

    print!("=== {} === {} nodes", cluster_label, total);
    if !ha_ok && !cluster.extra_entities.is_empty() {
        print!(" 🚫 switch off");
    } else if !net_ok {
        print!(" 🌐 interface down");
    } else if online == total {
        print!(" ✅ all online");
    } else {
        print!(" ({} online, {} offline)", online, offline);
    }
    if container_count > 0 {
        println!(", {} container(s)", container_count);
    } else {
        println!(", no containers");
    }

    for nh in &node_health {
        if !nh.online {
            eprintln!("  ✗ {} ({}) — offline", nh.name, nh.ib_ip);
            continue;
        }
        let net_status = if !nh.net_up {
            format!(" [net: {}]", nh.net_details.join(", "))
        } else {
            String::new()
        };
        if nh.containers.is_empty() {
            println!("  ○ {} ({}){}", nh.name, nh.ib_ip, net_status);
        } else {
            println!("  ● {} ({}){}", nh.name, nh.ib_ip, net_status);
            for line in &nh.containers {
                println!("    {line}");
            }
        }
    }
    Ok(())
}

/// Quick overview — recipes running on each cluster.
pub async fn health_summary(config: &Config) -> Result<()> {
    let mut seen = std::collections::HashSet::new();
    let mut handles = Vec::new();
    for cluster in &config.clusters {
        for node in &cluster.nodes {
            if seen.contains(&node.ssh_ip) {
                continue;
            }
            seen.insert(node.ssh_ip.clone());
            let config = config.clone();
            let node = node.clone();
            let name = node.name.clone();
            let cmd = "docker ps --filter label=sm121.recipe \
                       --format '{{.Label \"sm121.recipe\"}}|{{.Label \"sm121.cluster\"}}'"
                .to_string();
            handles.push(tokio::spawn(async move {
                let out = run_ssh(&config, &node, cmd, false)
                    .await
                    .unwrap_or_default();
                (name, out)
            }));
        }
    }

    let mut recipe_nodes: HashMap<String, std::collections::HashSet<String>> = HashMap::new();
    for h in futures::future::join_all(handles).await {
        let (node_name, output) = h.unwrap_or_default();
        for line in output.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let mut parts = line.splitn(2, '|');
            let recipe = parts.next().unwrap_or("").to_string();
            if recipe.is_empty() {
                continue;
            }
            recipe_nodes
                .entry(recipe)
                .or_default()
                .insert(node_name.clone());
        }
    }

    if recipe_nodes.is_empty() {
        println!("\nNo sm121 containers running.");
        return Ok(());
    }

    println!("\nRunning recipes:");
    for (recipe, nodes) in &recipe_nodes {
        let mut sorted: Vec<_> = nodes.iter().collect();
        sorted.sort();
        let node_list: Vec<&str> = sorted.iter().map(|s| s.as_str()).collect();
        println!(
            "  {} — {} node(s): {}",
            recipe,
            node_list.len(),
            node_list.join(", ")
        );
    }
    Ok(())
}