launcher/src/doctor.rs

use anyhow::{Context, Result};
use tracing::info;

use crate::config::*;
use crate::ssh::*;

/// Doctor check result for a single node/container.
pub struct DoctorResult {
    pub node_name: String,
    pub checks: Vec<DoctorCheck>,
}

pub(crate) struct DoctorCheck {
    name: String,
    status: DoctorStatus,
    detail: String,
}

#[derive(PartialEq)]
enum DoctorStatus {
    Ok,
    Warn,
    Fail,
}

impl std::fmt::Display for DoctorStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DoctorStatus::Ok => write!(f, "✓"),
            DoctorStatus::Warn => write!(f, "⚠"),
            DoctorStatus::Fail => write!(f, "✗"),
        }
    }
}

/// Run doctor checks on the cluster for a recipe.
pub async fn doctor(
    config: &Config,
    recipe_name: &str,
    cluster_name: &str,
    repair: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe = load_recipe(recipe_name).await?;
    println!(
        "\n🩺 Doctor: {} ({}) [backend: {}]",
        recipe.name, recipe.engine, recipe.distributed_backend
    );

    let (cluster, _nccl) = resolve_cluster(config, &recipe, cluster_name)?;

    // Check each node
    let mut all_results: Vec<DoctorResult> = Vec::new();
    let mut handles = Vec::new();

    for node in &cluster.nodes {
        let config = config.clone();
        let node = node.clone();
        let engine = recipe.engine.clone();
        let backend = recipe.distributed_backend.clone();
        let image = recipe.docker_image.clone();
        handles.push(tokio::spawn(async move {
            check_node(&config, &node, &engine, &backend, &image).await
        }));
    }

    let results = futures::future::join_all(handles).await;
    for r in results.into_iter().flatten() {
        all_results.push(r);
    }

    // Print results
    let mut has_failures = false;
    for result in &all_results {
        println!("\n--- {} ---", result.node_name);
        for check in &result.checks {
            println!("  {} {} {}", check.status, check.name, check.detail);
            if check.status == DoctorStatus::Fail {
                has_failures = true;
            }
        }
    }

    if has_failures {
        if repair {
            println!("\n🔧 Repairing image...");
            let repaired = run_cluster_repair(config, &cluster, &recipe, dry_run).await?;
            println!("✅ Repaired image: {}", repaired);
        } else {
            println!("\n⚠ Some checks failed. Run with --repair to auto-fix.");
            std::process::exit(1);
        }
    } else {
        println!("\n✅ All checks passed!");
    }

    Ok(())
}

/// Check a single node for required software.
async fn check_node(
    config: &Config,
    node: &Node,
    engine: &str,
    backend: &str,
    image: &str,
) -> DoctorResult {
    let mut checks = Vec::new();

    checks.push(node_check(config, node, "Docker", "docker --version").await);
    checks.push(node_check(config, node, "NVIDIA GPU", "nvidia-smi").await);
    checks.push(
        node_check(
            config,
            node,
            "NCCL libs",
            "ldconfig -p | grep -q nccl && echo OK",
        )
        .await,
    );

    checks
        .push(container_check(config, node, image, "Engine", format!("{} version", engine)).await);

    if backend == "ray" {
        checks.push(container_check(config, node, image, "Ray", "ray --version".to_string()).await);
    }

    checks.push(container_entrypoint_check(config, node, image).await);

    DoctorResult {
        node_name: node.name.clone(),
        checks,
    }
}

/// Run a command on the node via SSH and return a check result.
async fn node_check(config: &Config, node: &Node, name: &str, command: &str) -> DoctorCheck {
    match run_ssh(config, node, command.to_string(), false).await {
        Ok(output) if !output.trim().is_empty() => DoctorCheck {
            name: name.to_string(),
            status: DoctorStatus::Ok,
            detail: format!(
                "OK: {}",
                output.trim().lines().next().unwrap_or(output.trim())
            ),
        },
        Ok(_) => DoctorCheck {
            name: name.to_string(),
            status: DoctorStatus::Fail,
            detail: "not found".to_string(),
        },
        Err(e) => DoctorCheck {
            name: name.to_string(),
            status: DoctorStatus::Fail,
            detail: format!("error: {}", e),
        },
    }
}

/// Run a command inside a container on the node and return a check result.
async fn container_check(
    config: &Config,
    node: &Node,
    image: &str,
    name: &str,
    command: String,
) -> DoctorCheck {
    let docker_cmd = format!("docker run --rm {} {} 2>&1 | head -1", image, command);
    match run_ssh(config, node, docker_cmd, false).await {
        Ok(output) if !output.contains("not found") && !output.contains("No such file") => {
            DoctorCheck {
                name: name.to_string(),
                status: DoctorStatus::Ok,
                detail: format!("OK: {}", output.trim()),
            }
        }
        Ok(output) => DoctorCheck {
            name: name.to_string(),
            status: DoctorStatus::Fail,
            detail: format!("not found: {}", output.trim()),
        },
        Err(e) => DoctorCheck {
            name: name.to_string(),
            status: DoctorStatus::Fail,
            detail: format!("error: {}", e),
        },
    }
}

/// Check the container entrypoint.
async fn container_entrypoint_check(config: &Config, node: &Node, image: &str) -> DoctorCheck {
    let docker_cmd = format!(
        "docker inspect --format '{{{{json .Config.Entrypoint}}}}' {} 2>/dev/null || echo 'null'",
        image
    );
    match run_ssh(config, node, docker_cmd, false).await {
        Ok(output) => {
            let entrypoint = output.trim();
            if entrypoint == "null" || entrypoint == "[]" {
                DoctorCheck {
                    name: "Entrypoint".to_string(),
                    status: DoctorStatus::Ok,
                    detail: "default (shell) entrypoint".to_string(),
                }
            } else {
                DoctorCheck {
                    name: "Entrypoint".to_string(),
                    status: DoctorStatus::Warn,
                    detail: format!("custom: {}", entrypoint),
                }
            }
        }
        Err(e) => DoctorCheck {
            name: "Entrypoint".to_string(),
            status: DoctorStatus::Fail,
            detail: format!("error: {}", e),
        },
    }
}

/// Generate Dockerfile and build repaired image on head node.
async fn run_cluster_repair(
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
    dry_run: bool,
) -> Result<String> {
    let head_node = cluster.nodes.first().context("No nodes in cluster")?;
    let repaired_tag = format!("{}.repaired", recipe.docker_image);
    let dockerfile = generate_repair_dockerfile(&repaired_tag, recipe);

    let dockerfile_path = format!("recipes/{}/Dockerfile.repaired", "temp");
    tokio::fs::write(&dockerfile_path, &dockerfile)
        .await
        .context("Failed to write Dockerfile")?;

    let create_dockerfile = format!(
        "mkdir -p /tmp/sm121_docker && cat > /tmp/sm121_docker/Dockerfile << 'DOCKERFILE_END'\n{}DOCKERFILE_END",
        dockerfile
    );

    if !dry_run {
        run_ssh(config, head_node, create_dockerfile, false).await?;

        let build_cmd = format!("cd /tmp/sm121_docker && docker build -t {} .", repaired_tag);
        let output = run_ssh_with_retry(config, head_node, build_cmd, false, 1)
            .await
            .context("Failed to build repaired image")?;
        println!("Build output: {}", output.lines().last().unwrap_or(&output));
    } else {
        println!(
            "  (dry-run) Would build {} on {}",
            repaired_tag, head_node.name
        );
    }

    tokio::fs::remove_file(&dockerfile_path).await.ok();
    Ok(repaired_tag)
}

/// Generate a Dockerfile to repair a container.
fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> String {
    let mut lines = vec![format!("FROM {}", base_image)];

    if recipe.distributed_backend == "ray" {
        lines.push("RUN pip install ray[default] || pip install ray".to_string());
    }

    let engine_install = match recipe.engine.as_str() {
        "vllm" => "pip install vllm",
        "sglang" => "pip install sglang",
        _ => "",
    };
    if !engine_install.is_empty() {
        lines.push(format!("RUN {} || true", engine_install));
    }

    lines.push("ENTRYPOINT [\"/bin/bash\", \"-c\"]".to_string());
    lines.join("\n")
}

/// Check if a repaired image exists on the head node.
pub async fn try_repaired_image(
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
) -> Result<String> {
    let head_node = cluster.nodes.first().context("No nodes in cluster")?;
    let repaired_tag = format!("{}.repaired", recipe.docker_image);

    let check_cmd = format!("docker images -q {} | head -1", repaired_tag);

    match run_ssh(config, head_node, check_cmd, false).await {
        Ok(output) if !output.trim().is_empty() => {
            info!("🩺 Found repaired image: {}", repaired_tag);
            Ok(repaired_tag)
        }
        _ => Ok(recipe.docker_image.clone()),
    }
}

// ====================== Tests ======================

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn test_recipe() -> Recipe {
        Recipe {
            engine: "vllm".to_string(),
            name: "Test Model".to_string(),
            docker_image: "test/image:latest".to_string(),
            model: "test/model".to_string(),
            port: 8000,
            required_nodes: Some(2),
            pipeline_parallel_size: Some(2),
            tensor_parallel_size: 1,
            args: vec!["--max-model-len".to_string(), "4096".to_string()],
            env: HashMap::new(),
            cluster: None,
            entrypoint: vec![],
            distributed_backend: "mp".to_string(),
            before_serve: vec![],
        }
    }

    #[test]
    fn test_generate_repair_dockerfile_vllm() {
        let recipe = test_recipe();
        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);

        assert!(dockerfile.contains("FROM base/image:tag"));
        assert!(dockerfile.contains("pip install vllm"));
        assert!(!dockerfile.contains("ray"));
        assert!(dockerfile.contains("ENTRYPOINT"));
    }

    #[test]
    fn test_generate_repair_dockerfile_ray() {
        let mut recipe = test_recipe();
        recipe.distributed_backend = "ray".to_string();
        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);

        assert!(dockerfile.contains("FROM base/image:tag"));
        assert!(dockerfile.contains("pip install vllm"));
        assert!(dockerfile.contains("pip install ray"));
    }

    #[test]
    fn test_generate_repair_dockerfile_sglang() {
        let mut recipe = test_recipe();
        recipe.engine = "sglang".to_string();
        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);

        assert!(dockerfile.contains("pip install sglang"));
        assert!(!dockerfile.contains("vllm"));
    }
}