#!/usr/bin/env bash
#
# gpu-dev-power-install.sh
#
# Multi-distro NVIDIA CUDA + Docker + C/C++ build environment installer.
#
# Supported families:
#   - Parrot OS
#   - Ubuntu
#   - Debian
#   - Linux Mint
#   - Fedora
#   - openSUSE Leap / Tumbleweed
#
# Installs/configures:
#   - Package-managed NVIDIA driver
#   - CUDA Toolkit + nvcc
#   - NVIDIA Container Toolkit
#   - Docker CE + BuildKit/Buildx + Compose plugin
#   - GCC/G++/make/build-essential equivalents
#   - CMake, Ninja, Git, Python, pkg-config and common build utilities
#
# Design goals:
#   - Idempotent / safe to rerun
#   - Retries transient network/package failures
#   - Waits for APT locks instead of killing package managers
#   - Detects Secure Boot and Nouveau
#   - Does NOT auto-reboot
#   - Avoids NVIDIA .run installers
#   - Keeps driver and CUDA toolkit installation separate where practical
#
# Usage:
#   sudo bash gpu-dev-power-install.sh
#   sudo bash gpu-dev-power-install.sh --repair-driver
#   sudo bash gpu-dev-power-install.sh --skip-driver
#   sudo bash gpu-dev-power-install.sh --skip-cuda
#   sudo bash gpu-dev-power-install.sh --skip-docker
#   sudo bash gpu-dev-power-install.sh --skip-container-toolkit
#
set -Eeuo pipefail
shopt -s nullglob

LOG="${LOG:-/var/log/gpu-dev-power-install.log}"
REPAIR_DRIVER=0
SKIP_DRIVER=0
SKIP_CUDA=0
SKIP_DOCKER=0
SKIP_CONTAINER_TOOLKIT=0
NEEDS_REBOOT=0

for arg in "$@"; do
    case "$arg" in
        --repair-driver) REPAIR_DRIVER=1 ;;
        --skip-driver) SKIP_DRIVER=1 ;;
        --skip-cuda) SKIP_CUDA=1 ;;
        --skip-docker) SKIP_DOCKER=1 ;;
        --skip-container-toolkit) SKIP_CONTAINER_TOOLKIT=1 ;;
        -h|--help)
            sed -n '1,45p' "$0"
            exit 0
            ;;
        *)
            echo "Unknown option: $arg" >&2
            exit 2
            ;;
    esac
done

RED=$'\033[0;31m'
GREEN=$'\033[0;32m'
YELLOW=$'\033[1;33m'
BLUE=$'\033[0;34m'
NC=$'\033[0m'

mkdir -p "$(dirname "$LOG")"
touch "$LOG"

log()  { echo "${BLUE}[INFO]${NC} $*" | tee -a "$LOG"; }
ok()   { echo "${GREEN}[ OK ]${NC} $*" | tee -a "$LOG"; }
warn() { echo "${YELLOW}[WARN]${NC} $*" | tee -a "$LOG"; }
die()  { echo "${RED}[FAIL]${NC} $*" | tee -a "$LOG" >&2; exit 1; }

trap 'rc=$?; echo "${RED}[ERROR]${NC} line $LINENO: $BASH_COMMAND (rc=$rc)" | tee -a "$LOG" >&2' ERR

[[ $EUID -eq 0 ]] || die "Run as root: sudo bash $0"

if [[ -r /etc/os-release ]]; then
    # shellcheck disable=SC1091
    . /etc/os-release
else
    die "/etc/os-release not found."
fi

OS_ID="${ID,,}"
OS_ID_LIKE="${ID_LIKE:-}"
OS_VERSION="${VERSION_ID:-}"
OS_CODENAME="${VERSION_CODENAME:-}"
UBUNTU_CODENAME="${UBUNTU_CODENAME:-}"
ARCH="$(uname -m)"
PKG_FAMILY=""
DISTRO_FAMILY=""

case "$OS_ID" in
    parrot)
        DISTRO_FAMILY="parrot"
        PKG_FAMILY="apt"
        ;;
    ubuntu)
        DISTRO_FAMILY="ubuntu"
        PKG_FAMILY="apt"
        ;;
    debian)
        DISTRO_FAMILY="debian"
        PKG_FAMILY="apt"
        ;;
    linuxmint)
        DISTRO_FAMILY="mint"
        PKG_FAMILY="apt"
        ;;
    fedora)
        DISTRO_FAMILY="fedora"
        PKG_FAMILY="dnf"
        ;;
    opensuse*|sles)
        DISTRO_FAMILY="opensuse"
        PKG_FAMILY="zypper"
        ;;
    *)
        die "Unsupported distribution: ${PRETTY_NAME:-$OS_ID}. Supported: Parrot, Ubuntu, Debian, Mint, Fedora, openSUSE."
        ;;
esac

case "$ARCH" in
    x86_64|amd64) NVIDIA_ARCH="x86_64"; DEB_ARCH="amd64" ;;
    aarch64|arm64) NVIDIA_ARCH="sbsa"; DEB_ARCH="arm64" ;;
    *) die "Unsupported architecture for this script: $ARCH" ;;
esac

log "Detected: ${PRETTY_NAME:-$OS_ID} ($DISTRO_FAMILY / $PKG_FAMILY), architecture $ARCH"

retry() {
    local max="${RETRY_MAX:-5}" delay="${RETRY_DELAY:-4}" n=1
    until "$@"; do
        local rc=$?
        if (( n >= max )); then
            warn "Command failed after $max attempts (rc=$rc): $*"
            return "$rc"
        fi
        warn "Attempt $n/$max failed; retrying in ${delay}s: $*"
        sleep "$delay"
        n=$((n+1))
        delay=$((delay*2))
    done
}

url_exists() {
    curl -fsIL --connect-timeout 10 --max-time 25 "$1" >/dev/null 2>&1
}

wait_for_apt() {
    local waited=0
    while fuser /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock \
                /var/cache/apt/archives/lock /var/lib/apt/lists/lock \
                >/dev/null 2>&1; do
        (( waited >= 600 )) && die "APT/dpkg lock remained busy for 10 minutes. Inspect active apt/dpkg processes."
        warn "APT/dpkg is busy; waiting..."
        sleep 10
        waited=$((waited+10))
    done
}

apt_update() {
    wait_for_apt
    retry apt-get -o Acquire::Retries=5 -o DPkg::Lock::Timeout=300 update
}

apt_install() {
    wait_for_apt
    retry apt-get -o Acquire::Retries=5 -o DPkg::Lock::Timeout=300 \
        install -y --no-install-recommends "$@"
}

apt_repair() {
    wait_for_apt
    dpkg --configure -a || true
    apt-get -o DPkg::Lock::Timeout=300 -f install -y || true
    apt_update
}

dnf_install() {
    retry dnf -y --setopt=retries=5 install "$@"
}

zyp_install() {
    retry zypper --non-interactive --gpg-auto-import-keys install -y "$@"
}

install_base_tools() {
    log "Installing base build tools..."
    case "$PKG_FAMILY" in
        apt)
            export DEBIAN_FRONTEND=noninteractive
            export NEEDRESTART_MODE=a
            apt_repair
            apt_install ca-certificates curl wget gnupg gpg lsb-release pciutils usbutils \
                dkms build-essential gcc g++ make ninja-build pkg-config git git-lfs \
                unzip zip tar xz-utils jq python3 python3-pip python3-venv mokutil \
                linux-headers-"$(uname -r)" || {
                    warn "Current-kernel headers package failed; installing generic headers fallback."
                    apt_install linux-headers-generic || true
                }
            ;;
        dnf)
            retry dnf -y makecache
            dnf_install ca-certificates curl wget gnupg2 pciutils usbutils \
                gcc gcc-c++ make ninja-build cmake pkgconf-pkg-config git git-lfs \
                unzip zip tar xz jq python3 python3-pip mokutil \
                kernel-devel kernel-headers dkms dnf-plugins-core
            ;;
        zypper)
            retry zypper --non-interactive refresh
            zyp_install ca-certificates curl wget gpg2 pciutils usbutils \
                gcc gcc-c++ make ninja cmake pkg-config git git-lfs \
                unzip zip tar xz jq python3 python3-pip mokutil \
                kernel-default-devel kernel-devel dkms || true
            ;;
    esac
    ok "Base development tools installed."
}

check_gpu() {
    log "Detecting NVIDIA GPU..."
    if lspci 2>/dev/null | grep -qi nvidia; then
        NVIDIA_PRESENT=1
        lspci | grep -i nvidia | tee -a "$LOG"
        ok "NVIDIA PCI device detected."
    else
        NVIDIA_PRESENT=0
        warn "No NVIDIA PCI device detected. Toolkit installation may still be useful for build hosts."
    fi
}

check_secure_boot() {
    SECURE_BOOT="unknown"
    if command -v mokutil >/dev/null 2>&1; then
        if mokutil --sb-state 2>/dev/null | grep -qi enabled; then
            SECURE_BOOT="enabled"
            warn "Secure Boot is ENABLED. NVIDIA kernel modules may require signing/MOK enrollment."
        elif mokutil --sb-state 2>/dev/null | grep -qi disabled; then
            SECURE_BOOT="disabled"
            ok "Secure Boot is disabled."
        fi
    fi
}

blacklist_nouveau_if_loaded() {
    if lsmod | grep -q '^nouveau'; then
        warn "Nouveau is loaded; blacklisting it for the next boot."
        cat >/etc/modprobe.d/blacklist-nouveau.conf <<'EOF'
blacklist nouveau
options nouveau modeset=0
EOF
        if command -v update-initramfs >/dev/null 2>&1; then
            update-initramfs -u || true
        elif command -v dracut >/dev/null 2>&1; then
            dracut --force || true
        fi
        NEEDS_REBOOT=1
    else
        ok "Nouveau is not loaded."
    fi
}

install_nvidia_driver_apt() {
    case "$DISTRO_FAMILY" in
        ubuntu|mint)
            apt_install ubuntu-drivers-common
            local recommended=""
            recommended="$(ubuntu-drivers devices 2>/dev/null | awk '
              /recommended/ {
                for (i=1;i<=NF;i++) if ($i ~ /^nvidia-driver-[0-9]+/) { print $i; exit }
              }')"
            if [[ -n "$recommended" ]]; then
                log "Ubuntu-family recommended NVIDIA package: $recommended"
                if (( REPAIR_DRIVER )); then
                    apt-get install --reinstall -y "$recommended" || apt_install "$recommended"
                else
                    apt_install "$recommended"
                fi
            else
                warn "No recommended NVIDIA package detected; trying ubuntu-drivers autoinstall."
                retry ubuntu-drivers autoinstall
            fi
            ;;
        debian|parrot)
            # Debian/Parrot package naming. Repositories must expose non-free/non-free-firmware.
            if apt-cache show nvidia-driver >/dev/null 2>&1; then
                if (( REPAIR_DRIVER )); then
                    apt-get install --reinstall -y nvidia-driver || apt_install nvidia-driver
                else
                    apt_install nvidia-driver
                fi
                apt_install firmware-misc-nonfree || true
            else
                warn "nvidia-driver is not visible in APT."
                warn "On Debian/Parrot ensure contrib, non-free and non-free-firmware components are enabled."
                return 1
            fi
            ;;
    esac
}

install_nvidia_driver_fedora() {
    local fv="${OS_VERSION%%.*}"
    [[ "$fv" =~ ^[0-9]+$ ]] || return 1
    log "Enabling RPM Fusion repositories for Fedora $fv..."
    retry dnf -y install \
      "https://download1.rpmfusion.org/free/fedora/rpmfusion-free-release-${fv}.noarch.rpm" \
      "https://download1.rpmfusion.org/nonfree/fedora/rpmfusion-nonfree-release-${fv}.noarch.rpm"
    retry dnf -y makecache
    dnf_install akmod-nvidia xorg-x11-drv-nvidia-cuda
    log "Building/validating akmods; this can take a little while."
    command -v akmods >/dev/null 2>&1 && akmods --force || true
    NEEDS_REBOOT=1
}

install_nvidia_driver_opensuse() {
    log "Using NVIDIA CUDA repository driver packages for openSUSE."
    local repo="https://developer.download.nvidia.com/compute/cuda/repos/opensuse15/${NVIDIA_ARCH}/cuda-opensuse15.repo"
    if url_exists "$repo"; then
        zypper --non-interactive removerepo cuda-nvidia >/dev/null 2>&1 || true
        retry zypper --non-interactive addrepo -f "$repo" cuda-nvidia
        retry zypper --non-interactive --gpg-auto-import-keys refresh
        zyp_install cuda-drivers || {
            warn "cuda-drivers failed; leaving existing driver untouched."
            return 1
        }
        NEEDS_REBOOT=1
    else
        warn "NVIDIA openSUSE repository is unavailable for architecture $NVIDIA_ARCH."
        return 1
    fi
}

install_nvidia_driver() {
    (( SKIP_DRIVER )) && { warn "Skipping NVIDIA driver by request."; return 0; }
    (( NVIDIA_PRESENT == 0 )) && { warn "Skipping driver because no NVIDIA PCI GPU was found."; return 0; }

    log "Installing/updating package-managed NVIDIA driver..."
    case "$PKG_FAMILY" in
        apt) install_nvidia_driver_apt ;;
        dnf) install_nvidia_driver_fedora ;;
        zypper) install_nvidia_driver_opensuse ;;
    esac || warn "Driver installation was not fully successful; diagnostics will run at the end."
}

debian_major_guess() {
    local d=""
    if [[ -r /etc/debian_version ]]; then
        d="$(sed 's/[^0-9].*$//' /etc/debian_version)"
    fi
    if [[ "$d" =~ ^[0-9]+$ ]]; then
        echo "$d"
    else
        echo ""
    fi
}

ubuntu_version_for_derivative() {
    if [[ -n "$UBUNTU_CODENAME" ]]; then
        case "$UBUNTU_CODENAME" in
            jammy) echo "2204" ;;
            noble) echo "2404" ;;
            resolute) echo "2604" ;;
            *) echo "" ;;
        esac
        return
    fi
    case "$OS_VERSION" in
        22.04*) echo "2204" ;;
        24.04*) echo "2404" ;;
        26.04*) echo "2604" ;;
        *) echo "" ;;
    esac
}

choose_cuda_apt_repo() {
    local candidates=()
    if [[ "$DISTRO_FAMILY" == ubuntu || "$DISTRO_FAMILY" == mint ]]; then
        local uv
        uv="$(ubuntu_version_for_derivative)"
        [[ -n "$uv" ]] && candidates+=("ubuntu${uv}")
        candidates+=(ubuntu2604 ubuntu2404 ubuntu2204)
    else
        local dm
        dm="$(debian_major_guess)"
        [[ -n "$dm" ]] && candidates+=("debian${dm}")
        candidates+=(debian13 debian12)
    fi

    local seen="" c keyurl
    for c in "${candidates[@]}"; do
        [[ " $seen " == *" $c "* ]] && continue
        seen+=" $c"
        keyurl="https://developer.download.nvidia.com/compute/cuda/repos/${c}/${NVIDIA_ARCH}/cuda-keyring_1.1-1_all.deb"
        if url_exists "$keyurl"; then
            echo "$c"
            return 0
        fi
    done
    return 1
}

install_cuda_apt() {
    local repo_id
    repo_id="$(choose_cuda_apt_repo)" || {
        warn "Could not find a compatible NVIDIA CUDA APT repository."
        return 1
    }
    local base="https://developer.download.nvidia.com/compute/cuda/repos/${repo_id}/${NVIDIA_ARCH}"
    local deb="/tmp/cuda-keyring.deb"
    log "Using NVIDIA CUDA repository: $repo_id"
    retry curl -fL --retry 3 -o "$deb" "$base/cuda-keyring_1.1-1_all.deb"
    dpkg -i "$deb" || { apt_repair; dpkg -i "$deb"; }
    rm -f "$deb"
    apt_update
    # cuda-toolkit intentionally avoids using the broad 'cuda' meta-package to reduce driver churn.
    apt_install cuda-toolkit
}

install_cuda_fedora() {
    local fv="${OS_VERSION%%.*}"
    local candidates=("$fv")
    if [[ "$fv" =~ ^[0-9]+$ && "$fv" -gt 1 ]]; then candidates+=("$((fv-1))"); fi
    local v repo=""
    for v in "${candidates[@]}"; do
        local r="https://developer.download.nvidia.com/compute/cuda/repos/fedora${v}/${NVIDIA_ARCH}/cuda-fedora${v}.repo"
        if url_exists "$r"; then repo="$r"; break; fi
    done
    [[ -n "$repo" ]] || { warn "No compatible NVIDIA Fedora CUDA repository found."; return 1; }
    log "Adding CUDA repository: $repo"
    curl -fsSL "$repo" -o /etc/yum.repos.d/cuda.repo
    retry dnf -y makecache
    dnf_install cuda-toolkit
}

install_cuda_opensuse() {
    local repo="https://developer.download.nvidia.com/compute/cuda/repos/opensuse15/${NVIDIA_ARCH}/cuda-opensuse15.repo"
    url_exists "$repo" || { warn "No compatible NVIDIA openSUSE CUDA repository found."; return 1; }
    zypper --non-interactive removerepo cuda-toolkit-nvidia >/dev/null 2>&1 || true
    retry zypper --non-interactive addrepo -f "$repo" cuda-toolkit-nvidia
    retry zypper --non-interactive --gpg-auto-import-keys refresh
    zyp_install cuda-toolkit
}

configure_cuda_env() {
    local dirs=(/usr/local/cuda-[0-9]*)
    local newest=""
    if (( ${#dirs[@]} )); then
        newest="$(printf '%s\n' "${dirs[@]}" | sort -V | tail -1)"
    fi
    if [[ -d "$newest" ]]; then
        ln -sfn "$newest" /usr/local/cuda
    fi
    if [[ -d /usr/local/cuda ]]; then
        cat >/etc/profile.d/cuda.sh <<'EOF'
export CUDA_HOME=/usr/local/cuda
export PATH="${CUDA_HOME}/bin:${PATH}"
if [ -d "${CUDA_HOME}/lib64" ]; then
    export LD_LIBRARY_PATH="${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
fi
EOF
        chmod 0644 /etc/profile.d/cuda.sh
        export PATH="/usr/local/cuda/bin:$PATH"
    fi
}

install_cuda() {
    (( SKIP_CUDA )) && { warn "Skipping CUDA Toolkit by request."; return 0; }
    log "Installing latest CUDA Toolkit available for this distro family..."
    case "$PKG_FAMILY" in
        apt) install_cuda_apt ;;
        dnf) install_cuda_fedora ;;
        zypper) install_cuda_opensuse ;;
    esac || warn "CUDA Toolkit installation did not complete."
    configure_cuda_env
}

install_cmake() {
    log "Ensuring CMake is installed."
    case "$PKG_FAMILY" in
        apt)
            # Kitware officially publishes Ubuntu APT repos. Use them for Ubuntu/Mint when supported.
            local base_codename=""
            if [[ "$DISTRO_FAMILY" == ubuntu ]]; then
                base_codename="$OS_CODENAME"
            elif [[ "$DISTRO_FAMILY" == mint ]]; then
                base_codename="$UBUNTU_CODENAME"
            fi
            case "$base_codename" in
                focal|jammy|noble|resolute)
                    install -d -m 0755 /usr/share/keyrings
                    retry curl -fsSL https://apt.kitware.com/keys/kitware-archive-latest.asc \
                        -o /tmp/kitware.asc
                    gpg --dearmor --yes -o /usr/share/keyrings/kitware-archive-keyring.gpg /tmp/kitware.asc
                    rm -f /tmp/kitware.asc
                    cat >/etc/apt/sources.list.d/kitware.list <<EOF
deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ ${base_codename} main
EOF
                    apt_update
                    apt_install cmake
                    ;;
                *)
                    apt_install cmake
                    ;;
            esac
            ;;
        dnf) dnf_install cmake ;;
        zypper) zyp_install cmake ;;
    esac
}

docker_apt_base() {
    case "$DISTRO_FAMILY" in
        ubuntu) echo "ubuntu" ;;
        mint)
            # Main Mint is Ubuntu-based; LMDE would report a Debian-like base.
            if [[ "$OS_ID_LIKE" == *ubuntu* || -n "$UBUNTU_CODENAME" ]]; then echo "ubuntu"; else echo "debian"; fi
            ;;
        parrot|debian) echo "debian" ;;
        *) echo "debian" ;;
    esac
}

docker_apt_suite() {
    local base="$1"
    if [[ "$base" == ubuntu ]]; then
        if [[ -n "$UBUNTU_CODENAME" ]]; then echo "$UBUNTU_CODENAME"; else echo "$OS_CODENAME"; fi
    else
        # Derivatives such as Parrot must point Docker at the corresponding Debian suite.
        local dm
        dm="$(debian_major_guess)"
        case "$dm" in
            13) echo "trixie" ;;
            12) echo "bookworm" ;;
            11) echo "bullseye" ;;
            *) echo "${DEBIAN_CODENAME:-$OS_CODENAME}" ;;
        esac
    fi
}

install_docker_apt() {
    local base suite
    base="$(docker_apt_base)"
    suite="$(docker_apt_suite "$base")"
    [[ -n "$suite" ]] || { warn "Could not determine Docker upstream suite."; return 1; }

    for pkg in docker.io docker-compose docker-compose-v2 docker-doc docker-buildx podman-docker containerd runc; do
        dpkg -s "$pkg" >/dev/null 2>&1 && apt-get remove -y "$pkg" || true
    done

    install -m 0755 -d /etc/apt/keyrings
    retry curl -fsSL "https://download.docker.com/linux/${base}/gpg" -o /etc/apt/keyrings/docker.asc
    chmod a+r /etc/apt/keyrings/docker.asc

    cat >/etc/apt/sources.list.d/docker.sources <<EOF
Types: deb
URIs: https://download.docker.com/linux/${base}
Suites: ${suite}
Components: stable
Architectures: ${DEB_ARCH}
Signed-By: /etc/apt/keyrings/docker.asc
EOF
    apt_update
    apt_install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
}

install_docker_fedora() {
    dnf -y remove docker docker-client docker-client-latest docker-common docker-latest \
        docker-latest-logrotate docker-logrotate docker-selinux docker-engine-selinux docker-engine podman-docker || true
    retry dnf config-manager addrepo --from-repofile https://download.docker.com/linux/fedora/docker-ce.repo
    dnf_install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
}

install_docker_opensuse() {
    # Docker does not maintain a first-class openSUSE CE repo. Prefer the distro engine,
    # while still installing Buildx/Compose where available.
    warn "Docker CE has no first-class openSUSE repository; using openSUSE's Docker packages."
    zyp_install docker docker-compose docker-buildx || zyp_install docker
}

install_docker() {
    (( SKIP_DOCKER )) && { warn "Skipping Docker by request."; return 0; }
    log "Installing Docker engine + BuildKit/Buildx..."
    case "$PKG_FAMILY" in
        apt) install_docker_apt ;;
        dnf) install_docker_fedora ;;
        zypper) install_docker_opensuse ;;
    esac || { warn "Docker installation did not fully complete."; return 0; }

    systemctl enable --now docker || true

    local real_user="${SUDO_USER:-}"
    if [[ -n "$real_user" && "$real_user" != root ]]; then
        getent group docker >/dev/null || groupadd docker || true
        usermod -aG docker "$real_user" || true
        warn "Added '$real_user' to docker group; log out/in after installation."
    fi
}

install_nvidia_container_toolkit() {
    (( SKIP_CONTAINER_TOOLKIT )) && { warn "Skipping NVIDIA Container Toolkit by request."; return 0; }
    (( SKIP_DOCKER )) && { warn "Skipping NVIDIA Container Toolkit because Docker was skipped."; return 0; }

    log "Installing NVIDIA Container Toolkit..."
    case "$PKG_FAMILY" in
        apt)
            install -d -m 0755 /usr/share/keyrings
            retry curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey -o /tmp/nvidia-container.gpg
            gpg --dearmor --yes -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg /tmp/nvidia-container.gpg
            rm -f /tmp/nvidia-container.gpg
            retry curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
                -o /tmp/nvidia-container-toolkit.list
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
                /tmp/nvidia-container-toolkit.list >/etc/apt/sources.list.d/nvidia-container-toolkit.list
            rm -f /tmp/nvidia-container-toolkit.list
            apt_update
            apt_install nvidia-container-toolkit
            ;;
        dnf)
            retry curl -fsSL https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo \
                -o /etc/yum.repos.d/nvidia-container-toolkit.repo
            dnf_install nvidia-container-toolkit
            ;;
        zypper)
            zypper --non-interactive removerepo nvidia-container-toolkit >/dev/null 2>&1 || true
            retry zypper --non-interactive addrepo -f \
                https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo \
                nvidia-container-toolkit
            retry zypper --non-interactive --gpg-auto-import-keys refresh
            zyp_install nvidia-container-toolkit
            ;;
    esac

    if command -v nvidia-ctk >/dev/null 2>&1; then
        nvidia-ctk runtime configure --runtime=docker || warn "nvidia-ctk runtime configuration failed."
        systemctl restart docker || true
    fi
}

driver_diagnostics() {
    echo
    echo "================ NVIDIA DIAGNOSTICS ================" | tee -a "$LOG"
    echo "Kernel: $(uname -r)" | tee -a "$LOG"
    command -v dkms >/dev/null 2>&1 && dkms status | tee -a "$LOG" || true
    lsmod | grep -E '^(nvidia|nouveau)' | tee -a "$LOG" || true
    modinfo nvidia 2>/dev/null | head -20 | tee -a "$LOG" || true

    DRIVER_OK=0
    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi; then
        DRIVER_OK=1
        ok "NVIDIA driver is operational."
    else
        warn "nvidia-smi is not operational yet."
        modprobe nvidia 2>/dev/null || true
        sleep 2
        if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi; then
            DRIVER_OK=1
            ok "NVIDIA driver became operational after modprobe."
        else
            NEEDS_REBOOT=1
        fi
    fi
}

verify_stack() {
    echo
    echo "================ TOOLCHAIN VERSIONS ================" | tee -a "$LOG"
    gcc --version 2>/dev/null | head -1 | tee -a "$LOG" || true
    g++ --version 2>/dev/null | head -1 | tee -a "$LOG" || true
    make --version 2>/dev/null | head -1 | tee -a "$LOG" || true
    cmake --version 2>/dev/null | head -1 | tee -a "$LOG" || true
    ninja --version 2>/dev/null | tee -a "$LOG" || true
    nvcc --version 2>/dev/null | tail -4 | tee -a "$LOG" || true
    docker --version 2>/dev/null | tee -a "$LOG" || true
    docker buildx version 2>/dev/null | tee -a "$LOG" || true
    docker compose version 2>/dev/null | tee -a "$LOG" || true
    nvidia-ctk --version 2>/dev/null | tee -a "$LOG" || true

    if [[ "${DRIVER_OK:-0}" -eq 1 ]] && command -v docker >/dev/null 2>&1; then
        log "Testing GPU access from a CUDA container..."
        if docker run --rm --gpus all nvidia/cuda:13.0.2-base-ubuntu24.04 nvidia-smi; then
            ok "Docker GPU passthrough works."
        else
            warn "GPU container test failed. A reboot may be required; see $LOG."
        fi
    fi
}

main() {
    install_base_tools
    check_gpu
    check_secure_boot
    blacklist_nouveau_if_loaded
    install_nvidia_driver
    install_cuda
    install_cmake
    install_docker
    install_nvidia_container_toolkit
    driver_diagnostics
    verify_stack

    echo
    echo "============================================================"
    echo "Installation complete."
    echo "Log: $LOG"
    echo
    if (( NEEDS_REBOOT )); then
        warn "A reboot is recommended/required before final NVIDIA validation."
    else
        ok "No mandatory reboot was detected, though rebooting after a driver update is still wise."
    fi
    if [[ "$SECURE_BOOT" == enabled ]]; then
        warn "Secure Boot is enabled. Complete any MOK enrollment prompt during reboot."
    fi
    echo
    echo "After reboot run:"
    echo "  nvidia-smi"
    echo "  nvcc --version"
    echo "  cmake --version"
    echo "  docker buildx version"
    echo "  docker run --rm --gpus all nvidia/cuda:13.0.2-base-ubuntu24.04 nvidia-smi"
    echo "============================================================"
}

main "$@"