#!/bin/bash
set -euo pipefail

# =============================================================================
# NVIDIA Container Toolkit + Docker Installation Script
# Enables GPU-accelerated Docker containers
# Tested on Ubuntu 20.04/22.04/24.04, Debian 11/12
# =============================================================================

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

log()   { echo -e "${GREEN}[INFO]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
error() { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }

# --- Pre-flight checks -------------------------------------------------------
[[ $EUID -ne 0 ]] && error "Please run as root or with sudo."

log "Checking for NVIDIA GPU..."
if ! lspci | grep -qi nvidia; then
    error "No NVIDIA GPU detected. Aborting."
fi

log "Checking for NVIDIA driver..."
if ! nvidia-smi &>/dev/null; then
    warn "NVIDIA driver not found. Attempting to install recommended driver..."
    apt-get update -qq
    apt-get install -y ubuntu-drivers-common
    ubuntu-drivers autoinstall
    warn "NVIDIA driver installed. A REBOOT may be required before proceeding."
    read -rp "Continue anyway? (y/N): " choice
    [[ "$choice" != [yY] ]] && { log "Reboot and re-run this script."; exit 0; }
fi

nvidia-smi
log "NVIDIA driver is working."

# --- Install Docker -----------------------------------------------------------
if ! command -v docker &>/dev/null; then
    log "Docker not found. Installing Docker Engine..."
    apt-get update -qq
    apt-get install -y ca-certificates curl gnupg

    install -m 0755 -d /etc/apt/keyrings
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | \
        gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    chmod a+r /etc/apt/keyrings/docker.gpg

    # Detect distro (works for Ubuntu & Debian)
    DISTRO_ID=$(. /etc/os-release && echo "$ID")
    DISTRO_CODENAME=$(. /etc/os-release && echo "$VERSION_CODENAME")

    echo \
      "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] \
      https://download.docker.com/linux/${DISTRO_ID} ${DISTRO_CODENAME} stable" | \
      tee /etc/apt/sources.list.d/docker.list > /dev/null

    apt-get update -qq
    apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

    systemctl enable --now docker
    log "Docker installed successfully."
else
    log "Docker is already installed: $(docker --version)"
fi

# --- Install NVIDIA Container Toolkit ----------------------------------------
log "Adding NVIDIA Container Toolkit repository..."

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list > /dev/null

apt-get update -qq
apt-get install -y nvidia-container-toolkit

# --- Configure Docker runtime -------------------------------------------------
log "Configuring Docker to use NVIDIA runtime..."
nvidia-ctk runtime configure --runtime=docker

systemctl restart docker
log "Docker restarted with NVIDIA runtime."

# --- Verify -------------------------------------------------------------------
log "Running verification: docker run --gpus all nvidia/cuda nvidia-smi"
echo "---"
docker run --rm --gpus all nvidia/cuda:12.6.3-base-ubuntu24.04 nvidia-smi

echo ""
log "============================================="
log " NVIDIA Docker setup complete!"
log " Usage: docker run --gpus all <image>"
log "============================================="