#!/bin/bash
# Complete downgrade of all nvidia-container packages to 1.17.9-1

set -e

GREEN='\033[0;32m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'

echo -e "${BLUE}╔═══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║   Downgrade All NVIDIA Container Packages           ║${NC}"
echo -e "${BLUE}╚═══════════════════════════════════════════════════════╝${NC}"
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}✗ Please run as root (use sudo)${NC}"
    exit 1
fi

VERSION="1.17.9-1"

echo -e "${BLUE}==> Step 1: Remove all nvidia-container packages${NC}"
apt-get remove -y nvidia-container-toolkit nvidia-container-toolkit-base libnvidia-container-tools libnvidia-container1 || true
echo -e "${GREEN}✓ Removed packages${NC}"
echo ""

echo -e "${BLUE}==> Step 2: Clean up${NC}"
apt-get autoremove -y
echo -e "${GREEN}✓ Cleaned up${NC}"
echo ""

echo -e "${BLUE}==> Step 3: Install all packages at version $VERSION${NC}"
apt-get install -y \
  nvidia-container-toolkit=$VERSION \
  nvidia-container-toolkit-base=$VERSION \
  libnvidia-container-tools=$VERSION \
  libnvidia-container1=$VERSION
echo -e "${GREEN}✓ Installed all packages at $VERSION${NC}"
echo ""

echo -e "${BLUE}==> Step 4: Hold packages to prevent auto-upgrade${NC}"
apt-mark hold nvidia-container-toolkit nvidia-container-toolkit-base libnvidia-container-tools libnvidia-container1
echo -e "${GREEN}✓ Packages held${NC}"
echo ""

echo -e "${BLUE}==> Step 5: Configure Docker${NC}"
nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json
echo -e "${GREEN}✓ Docker configured${NC}"
echo ""

echo -e "${BLUE}==> Step 6: Rebuild library cache${NC}"
ldconfig
echo -e "${GREEN}✓ Library cache rebuilt${NC}"
echo ""

echo -e "${BLUE}==> Step 7: Restart Docker${NC}"
systemctl restart docker
sleep 3
echo -e "${GREEN}✓ Docker restarted${NC}"
echo ""

echo -e "${BLUE}==> Step 8: Test GPU access${NC}"
echo ""

if docker run --rm --gpus all nvidia/cuda:11.8.0-runtime-ubuntu20.04 nvidia-smi; then
    echo ""
    echo -e "${GREEN}╔═══════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}║                    SUCCESS!                           ║${NC}"
    echo -e "${GREEN}║        GPU is now accessible in Docker!              ║${NC}"
    echo -e "${GREEN}╚═══════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo "Installed versions:"
    dpkg -l | grep nvidia-container
    echo ""
    exit 0
else
    echo ""
    echo -e "${RED}╔═══════════════════════════════════════════════════════╗${NC}"
    echo -e "${RED}║                    FAILED                             ║${NC}"
    echo -e "${RED}║     GPU test still failing                           ║${NC}"
    echo -e "${RED}╚═══════════════════════════════════════════════════════╝${NC}"
    echo ""
    exit 1
fi