#!/usr/bin/env bash
# cuda_toolkit_check — Detecta componentes CUDA instalados en el sistema.
#
# Emite pares key=value a stdout:
#   nvcc=<version|missing>
#   nvidia_smi=<present|missing>
#   driver_version=<version|missing>
#   cuda_libs=<path|missing>
#   overall=<ok|partial|missing>
#
# Exit code 0 siempre (funcion informativa, no fatal).
# Idempotente: se puede invocar multiples veces sin efectos secundarios.

cuda_toolkit_check() {
    local nvcc_ver="missing"
    local nvidia_smi_status="missing"
    local driver_version="missing"
    local cuda_libs_path="missing"

    # --- nvcc ---
    if command -v nvcc &>/dev/null; then
        # nvcc --version imprime algo como:
        #   Cuda compilation tools, release 12.4, V12.4.131
        local raw
        raw="$(nvcc --version 2>&1)"
        # Extraer "12.4" de "release 12.4,"
        local ver
        ver="$(echo "$raw" | grep -oP 'release \K[0-9]+\.[0-9]+')"
        nvcc_ver="${ver:-present}"
    fi

    # --- nvidia-smi + driver_version ---
    if command -v nvidia-smi &>/dev/null; then
        nvidia_smi_status="present"
        # nvidia-smi --query-gpu=driver_version --format=csv,noheader retorna la version
        local drv
        drv="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -n1 | tr -d ' ')"
        if [[ -n "$drv" ]]; then
            driver_version="$drv"
        fi
    fi

    # --- cuda_libs: buscar en rutas canonicas ---
    local search_dirs=(
        "/usr/local/cuda"
        "/usr/local/cuda-"*
        "/opt/cuda"
        "/opt/cuda-"*
        "/usr/lib/x86_64-linux-gnu/libcuda.so"*
        "/usr/lib/aarch64-linux-gnu/libcuda.so"*
    )

    for candidate in "${search_dirs[@]}"; do
        # shellcheck disable=SC2206
        # Expandir globs: si el candidato no existe el glob no expande
        for path in $candidate; do
            if [[ -e "$path" ]]; then
                # Normalizar: tomar solo el directorio raiz /usr/local/cuda*
                local base
                base="${path%%/lib*}"
                cuda_libs_path="$base"
                break 2
            fi
        done
    done

    # Si no encontramos directorio CUDA pero si libcuda.so en rutas de lib estandar
    if [[ "$cuda_libs_path" == "missing" ]]; then
        local libcuda
        libcuda="$(ldconfig -p 2>/dev/null | grep 'libcuda\.so' | head -n1 | awk '{print $NF}')"
        if [[ -n "$libcuda" ]]; then
            cuda_libs_path="$(dirname "$libcuda")"
        fi
    fi

    # --- overall ---
    local found_count=0
    [[ "$nvcc_ver"          != "missing" ]] && ((found_count++))
    [[ "$nvidia_smi_status" != "missing" ]] && ((found_count++))
    [[ "$cuda_libs_path"    != "missing" ]] && ((found_count++))

    local overall
    if   [[ $found_count -eq 0 ]]; then overall="missing"
    elif [[ $found_count -eq 3 ]]; then overall="ok"
    else                                  overall="partial"
    fi

    # --- emitir resultados ---
    echo "nvcc=${nvcc_ver}"
    echo "nvidia_smi=${nvidia_smi_status}"
    echo "driver_version=${driver_version}"
    echo "cuda_libs=${cuda_libs_path}"
    echo "overall=${overall}"
}

# Ejecutar si se invoca directamente
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    cuda_toolkit_check "$@"
fi