rusty_stack/

component_status.rs

use crate::benchmark_logs;
use crate::state::{Category, Component};
use std::collections::HashSet;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

#[derive(Debug, Clone)]
pub struct VerificationCommand {
    pub label: String,
    pub target_id: String,
    pub program: String,
    pub args: Vec<String>,
    pub modules: Vec<String>,
}

pub fn python_interpreters() -> Vec<String> {
    let mut candidates = Vec::new();
    if let Ok(value) = env::var("MLSTACK_PYTHON_BIN") {
        push_python_candidate(&mut candidates, value);
    }
    if let Ok(value) = env::var("AITER_VENV_PYTHON") {
        push_python_candidate(&mut candidates, value);
    }
    if let Ok(value) = env::var("UV_PYTHON") {
        push_python_candidate(&mut candidates, value);
    }
    for key in [
        "WANDB_VENV_PYTHON",
        "MEGATRON_VENV_PYTHON",
        "MPI4PY_VENV_PYTHON",
        "FLASH_ATTENTION_VENV_PYTHON",
        "PYTORCH_VENV_PYTHON",
        "DEEPSPEED_VENV_PYTHON",
        "VLLM_VENV_PYTHON",
    ] {
        if let Ok(value) = env::var(key) {
            push_python_candidate(&mut candidates, value);
        }
    }
    let home = resolve_component_user_home();
    let venv = Path::new(&home)
        .join("rocm_venv")
        .join("bin")
        .join("python");
    if venv.exists() {
        candidates.push(venv.to_string_lossy().to_string());
    }

    // Anaconda/Miniconda candidates
    let conda_path = Path::new(&home)
        .join("anaconda3")
        .join("bin")
        .join("python");
    if conda_path.exists() {
        candidates.push(conda_path.to_string_lossy().to_string());
    }
    let miniconda_path = Path::new(&home)
        .join("miniconda3")
        .join("bin")
        .join("python");
    if miniconda_path.exists() {
        candidates.push(miniconda_path.to_string_lossy().to_string());
    }

    if Path::new("/usr/local/bin/python3").exists() {
        candidates.push("/usr/local/bin/python3".to_string());
    }
    if command_exists("python3.13") {
        candidates.push("python3.13".to_string());
    }
    if command_exists("python3.12") {
        candidates.push("python3.12".to_string());
    }
    if command_exists("python3") {
        candidates.push("python3".to_string());
    }
    if command_exists("python") {
        candidates.push("python".to_string());
    }

    let mut roots = vec![PathBuf::from(&home)];
    if let Ok(cwd) = env::current_dir() {
        roots.push(cwd.clone());
        roots.push(cwd.join("rusty-stack"));
    }
    for root in roots {
        push_component_venv_candidates(&mut candidates, &root);
    }

    dedupe_python_candidates_in_order(candidates)
}

fn dedupe_python_candidates_in_order(candidates: Vec<String>) -> Vec<String> {
    let mut seen = HashSet::new();
    let mut ordered = Vec::new();
    for candidate in candidates {
        if seen.insert(candidate.clone()) {
            ordered.push(candidate);
        }
    }
    ordered
}

fn push_component_venv_candidates(candidates: &mut Vec<String>, root: &Path) {
    for rel in [
        "wandb_venv/bin/python",
        "megatron_rocm_venv/bin/python",
        "Megatron-LM/megatron_rocm_venv/bin/python",
        "mpi4py_venv/bin/python",
        "flash_attention_venv/bin/python",
        "vllm_venv/bin/python",
        ".mlstack/venvs/aiter/bin/python",
        ".mlstack/venvs/vllm/bin/python",
        "pytorch_rocm_venv/bin/python",
    ] {
        let candidate = root.join(rel);
        if candidate.exists() {
            candidates.push(candidate.to_string_lossy().to_string());
        }
    }
}

fn push_python_candidate(candidates: &mut Vec<String>, value: String) {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return;
    }
    let candidate = trimmed.to_string();
    if Path::new(&candidate).exists() || command_exists(&candidate) {
        candidates.push(candidate);
    }
}

pub fn is_component_installed(component: &Component, python_candidates: &[String]) -> bool {
    if component.category == Category::Maintenance {
        return false;
    }
    is_component_installed_by_id(&component.id, python_candidates)
}

pub fn is_component_installed_by_id(component_id: &str, python_candidates: &[String]) -> bool {
    let home = resolve_component_user_home();
    let has_benchmark_logs = has_benchmark_log_dirs();
    match component_id {
        // ROCm requires BOTH version file AND functional rocminfo (not just binary existence)
        // This prevents false positives from partial downloads
        "rocm" => {
            // Check for version file in multiple locations
            let version_exists = path_exists("/opt/rocm/.info/version")
                || path_exists("/opt/rocm/version")
                || path_exists("/usr/lib/rocm/.info/version");

            // Check if rocminfo actually works (not just exists)
            let rocminfo_path = crate::installers::common::utils::resolve_rocminfo_path();
            let rocminfo_functional = if let Ok(output) = std::process::Command::new(&rocminfo_path)
                .arg("--version")
                .output()
            {
                output.status.success()
            } else {
                false
            };

            // Require both: version file AND functional rocminfo
            // This ensures we don't report partial installs as complete
            version_exists && rocminfo_functional
        }
        "pytorch" => python_any(python_candidates, &["torch"]),
        "triton" => python_any(python_candidates, &["triton"]),
        "mpi4py" => python_any(python_candidates, &["mpi4py"]),
        "deepspeed" => python_any(python_candidates, &["deepspeed"]),
        "ml-stack-core" => {
            python_any(python_candidates, &["stans_ml_stack"]) || repo_has_ml_stack_core()
        }
        // Canonical "flash-attn" id (used by verify/migrate): installed if the
        // flash_attn package is importable under ANY backend. The concrete TUI
        // components (flash-attn-triton / flash-attn-ck) below use the stricter
        // backend-marker check.
        "flash-attn" => python_any(python_candidates, &["flash_attn", "flash_attn_2"]),
        // Flash Attention backends install the SAME `flash_attn` package, so
        // importability alone can't tell them apart. "installed" requires the
        // module importable AND the backend marker
        // (~/.mlstack/flash-attention/.backend) to match — this is what stops a
        // leftover build dir or a manual/unknown install from false-reporting as
        // installed. Only the last-installed backend (the marker value) reports
        // installed (mutual exclusivity).
        "flash-attn-triton" => {
            python_any(python_candidates, &["flash_attn", "flash_attn_2"])
                && backend_marker_is("triton")
        }
        "flash-attn-ck" => {
            python_any(python_candidates, &["flash_attn", "flash_attn_2"])
                && backend_marker_is("ck")
        }
        "megatron" => {
            path_exists(home_path(&home, &["megatron", "Megatron-LM"]))
                || python_any(python_candidates, &["megatron"])
        }
        "vllm" => {
            python_candidates
                .iter()
                .any(|python| python_exec(python, vllm_runtime_check_snippet()))
                || path_exists(home_path(&home, &["vllm_build"]))
                || path_exists(home_path(&home, &["vllm_py313"]))
        }
        "aiter" => python_any(python_candidates, &["aiter"]),
        "vllm-studio" => {
            command_exists("vllm-studio") || path_exists(home_path(&home, &["vllm-studio"]))
        }
        "onnx" => {
            python_candidates
                .iter()
                .any(|python| python_exec(python, onnxruntime_present_snippet()))
                || path_exists(home_path(&home, &["onnxruntime_build"]))
        }
        "bitsandbytes" => {
            python_any(python_candidates, &["bitsandbytes"])
                || path_exists(home_path(&home, &["ml_stack", "bitsandbytes"]))
        }
        // rocm-smi - check multiple locations and verify it works
        "rocm-smi" => {
            let rocm_smi_paths = [
                "/opt/rocm/bin/rocm-smi",
                "/usr/bin/rocm-smi",
                "/usr/local/bin/rocm-smi",
            ];
            for path in &rocm_smi_paths {
                if path_exists(path) {
                    // Verify it actually works
                    if let Ok(output) = std::process::Command::new(path).arg("--showuse").output() {
                        if output.status.success() {
                            return true;
                        }
                    }
                }
            }
            // Also check if rocm-smi is in PATH and works
            if command_exists("rocm-smi") {
                if let Ok(output) = std::process::Command::new("rocm-smi")
                    .arg("--showuse")
                    .output()
                {
                    return output.status.success();
                }
            }
            false
        }
        "migraphx" => {
            // Python bindings (Debian/ROCm image) OR the C++ driver (Arch/CachyOS,
            // where pip `migraphx` is unavailable). Functional only — a driver that
            // can't load its shared libs (libprotobuf/libabsl SONAME desync from a
            // partial upgrade) is NOT installed. The old `migraphx_build` dir check
            // was a false-positive surface (leftover build dir ≠ functional), removed.
            python_any(python_candidates, &["migraphx"])
                || crate::installers::components::migraphx_multi::migraphx_driver_functional()
        }
        "pytorch-profiler" => python_any(python_candidates, &["torch"]),
        "wandb" => python_any(python_candidates, &["wandb"]),
        "fastvideo" => python_candidates.iter().any(|python| {
            python_exec(
                python,
                crate::installers::components::fastvideo::fastvideo_verification_snippet(),
            )
        }),
        "llama-cpp" => {
            // Use the detection contract from llama_cpp.rs: llama-cli --help must succeed
            crate::installers::components::llama_cpp::is_llama_cli_functional(&home)
        }
        "permanent-env" => env_file_has_permanent(&home_path(&home, &[".mlstack_env"])),
        "basic-env" => path_exists(home_path(&home, &[".mlstack_env"])),
        "enhanced-env" => env_file_has_enhanced(&home_path(&home, &[".mlstack_env"])),
        "vllm-performance" => has_benchmark_logs,
        "deepspeed-performance" => has_benchmark_logs,
        "megatron-performance" => has_benchmark_logs,
        "all-benchmarks" => has_benchmark_logs,
        "comfyui" => {
            // Check for ComfyUI installation
            // Default location: $HOME/ComfyUI
            // User location: /mnt/e0f7c1a8-b834-4827-b579-0251b006bc1f/ComfyUI/
            // Check for .git folder and main.py
            let default_path = home_path(&home, &["ComfyUI"]);
            path_exists(default_path.join(".git")) || path_exists(default_path.join("main.py"))
        }
        "textgen" => {
            // Check for text-generation-webui installation
            // Default location: $HOME/text-generation-webui
            let default_path = home_path(&home, &["text-generation-webui"]);
            path_exists(default_path.join(".git")) || path_exists(default_path.join("server.py"))
        }
        _ => false,
    }
}

pub fn verification_commands(
    component_id: &str,
    python_candidates: &[String],
) -> Vec<VerificationCommand> {
    match component_id {
        "verify-basic" => basic_verification_commands(python_candidates),
        "verify-enhanced" => enhanced_verification_commands(python_candidates),
        "verify-build" => build_verification_commands(python_candidates),
        _ => Vec::new(),
    }
}

pub fn component_verification_commands(
    component_id: &str,
    python_candidates: &[String],
) -> Vec<VerificationCommand> {
    match component_id {
        "rocm" => vec![shell_command("ROCm info", "rocm", "rocminfo", &[])],
        "pytorch" => vec![python_command(
            "PyTorch",
            "pytorch",
            &["torch"],
            python_candidates,
            pytorch_diagnostic_snippet(),
        )],
        "triton" => vec![python_command(
            "Triton",
            "triton",
            &["triton"],
            python_candidates,
            "import triton,sys; import triton.language as tl; import importlib; ok=False\nfor _m in ('triton.runtime.driver','triton.runtime.backends.amd','triton.backends.amd'):\n    try:\n        importlib.import_module(_m); ok=True; break\n    except Exception:\n        pass\nprint(f'Triton {triton.__version__} backend_loaded={ok}'); sys.exit(0 if ok else 1)",
        )],
        "mpi4py" => vec![python_command(
            "MPI4Py",
            "mpi4py",
            &["mpi4py"],
            python_candidates,
            "from mpi4py import MPI; import sys\ntry:\n    rank=MPI.COMM_WORLD.Get_rank(); ok=rank>=0\nexcept Exception:\n    rank=-1; ok=False\nprint(f'MPI4Py rank={rank} ok={ok}'); sys.exit(0 if ok else 1)",
        )],
        "deepspeed" => vec![python_command(
            "DeepSpeed",
            "deepspeed",
            &["deepspeed"],
            python_candidates,
            "import deepspeed,sys; ok=False\ntry:\n    from deepspeed.ops.adam import DeepSpeedCPUAdam; ok=True\nexcept Exception:\n    pass\nprint(f'DeepSpeed {deepspeed.__version__} ops_loaded={ok}'); sys.exit(0 if ok else 1)",
        )],
        "ml-stack-core" => vec![python_command(
            "ML Stack Core",
            "ml-stack-core",
            &["stans_ml_stack"],
            python_candidates,
            "import stans_ml_stack,sys; import importlib; ok=True\nfor _s in ('core','utils','cli','installers'):\n    try:\n        importlib.import_module(f'stans_ml_stack.{_s}')\n    except Exception:\n        ok=False; break\nver=getattr(stans_ml_stack,'__version__','ok')\nprint(f'ML Stack Core {ver} submodule_loaded={ok}'); sys.exit(0 if ok else 1)",
        )],
        // Canonical id (verify/migrate): flash_attn importable under ANY backend.
        "flash-attn" => vec![flash_attention_verification_command(
            "Flash Attention",
            "flash-attn",
            None,
            python_candidates,
        )],
        // Split TUI ids: require the backend marker to match, so a successful
        // install verifies + a stale/manual one does not.
        "flash-attn-triton" => vec![flash_attention_verification_command(
            "Flash Attention (Triton)",
            "flash-attn-triton",
            Some("triton"),
            python_candidates,
        )],
        "flash-attn-ck" => vec![flash_attention_verification_command(
            "Flash Attention (CK)",
            "flash-attn-ck",
            Some("ck"),
            python_candidates,
        )],
        "megatron" => vec![python_command(
            "Megatron-LM",
            "megatron",
            &["megatron"],
            python_candidates,
            "import megatron,sys; ok=False\ntry:\n    from megatron.core import tensor_parallel; ok=True\nexcept Exception:\n    pass\nprint(f'Megatron-LM functional={ok}'); sys.exit(0 if ok else 1)",
        )],
        "vllm" => vec![python_command(
            "vLLM",
            "vllm",
            &["vllm"],
            python_candidates,
            vllm_runtime_check_snippet(),
        )],
        "aiter" => vec![python_command(
            "AITER",
            "aiter",
            &["aiter"],
            python_candidates,
            "import aiter,sys; ok=False\ntry:\n    from aiter import op as _op; ok=True\nexcept Exception:\n    pass\nmods=[x for x in dir(aiter) if not x.startswith('_')]\nver=getattr(aiter,'__version__','?')\nprint(f'AITER {ver} op_loaded={ok} exports={len(mods)}'); sys.exit(0 if ok else 1)",
        )],
        "vllm-studio" => vec![shell_command(
            "vLLM Studio",
            "vllm-studio",
            "bun",
            &["--version"],
        )],
        "onnx" => vec![python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "import onnxruntime as ort; import sys; print('Version:', ort.__version__); providers=ort.get_available_providers(); print('Providers:', providers); gpu=[p for p in providers if p in ('ROCMExecutionProvider','MIGraphXExecutionProvider','CUDAExecutionProvider')]; sys.exit(0 if gpu else 1)",
        )],
        "bitsandbytes" => vec![python_command(
            "BITSANDBYTES",
            "bitsandbytes",
            &["bitsandbytes"],
            python_candidates,
            "import bitsandbytes as bnb; import pathlib; import sys; print('Version:', getattr(bnb, '__version__', 'unknown')); path=pathlib.Path(bnb.__file__).parent; libs=list(path.glob('libbitsandbytes_rocm*.so')); print('ROCm Libs:', libs); sys.exit(0 if libs else 1)",
        )],
        "rocm-smi" => vec![shell_command(
            "ROCm SMI",
            "rocm-smi",
            "rocm-smi",
            &["--showproductname"],
        )],
        "migraphx" => vec![python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            migraphx_verify_snippet(),
        )],
        "pytorch-profiler" => vec![python_command(
            "PyTorch Profiler",
            "pytorch-profiler",
            &["torch"],
            python_candidates,
            "from torch.profiler import profile; print(profile)",
        )],
        "wandb" => vec![python_command(
            "Weights & Biases",
            "wandb",
            &["wandb"],
            python_candidates,
            "import wandb,sys; ok=False\ntry:\n    _r=wandb.init(mode='disabled'); ok=_r is not None; _r.finish()\nexcept Exception:\n    pass\nprint(f'wandb {wandb.__version__} init_run_ok={ok}'); sys.exit(0 if ok else 1)",
        )],
        "fastvideo" => vec![python_command(
            "FastVideo",
            "fastvideo",
            &["fastvideo"],
            python_candidates,
            crate::installers::components::fastvideo::fastvideo_verification_snippet(),
        )],
        "llama-cpp" => vec![shell_command(
            "llama.cpp",
            "llama-cpp",
            "llama-cli",
            &["--help"],
        )],
        "basic-env" => vec![shell_command(
            "Environment file",
            "basic-env",
            "bash",
            &[
                "-c",
                "test -f \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        "enhanced-env" => vec![shell_command(
            "Enhanced env",
            "enhanced-env",
            "bash",
            &[
                "-c",
                "grep -q \"Enhanced ML Stack Environment Setup Script\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        "permanent-env" => vec![shell_command(
            "Permanent env",
            "permanent-env",
            "bash",
            &[
                "-c",
                "grep -qiE \"Permanent ROCm Environment|Permanent ROCm Env|MLSTACK_PYTHON_BIN\" \"$HOME/.mlstack_env\" && echo OK || exit 1",
            ],
        )],
        // Benchmark components - verify by checking log files exist
        "mlperf-inference" => vec![benchmark_log_check_command(
            "MLPerf benchmark logs",
            "mlperf-inference",
            "mlperf_inference",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "rocm-benchmarks" => vec![benchmark_log_check_command(
            "ROCm benchmark logs",
            "rocm-benchmarks",
            "rocm_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "gpu-memory-bandwidth" => vec![benchmark_log_check_command(
            "Memory bandwidth logs",
            "gpu-memory-bandwidth",
            "gpu_memory_bandwidth",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "rocm-smi-bench" => vec![benchmark_log_check_command(
            "ROCm SMI logs",
            "rocm-smi-bench",
            "rocm_smi_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "pytorch-performance" => vec![benchmark_log_check_command(
            "PyTorch benchmark logs",
            "pytorch-performance",
            "pytorch_performance",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "vllm-performance" => vec![benchmark_log_check_command(
            "vLLM benchmark logs",
            "vllm-performance",
            "vllm_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "deepspeed-performance" => vec![benchmark_log_check_command(
            "DeepSpeed benchmark logs",
            "deepspeed-performance",
            "deepspeed_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "megatron-performance" => vec![benchmark_log_check_command(
            "Megatron benchmark logs",
            "megatron-performance",
            "megatron_benchmarks",
            "Benchmark logs found",
            "No benchmark logs yet",
        )],
        "all-benchmarks" => vec![benchmark_log_check_command(
            "Full suite logs",
            "all-benchmarks",
            "full_benchmarks",
            "Suite logs found",
            "No suite logs yet",
        )],
        "comfyui" => vec![{
            // F2: verify against the env's interpreter, not a hardcoded python3.
            // Resolve the same python the installers pin (MLSTACK_PYTHON_BIN /
            // UV_PYTHON / first candidate), then import from the install dir.
            let py = python_candidates
                .first()
                .cloned()
                .or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
                .or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
                .unwrap_or_else(|| "python3".to_string());
            shell_command(
                "ComfyUI",
                "comfyui",
                "bash",
                &[
                    "-c",
                    &format!(
                        "cd \"$HOME/ComfyUI\" 2>/dev/null && {py} -c 'import folder_paths' 2>/dev/null && echo 'comfyui functional' || exit 1"
                    ),
                ],
            )
        }],
        "textgen" => vec![{
            // F2: verify against the env's interpreter, not a hardcoded python3.
            let py = python_candidates
                .first()
                .cloned()
                .or_else(|| env::var("MLSTACK_PYTHON_BIN").ok().filter(|s| !s.is_empty()))
                .or_else(|| env::var("UV_PYTHON").ok().filter(|s| !s.is_empty()))
                .unwrap_or_else(|| "python3".to_string());
            shell_command(
                "text-generation-webui",
                "textgen",
                "bash",
                &[
                    "-c",
                    &format!(
                        "cd \"$HOME/text-generation-webui\" 2>/dev/null && {py} -c 'import server' 2>/dev/null && echo 'textgen functional' || exit 1"
                    ),
                ],
            )
        }],
        _ => Vec::new(),
    }
}

pub fn modules_available(modules: &[String], python_candidates: &[String]) -> bool {
    if modules.is_empty() {
        return true;
    }
    let module_refs: Vec<&str> = modules.iter().map(|s| s.as_str()).collect();
    python_any(python_candidates, &module_refs)
}

fn basic_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    vec![
        shell_command("ROCm info", "rocm", "rocminfo", &[]),
        python_command(
            "PyTorch",
            "pytorch",
            &["torch"],
            python_candidates,
            pytorch_diagnostic_snippet(),
        ),
        python_command(
            "Triton",
            "triton",
            &["triton"],
            python_candidates,
            "import triton; import triton._C; print(f'Triton {triton.__version__} ok')",
        ),
        python_command(
            "MPI4Py",
            "mpi4py",
            &["mpi4py"],
            python_candidates,
            "from mpi4py import MPI; import sys\ntry:\n    rank=MPI.COMM_WORLD.Get_rank(); ok=rank>=0\nexcept Exception:\n    rank=-1; ok=False\nprint(f'MPI4Py rank={rank} ok={ok}'); sys.exit(0 if ok else 1)",
        ),
        python_command(
            "DeepSpeed",
            "deepspeed",
            &["deepspeed", "einops"],
            python_candidates,
            "import deepspeed; import einops; print(deepspeed.__version__)",
        ),
    ]
}

fn enhanced_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    let mut steps = basic_verification_commands(python_candidates);
    steps.extend(vec![
        flash_attention_verification_command(
            "Flash Attention",
            "flash-attn",
            None,
            python_candidates,
        ),
        python_command(
            "vLLM",
            "vllm",
            &["vllm"],
            python_candidates,
            vllm_runtime_check_snippet(),
        ),
        python_command(
            "AITER",
            "aiter",
            &["aiter"],
            python_candidates,
            "import aiter; print(getattr(aiter, '__version__', 'ok'))",
        ),
        python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "import onnxruntime as ort; import pathlib; import os; import sys; print('Version:', ort.__version__); base=pathlib.Path(ort.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); [os.environ.update({'ORT_ROCM_EP_PROVIDER_PATH': str(l)}) for l in libs[:1]]; providers=ort.get_available_providers(); print('Providers:', providers); sys.exit(0 if 'ROCMExecutionProvider' in providers else 1)",
        ),
        python_command(
            "BITSANDBYTES",
            "bitsandbytes",
            &["bitsandbytes"],
            python_candidates,
            "import bitsandbytes as bnb; import pathlib; import sys; print('Version:', getattr(bnb, '__version__', 'unknown')); path=pathlib.Path(bnb.__file__).parent; libs=list(path.glob('libbitsandbytes_rocm*.so')); print('ROCm Libs:', libs); sys.exit(0 if libs else 1)",
        ),
        shell_command("ROCm SMI", "rocm-smi", "rocm-smi", &["--showproductname"]),
        python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            migraphx_verify_snippet(),
        ),
        python_command(
            "PyTorch Profiler",
            "pytorch-profiler",
            &["torch"],
            python_candidates,
            "from torch.profiler import profile; print(profile)",
        ),
        python_command(
            "Weights & Biases",
            "wandb",
            &["wandb"],
            python_candidates,
            "import wandb,sys; ok=False\ntry:\n    _r=wandb.init(mode='disabled'); ok=_r is not None; _r.finish()\nexcept Exception:\n    pass\nprint(f'wandb {wandb.__version__} init_run_ok={ok}'); sys.exit(0 if ok else 1)",
        ),
        python_command(
            "FastVideo",
            "fastvideo",
            &["fastvideo"],
            python_candidates,
            crate::installers::components::fastvideo::fastvideo_verification_snippet(),
        ),
        shell_command("vLLM Studio", "vllm-studio", "bun", &["--version"]),
    ]);
    steps
}

fn build_verification_commands(python_candidates: &[String]) -> Vec<VerificationCommand> {
    let mut steps = basic_verification_commands(python_candidates);
    steps.extend(vec![
        python_command(
            "ONNX Runtime",
            "onnx",
            &["onnxruntime"],
            python_candidates,
            "import onnxruntime as ort; import pathlib; import os; import sys; print('Version:', ort.__version__); base=pathlib.Path(ort.__file__).parent; libs=list(base.rglob('libonnxruntime_providers_rocm.so')); [os.environ.update({'ORT_ROCM_EP_PROVIDER_PATH': str(l)}) for l in libs[:1]]; providers=ort.get_available_providers(); print('Providers:', providers); sys.exit(0 if 'ROCMExecutionProvider' in providers else 1)",
        ),
        flash_attention_verification_command(
            "Flash Attention (build)",
            "flash-attn",
            None,
            python_candidates,
        ),
        python_command(
            "MIGraphX",
            "migraphx",
            &["migraphx"],
            python_candidates,
            migraphx_verify_snippet(),
        ),
    ]);
    steps
}

fn python_command(
    label: &str,
    target_id: &str,
    modules: &[&str],
    python_candidates: &[String],
    code: &str,
) -> VerificationCommand {
    let program = select_python_for_modules(modules, python_candidates)
        .or_else(|| python_candidates.first().cloned())
        .unwrap_or_else(|| "python3".to_string());

    VerificationCommand {
        label: label.to_string(),
        target_id: target_id.to_string(),
        program,
        args: vec!["-c".to_string(), code.to_string()],
        modules: modules.iter().map(|m| m.to_string()).collect(),
    }
}

/// PyTorch diagnostic verification snippet.
///
/// This performs a multi-tier check:
/// 1. Can we import torch? (basic install check)
/// 2. Is torch.version.hip set? (ROCm build check)
/// 3. Is torch.cuda.is_available()? (runtime HIP check)
/// 4. Diagnostics for why HIP might not be available
///
/// Verification reflects FUNCTIONALITY, not just import: it exits non-zero if
/// torch is not a ROCm build (`torch.version.hip` is None) OR the HIP runtime
/// is unavailable (`torch.cuda.is_available()` is False). The diagnostic tiers
/// print actionable guidance. Whether a pytorch-verify failure should block
/// sibling components is an orchestrator-layer concern — NOT a reason to lie
/// about pytorch's status. (Stage 4 tenet: verify functionality, fail loud.)
fn pytorch_diagnostic_snippet() -> &'static str {
    r#"import torch, sys, os, ctypes.util

# Tier 1: Basic torch import (already verified by module check)
print(f'torch version: {torch.__version__}')

# Tier 2: Check if this is a ROCm build
hip_version = getattr(torch.version, 'hip', None)
cuda_version = getattr(torch.version, 'cuda', None)
print(f'hip: {hip_version}')
print(f'cuda_build: {cuda_version}')

# Tier 3: Runtime HIP/CUDA availability
try:
    hip_available = torch.cuda.is_available()
    # 'cuda' here is AVAILABILITY (CUDA-via-HIP on ROCm), NOT torch.version.cuda
    # (which is None on a ROCm build — that's the CUDA *toolkit* version torch was
    # built against, not whether CUDA-style code can run). is_available()==True
    # means CUDA works through HIP, which is the user-facing answer.
    print(f'cuda: {hip_available}')
    if hip_available:
        device_count = torch.cuda.device_count()
        print(f'device_count: {device_count}')
        if device_count > 0:
            print(f'device_name: {torch.cuda.get_device_name(0)}')
except Exception as e:
    hip_available = False
    print(f'torch.cuda.is_available() error: {e}')

# Tier 4: Diagnostics if HIP not available
if not hip_available:
    print('--- ROCm HIP Diagnostics ---')

    # Check LD_LIBRARY_PATH
    ld_path = os.environ.get('LD_LIBRARY_PATH', '')
    print(f'LD_LIBRARY_PATH: {ld_path if ld_path else "(not set)"}')

    # Check ROCm_PATH
    rocm_path = os.environ.get('ROCM_PATH', '')
    print(f'ROCM_PATH: {rocm_path if rocm_path else "(not set)"}')

    # Check if ROCm libraries exist
    rocm_lib_dirs = ['/opt/rocm/lib', '/opt/rocm/hip/lib']
    for d in rocm_lib_dirs:
        exists = os.path.isdir(d)
        print(f'{d}: {"exists" if exists else "not found"}')

    # Check if libamdhip64.so can be found
    hip_lib = ctypes.util.find_library('amdhip64')
    print(f'libamdhip64: {hip_lib if hip_lib else "not found"}')

    # Check Python executable path
    print(f'python: {sys.executable}')

    # Provide actionable guidance
    if hip_version is None and cuda_version is None:
        print('HINT: torch.version.hip and torch.version.cuda are both None.')
        print('  This torch build may not have GPU support.')
        print('  Reinstall with ROCm index: pip install torch --index-url https://download.pytorch.org/whl/rocm7.2/')
    elif hip_version is not None and not hip_available:
        print('HINT: torch.version.hip is set but HIP runtime not available.')
        print('  Possible causes:')
        print('  1. LD_LIBRARY_PATH does not include /opt/rocm/lib')
        print('  2. Running in a container/environment without GPU access')
        print('  3. ROCm drivers not loaded (check with: rocminfo)')
        print('  Fix: source ~/.mlstack_env or set LD_LIBRARY_PATH=/opt/rocm/lib')

# Stage 4: exit code reflects FUNCTIONAL HIP availability on this ROCm stack.
# A CPU-only/CUDA torch build, or a ROCm build whose HIP runtime is
# unreachable, must FAIL verification — not silently pass. The diagnostics
# above guide the fix; downstream gating belongs in the orchestrator layer,
# not a lying exit code.
ok = (hip_version is not None) and hip_available
sys.exit(0 if ok else 1)"#
}

fn vllm_runtime_check_snippet() -> &'static str {
    r#"import importlib
import vllm
import cachetools
import cbor2
import gguf
import pybase64
import ijson
import mistral_common
import openai_harmony

loaded = []
errs = []
for name in ("vllm._C", "vllm._rocm_C"):
    try:
        importlib.import_module(name)
        loaded.append(name)
    except Exception as exc:
        errs.append(f"{name}: {exc}")

if not loaded:
    raise SystemExit("vLLM native extension load failure: " + " | ".join(errs))

print(vllm.__version__)"#
}

fn shell_command(
    label: &str,
    target_id: &str,
    program: &str,
    args: &[&str],
) -> VerificationCommand {
    VerificationCommand {
        label: label.to_string(),
        target_id: target_id.to_string(),
        program: program.to_string(),
        args: args.iter().map(|s| s.to_string()).collect(),
        modules: Vec::new(),
    }
}

fn has_benchmark_log_dirs() -> bool {
    benchmark_logs::benchmark_log_directories()
        .iter()
        .any(|dir| {
            fs::read_dir(dir)
                .ok()
                .map(|entries| {
                    entries
                        .filter_map(|entry| entry.ok())
                        .any(|entry| entry.path().is_file())
                })
                .unwrap_or(false)
        })
}

fn benchmark_log_check_command(
    label: &str,
    target_id: &str,
    pattern: &str,
    success_message: &str,
    failure_message: &str,
) -> VerificationCommand {
    let script = benchmark_log_check_script(pattern, success_message, failure_message);
    shell_command(label, target_id, "bash", &["-c", &script])
}

fn benchmark_log_check_script(
    pattern: &str,
    success_message: &str,
    failure_message: &str,
) -> String {
    format!(
        "dirs=(\"${{MLSTACK_LOG_DIR:-}}\" \"$HOME/.rusty-stack/logs\" \"${{TMPDIR:-/tmp}}/rusty-stack/logs\")\n\
         found=0\n\
         for dir in \"${{dirs[@]}}\"; do\n\
             if [ -n \"$dir\" ] && [ -d \"$dir\" ]; then\n\
                 if ls \"$dir\" | grep -q \"{}\"; then\n\
                     echo \"{}\"\n\
                     found=1\n\
                     break\n\
                 fi\n\
             fi\n\
         done\n\
         if [ \"$found\" -ne 1 ]; then\n\
             echo \"{}\"\n\
         fi",
        pattern, success_message, failure_message
    )
}

fn python_any(python_candidates: &[String], modules: &[&str]) -> bool {
    python_candidates.iter().any(|python| {
        modules
            .iter()
            .any(|module| python_has_module(python, module))
    })
}

fn python_has_module(python: &str, module: &str) -> bool {
    python_exec(
        python,
        &format!(
            "import importlib.util, sys; sys.exit(0 if importlib.util.find_spec('{module}') else 1)"
        ),
    )
}

/// Detection snippet for `onnxruntime`: requires a REAL installed package, not
/// a CWD namespace-package shadow.
///
/// `python_has_module`/`find_spec` alone is fooled by a repo-local
/// `onnxruntime/` source directory (e.g. a clone checked out inside the repo):
/// it resolves as a namespace package (`find_spec` returns a spec, but
/// `origin` is `None`), so a bare directory reports "installed" while
/// `import onnxruntime` yields a hollow module with no `__version__`/providers.
/// This strips CWD from `sys.path` (so the env package wins) AND requires a
/// non-`None` `origin` (a real `__init__.py`-backed package). Same footgun
/// fixed for migraphx in `migraphx_verify_snippet`.
fn onnxruntime_present_snippet() -> &'static str {
    "import sys, os, importlib.util; \
     sys.path = [p for p in sys.path if p and os.path.abspath(p.rstrip('/')) != os.path.abspath(os.getcwd())]; \
     s = importlib.util.find_spec('onnxruntime'); \
     sys.exit(0 if (s is not None and s.origin is not None) else 1)"
}

fn python_exec(python: &str, code: &str) -> bool {
    let mut cmd = Command::new(python);
    cmd.arg("-c")
        .arg(code)
        .stdout(Stdio::null())
        .stderr(Stdio::null());

    if let Ok(existing) = env::var("PYTHONPATH") {
        let mut path = existing;
        let extra_paths = python_search_paths();
        if !extra_paths.is_empty() {
            if !path.is_empty() {
                path.push(':');
            }
            path.push_str(&extra_paths.join(":"));
        }
        cmd.env("PYTHONPATH", path);
    } else {
        let extra_paths = python_search_paths();
        if !extra_paths.is_empty() {
            cmd.env("PYTHONPATH", extra_paths.join(":"));
        }
    }

    let home = resolve_component_user_home();
    let ld_path = component_status_ld_library_path(&home);
    if !ld_path.is_empty() {
        cmd.env("LD_LIBRARY_PATH", ld_path);
    }

    cmd.status().map(|status| status.success()).unwrap_or(false)
}

fn push_unique_path(paths: &mut Vec<String>, value: &str) {
    let value = value.trim();
    if value.is_empty() {
        return;
    }
    if paths.iter().any(|existing| existing == value) {
        return;
    }
    paths.push(value.to_string());
}

fn component_user_home_from_passwd(user_name: &str) -> Option<String> {
    let target = user_name.trim();
    if target.is_empty() {
        return None;
    }
    let passwd = fs::read_to_string("/etc/passwd").ok()?;
    passwd.lines().find_map(|line| {
        if line.trim().is_empty() || line.starts_with('#') {
            return None;
        }
        let fields: Vec<&str> = line.split(':').collect();
        if fields.len() < 6 || fields[0] != target {
            return None;
        }
        let home = fields[5].trim();
        if home.is_empty() {
            None
        } else {
            Some(home.to_string())
        }
    })
}

fn component_passwd_homes_with_mlstack() -> Vec<String> {
    let mut homes = Vec::new();
    let Ok(passwd) = fs::read_to_string("/etc/passwd") else {
        return homes;
    };

    for line in passwd.lines() {
        if line.trim().is_empty() || line.starts_with('#') {
            continue;
        }
        let fields: Vec<&str> = line.split(':').collect();
        if fields.len() < 6 {
            continue;
        }
        let home = fields[5].trim();
        if home.is_empty() {
            continue;
        }
        if Path::new(home).join(".mlstack").is_dir()
            || Path::new(home).join(".mlstack_env").is_file()
        {
            push_unique_path(&mut homes, home);
        }
    }

    homes
}

fn component_candidate_homes(preferred_home: &str) -> Vec<String> {
    let mut homes = Vec::new();
    push_unique_path(&mut homes, preferred_home);

    if let Ok(value) = env::var("MLSTACK_USER_HOME") {
        push_unique_path(&mut homes, &value);
    }
    if let Ok(value) = env::var("HOME") {
        push_unique_path(&mut homes, &value);
    }

    for key in ["SUDO_USER", "USER", "LOGNAME"] {
        if let Ok(user_name) = env::var(key) {
            if let Some(home) = component_user_home_from_passwd(&user_name) {
                push_unique_path(&mut homes, &home);
            }
        }
    }

    for home in component_passwd_homes_with_mlstack() {
        push_unique_path(&mut homes, &home);
    }

    homes
}

fn resolve_component_user_home() -> String {
    if let Ok(home) = env::var("MLSTACK_USER_HOME") {
        return home;
    }

    let fallback = env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
    let candidates = component_candidate_homes(&fallback);

    for home in &candidates {
        if Path::new(home).join(".mlstack_env").is_file() {
            return home.clone();
        }
    }
    for home in &candidates {
        if Path::new(home).join(".mlstack").is_dir() {
            return home.clone();
        }
    }

    candidates.into_iter().next().unwrap_or(fallback)
}

fn component_status_mpi_compat_dirs(home: &str) -> Vec<String> {
    let mut dirs = Vec::new();
    for home in component_candidate_homes(home) {
        let mlstack_dir = Path::new(&home).join(".mlstack");
        let primary = mlstack_dir.join("libmpi-compat");
        if primary.is_dir() {
            dirs.push(primary.to_string_lossy().to_string());
        }

        if let Ok(entries) = fs::read_dir(&mlstack_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if !path.is_dir() {
                    continue;
                }
                let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
                    continue;
                };
                if name.starts_with("libmpi-compat-user-") {
                    dirs.push(path.to_string_lossy().to_string());
                }
            }
        }
    }

    dirs.sort();
    dirs.dedup();
    dirs
}

fn component_status_ld_library_path(home: &str) -> String {
    let mut paths = Vec::new();

    for dir in component_status_mpi_compat_dirs(home) {
        push_unique_path(&mut paths, &dir);
    }

    for rocm_path in ["/opt/rocm/lib", "/opt/rocm/hip/lib", "/opt/rocm/opencl/lib"] {
        if Path::new(rocm_path).exists() {
            push_unique_path(&mut paths, rocm_path);
        }
    }

    if let Ok(existing) = env::var("LD_LIBRARY_PATH") {
        for part in existing.split(':') {
            push_unique_path(&mut paths, part);
        }
    }

    paths.join(":")
}

fn command_exists(command: &str) -> bool {
    Command::new("bash")
        .arg("-c")
        .arg(format!("command -v {}", command))
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .map(|status| status.success())
        .unwrap_or(false)
}

pub fn python_search_paths() -> Vec<String> {
    let home = resolve_component_user_home();
    let mut paths = vec![
        format!("{}/pytorch", home),
        format!("{}/.local/lib/python3.13/site-packages", home),
        format!("{}/.local/lib/python3.12/site-packages", home),
        format!("{}/rocm_venv/lib/python3.13/site-packages", home),
        format!("{}/rocm_venv/lib/python3.12/site-packages", home),
        format!("{}/megatron/Megatron-LM", home),
        format!("{}/ml_stack/bitsandbytes/bitsandbytes", home),
    ];

    // Dynamically add /opt/rocm/lib/python*/site-packages and /opt/rocm/lib
    let rocm_lib = Path::new("/opt/rocm/lib");
    if rocm_lib.exists() {
        paths.push(rocm_lib.to_string_lossy().to_string());
        if let Ok(entries) = fs::read_dir(rocm_lib) {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_dir() {
                    if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
                        if name.starts_with("python") {
                            let sp = path.join("site-packages");
                            if sp.exists() {
                                paths.push(sp.to_string_lossy().to_string());
                            }
                        }
                    }
                }
            }
        }
    }

    paths
}

fn select_python_for_modules(modules: &[&str], python_candidates: &[String]) -> Option<String> {
    for key in ["MLSTACK_PYTHON_BIN", "UV_PYTHON"] {
        if let Ok(value) = env::var(key) {
            let value = value.trim();
            if !value.is_empty() {
                if modules.is_empty() {
                    return Some(value.to_string());
                }
                for module in modules {
                    if python_has_module(value, module) {
                        return Some(value.to_string());
                    }
                }
            }
        }
    }

    for python in python_candidates {
        for module in modules {
            if python_has_module(python, module) {
                return Some(python.clone());
            }
        }
    }
    None
}

fn path_exists<P: AsRef<Path>>(path: P) -> bool {
    path.as_ref().exists()
}

fn home_path(home: &str, parts: &[&str]) -> PathBuf {
    let mut path = PathBuf::from(home);
    for part in parts {
        path.push(part);
    }
    path
}

/// Read the active Flash Attention backend marker
/// (`~/.mlstack/flash-attention/.backend`) and return `true` iff it equals
/// `name`. The installer writes this file on success; detection requires it so
/// that a leftover build directory or a manual/unknown install can NEVER
/// false-report Flash Attention as "installed".
fn backend_marker_is(name: &str) -> bool {
    let home = resolve_component_user_home();
    let marker = home_path(&home, &[".mlstack", "flash-attention", ".backend"]);
    std::fs::read_to_string(&marker)
        .ok()
        .map(|s| s.trim() == name)
        .unwrap_or(false)
}

/// Multi-modal MIGraphX verification snippet (run via `python -c`).
///
/// Two hardening fixes over the naive `import migraphx; hasattr(parse_onnx)`:
///
/// 1. **CWD excluded from `sys.path`** — otherwise a repo-local `migraphx`
///    stub (e.g. `core/migraphx/migraphx/__init__.py`) shadows a real install
///    when rusty-stack runs from the repo root, importing cleanly but without
///    `parse_onnx`, producing a false `parse_onnx=False` failure.
/// 2. **Driver fallback on *any* incomplete python result** — the previous
///    logic only fell back to `migraphx-driver --version` on `ImportError`. A
///    stub that imports without `parse_onnx` never reached the fallback. Now
///    the C++ driver is tried whenever the python bindings are missing OR
///    incomplete. On Arch/CachyOS (SystemOnly support: no pip `migraphx`
///    wheel) the driver IS the install target, so a clean driver load = pass.
///    This still catches a broken shared-library state (SONAME desync): the
///    driver won't load, so the check can never false-pass.
fn migraphx_verify_snippet() -> &'static str {
    r#"
import os, sys, subprocess
# (1) Drop CWD / '' entries so a repo-local migraphx stub can't shadow a real install.
sys.path = [p for p in sys.path if p and os.path.abspath(p.rstrip('/')) != os.path.abspath(os.getcwd())]
ok = False
ver = 'ok'
via = 'none'
try:
    import migraphx
    if hasattr(migraphx, 'parse_onnx'):
        ok = True
        via = 'python'
        ver = getattr(migraphx, '__version__', 'ok')
    else:
        ver = getattr(migraphx, '__version__', 'stub')
except ImportError:
    ver = 'absent'
# (2) Fall through to the C++ driver whenever python bindings are missing OR
#     incomplete (no parse_onnx). On Arch/CachyOS the driver is the install target.
if not ok:
    r = subprocess.run(['/opt/rocm/bin/migraphx-driver', '--version'], capture_output=True)
    if r.returncode == 0:
        ok = True
        via = 'driver'
        ver = 'driver'
print(f'MIGraphX {ver} parse_onnx={ok} via={via}')
sys.exit(0 if ok else 1)
"#
}

/// Flash Attention verification command (run via `python -c`).
///
/// Verifies `flash_attn` importability (the hard minimum — both ROCm backends
/// install the same `flash_attn` package). When `expected_backend` is `Some`,
/// ALSO requires the `~/.mlstack/flash-attention/.backend` marker to match, so a
/// successful `flash-attn-triton`/`flash-attn-ck` install reaches final success
/// while a stale or manual install cannot. Never accepts the legacy
/// `flash_attention_amd` module — detection dropped it, so verification must too.
fn flash_attention_verification_command(
    label: &str,
    target_id: &str,
    expected_backend: Option<&str>,
    python_candidates: &[String],
) -> VerificationCommand {
    let marker_block = match expected_backend {
        Some(name) => format!(
            "home=pathlib.Path(os.environ.get('HOME',''))\nmarker=home/'.mlstack'/'flash-attention'/'.backend'\nif not marker.exists() or marker.read_text().strip()!='{name}':\n    print(f'backend marker mismatch (expected {name})');sys.exit(1)\n"
        ),
        None => String::new(),
    };
    let code = format!(
        "import importlib,sys,os,pathlib\nm=importlib.import_module('flash_attn')\nver=getattr(m,'__version__','unknown')\nhas_func=hasattr(m,'flash_attn_func') or hasattr(m,'flash_attn_qkvpacked_func')\n{marker_block}print(f'flash_attn {{ver}} func={{has_func}}')"
    );
    python_command(label, target_id, &["flash_attn"], python_candidates, &code)
}

fn env_file_has_enhanced(path: &Path) -> bool {
    if let Ok(contents) = fs::read_to_string(path) {
        return contents.contains("Enhanced ML Stack Environment Setup Script");
    }
    false
}

fn env_file_has_permanent(path: &Path) -> bool {
    if let Ok(contents) = fs::read_to_string(path) {
        let has_marker = contents.contains("Permanent ROCm Environment")
            || contents.contains("Permanent ROCm Env")
            || (contents.contains("ML Stack Environment File")
                && contents.contains("MLSTACK_PYTHON_BIN"));
        if !has_marker {
            return false;
        }
        // Reject gfx000 — invalid placeholder arch
        // Validate arch pattern: gfx(9|10|11|12)NN
        let valid_arch_prefixes = ["gfx9", "gfx10", "gfx11", "gfx12"];
        let mut has_gpu_arch_line = false;
        for line in contents.lines() {
            let trimmed = line.trim();
            // The env file is a sourced shell script, so assignments carry an
            // `export ` prefix (e.g. `export GPU_ARCH=gfx1100`). Strip it before
            // matching, or the check never finds the arch line and a perfectly
            // valid env reports "not configured". Note `GPU_ARCHS=` / other
            // `*_ARCH=` lines do NOT start with `GPU_ARCH=` and are correctly
            // ignored.
            let assign = trimmed.strip_prefix("export ").unwrap_or(trimmed);
            if assign.starts_with("GPU_ARCH=") {
                has_gpu_arch_line = true;
                if let Some(rest) = assign.split("GPU_ARCH=").nth(1) {
                    let arch = rest.split_whitespace().next().unwrap_or(rest);
                    if arch == "gfx000" {
                        return false;
                    }
                    // Check if arch starts with a valid prefix
                    if !valid_arch_prefixes
                        .iter()
                        .any(|prefix| arch.starts_with(prefix))
                    {
                        return false;
                    }
                }
            }
        }
        // Require a GPU_ARCH= line to be present
        has_gpu_arch_line
    } else {
        false
    }
}

fn repo_has_ml_stack_core() -> bool {
    if let Ok(root) = env::var("MLSTACK_REPO_ROOT") {
        return path_exists(Path::new(&root).join("stans_ml_stack"));
    }
    false
}

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

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

    #[test]
    fn test_fastvideo_verification_uses_shared_compiled_kernel_smoke_test() {
        let candidates = vec!["python3".to_string()];
        let commands = component_verification_commands("fastvideo", &candidates);
        assert_eq!(commands.len(), 1);

        let command = &commands[0];
        assert_eq!(command.label, "FastVideo");
        assert_eq!(command.target_id, "fastvideo");
        assert_eq!(command.modules, vec!["fastvideo".to_string()]);
        assert_eq!(
            command.args,
            vec![
                "-c".to_string(),
                crate::installers::components::fastvideo::fastvideo_verification_snippet()
                    .to_string(),
            ]
        );

        let snippet = &command.args[1];
        for anchor in [
            "from fastvideo_kernel import int8_quant",
            "range(int(discovered.stdout.strip()))",
            "int8_quant(source)",
            "torch.cuda.synchronize()",
            "raise SystemExit",
        ] {
            assert!(
                snippet.contains(anchor),
                "missing verification anchor: {anchor}"
            );
        }

        let enhanced = enhanced_verification_commands(&candidates);
        let enhanced_fastvideo: Vec<_> = enhanced
            .iter()
            .filter(|candidate| candidate.target_id == "fastvideo")
            .collect();
        assert_eq!(enhanced_fastvideo.len(), 1);
        assert_eq!(
            enhanced_fastvideo[0].args[1],
            crate::installers::components::fastvideo::fastvideo_verification_snippet()
        );
        assert_eq!(enhanced_fastvideo[0].modules, vec!["fastvideo".to_string()]);
    }

    #[test]
    fn test_fastvideo_installation_status_requires_compiled_kernel_smoke_test() {
        use std::os::unix::fs::PermissionsExt;
        use std::time::{SystemTime, UNIX_EPOCH};

        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system clock must be after Unix epoch")
            .as_nanos();
        let temp_dir = env::temp_dir().join(format!(
            "rusty-stack-fastvideo-status-{}-{unique}",
            std::process::id()
        ));
        fs::create_dir_all(&temp_dir).expect("create temporary FastVideo status test directory");

        let import_only = temp_dir.join("import-only-python");
        fs::write(
            &import_only,
            "#!/bin/sh\ncase \"$2\" in\n  *\"find_spec\"*\"fastvideo\"*) exit 0 ;;\n  *) exit 1 ;;\nesac\n",
        )
        .expect("write import-only fake Python");
        let mut permissions = fs::metadata(&import_only)
            .expect("stat import-only fake Python")
            .permissions();
        permissions.set_mode(0o755);
        fs::set_permissions(&import_only, permissions)
            .expect("make import-only fake Python executable");

        assert!(
            !is_component_installed_by_id(
                "fastvideo",
                &[import_only.to_string_lossy().into_owned()]
            ),
            "an import spec alone must not report FastVideo installed"
        );

        let functional = temp_dir.join("functional-python");
        fs::write(
            &functional,
            "#!/bin/sh\ncase \"$2\" in\n  *\"import fastvideo\"*\"from fastvideo_kernel import int8_quant\"*\"int8_quant(source)\"*) exit 0 ;;\n  *) exit 1 ;;\nesac\n",
        )
        .expect("write functional fake Python");
        let mut permissions = fs::metadata(&functional)
            .expect("stat functional fake Python")
            .permissions();
        permissions.set_mode(0o755);
        fs::set_permissions(&functional, permissions)
            .expect("make functional fake Python executable");

        assert!(
            is_component_installed_by_id("fastvideo", &[functional.to_string_lossy().into_owned()]),
            "the shared compiled-kernel smoke must report FastVideo installed"
        );

        fs::remove_dir_all(&temp_dir).expect("remove temporary FastVideo status test directory");
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_exits_strict() {
        // Stage 4: the snippet must FAIL (exit 1) when torch is not a ROCm build
        // or the HIP runtime is unavailable — it must NOT soft-pass with an
        // unconditional exit 0. Downstream gating is an orchestrator concern.
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("(hip_version is not None) and hip_available"),
            "Should gate the exit on functional HIP availability"
        );
        assert!(
            snippet.contains("sys.exit(0 if ok else 1)"),
            "Should exit non-zero when HIP is unavailable (fail loud)"
        );
        assert!(
            !snippet.contains("\nsys.exit(0)\n"),
            "Should NOT have an unconditional soft-pass sys.exit(0)"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_checks_hip() {
        // The snippet should check for HIP availability to provide diagnostics
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("torch.version"),
            "Should check torch.version for HIP info"
        );
        assert!(
            snippet.contains("torch.cuda.is_available()"),
            "Should check torch.cuda.is_available() for runtime HIP"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_provides_hints() {
        // The snippet should provide actionable hints when HIP is not available
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("HINT:"),
            "Should provide HINT messages for troubleshooting"
        );
        assert!(
            snippet.contains("LD_LIBRARY_PATH"),
            "Should mention LD_LIBRARY_PATH as a possible cause"
        );
        assert!(
            snippet.contains("/opt/rocm/lib"),
            "Should mention ROCm lib path"
        );
    }

    #[test]
    fn test_pytorch_diagnostic_snippet_checks_rocm_libs() {
        // The snippet should check for ROCm library existence
        let snippet = pytorch_diagnostic_snippet();
        assert!(
            snippet.contains("libamdhip64"),
            "Should check for libamdhip64 (core HIP library)"
        );
        assert!(
            snippet.contains("/opt/rocm/lib"),
            "Should check /opt/rocm/lib directory"
        );
    }

    #[test]
    fn test_pytorch_component_verification_uses_diagnostic() {
        // Verify that the pytorch component verification command uses
        // the diagnostic snippet (not the old strict check)
        let candidates = vec!["python3".to_string()];
        let cmds = component_verification_commands("pytorch", &candidates);
        assert_eq!(
            cmds.len(),
            1,
            "Should have exactly one verification command"
        );
        let cmd = &cmds[0];
        assert_eq!(cmd.label, "PyTorch");
        assert_eq!(cmd.target_id, "pytorch");
        // Stage 4: must use the strict FUNCTIONAL check — fail when torch has
        // no ROCm build or the HIP runtime is unavailable (no soft-pass).
        let empty = String::new();
        let code = cmd.args.get(1).unwrap_or(&empty);
        assert!(
            code.contains("(hip_version is not None) and hip_available"),
            "Should gate the exit on functional HIP availability (Stage 4)"
        );
        assert!(
            code.contains("sys.exit(0 if ok else 1)"),
            "Should exit non-zero when HIP is unavailable (Stage 4: fail loud)"
        );
    }

    #[test]
    fn test_pytorch_basic_verification_uses_diagnostic() {
        // Verify that basic verification also uses the strict functional check
        let candidates = vec!["python3".to_string()];
        let cmds = basic_verification_commands(&candidates);
        let pytorch_cmd = cmds.iter().find(|c| c.target_id == "pytorch");
        assert!(
            pytorch_cmd.is_some(),
            "Should have a PyTorch verification command"
        );
        let cmd = pytorch_cmd.unwrap();
        let empty = String::new();
        let code = cmd.args.get(1).unwrap_or(&empty);
        assert!(
            code.contains("(hip_version is not None) and hip_available"),
            "Basic verification should gate exit on functional HIP (Stage 4)"
        );
        assert!(
            code.contains("sys.exit(0 if ok else 1)"),
            "Basic verification should fail loud when HIP unavailable (Stage 4)"
        );
    }

    #[test]
    fn test_pytorch_installation_check_only_imports() {
        // The is_component_installed_by_id check for pytorch should only
        // check if the torch module is importable (not check HIP).
        // This is the correct basic detection — HIP availability is a
        // verification concern, not an installation concern.
        let candidates = vec!["python3".to_string()];
        // We can't assert the result since it depends on the system state,
        // but we can verify the function doesn't panic
        let _ = is_component_installed_by_id("pytorch", &candidates);
    }

    #[test]
    fn test_verification_command_structure() {
        // Verify the VerificationCommand struct for pytorch is well-formed
        let candidates = vec!["python3".to_string()];
        let cmds = component_verification_commands("pytorch", &candidates);
        assert_eq!(cmds.len(), 1);
        let cmd = &cmds[0];
        assert!(!cmd.label.is_empty());
        assert!(!cmd.target_id.is_empty());
        assert!(!cmd.program.is_empty());
        assert_eq!(cmd.args.len(), 2, "Should have -c and the code");
        assert_eq!(cmd.args[0], "-c");
        assert!(!cmd.args[1].is_empty());
        assert!(cmd.modules.contains(&"torch".to_string()));
    }

    #[test]
    fn test_migraphx_snippet_excludes_cwd_and_falls_back_to_driver() {
        // Regression: a repo-local `migraphx` stub imported cleanly but had no
        // parse_onnx, and the old snippet only fell back to the driver on
        // ImportError — so the stub produced a false failure. The new snippet
        // must (a) strip CWD from sys.path and (b) run the driver whenever the
        // python bindings are missing OR incomplete.
        let s = migraphx_verify_snippet();

        // (a) CWD excluded from sys.path (no repo-stub shadowing).
        assert!(
            s.contains("getcwd()"),
            "snippet must drop CWD from sys.path; got: {s}"
        );
        assert!(
            s.contains("sys.path = [p for p in sys.path"),
            "snippet must filter sys.path; got: {s}"
        );

        // (b) Driver fallback is OUTSIDE the try/except — keyed on `if not ok`,
        // so an incomplete import (stub without parse_onnx) still reaches it.
        assert!(
            s.contains("if not ok:"),
            "snippet must fall back to the driver whenever python bindings are incomplete"
        );
        assert!(s.contains("migraphx-driver"));
        assert!(s.contains("--version"));
        // parse_onnx is only accepted inside the successful python branch.
        assert!(s.contains("hasattr(migraphx, 'parse_onnx')"));
        // Exit reflects final ok state (driver fallback can flip it to pass).
        assert!(s.contains("sys.exit(0 if ok else 1)"));
    }

    #[test]
    fn test_onnx_present_snippet_strips_cwd_and_requires_real_package() {
        // Regression: a repo-local `onnxruntime/` source directory resolves as a
        // namespace package (find_spec returns a spec, origin is None), so the
        // old `python_any(&["onnxruntime"])` reported "installed" while
        // `import onnxruntime` had no __version__/providers. The snippet must
        // (a) strip CWD from sys.path and (b) require a real package (origin).
        let s = onnxruntime_present_snippet();
        assert!(
            s.contains("getcwd()"),
            "snippet must drop CWD from sys.path; got: {s}"
        );
        assert!(
            s.contains("sys.path = [p for p in sys.path"),
            "snippet must filter sys.path; got: {s}"
        );
        assert!(
            s.contains("origin is not None"),
            "snippet must require a real package (non-None origin), not a namespace shadow; got: {s}"
        );
    }
}

#[cfg(test)]
mod component_status_tests;