#!/bin/bash
# GPU Docker Passthrough - Clean Slate Fix Script
# Removes messy configs and installs working nvidia-container-toolkit version

set -e

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

echo -e "${BLUE}╔═══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║   GPU Docker Passthrough - Clean Slate Fix          ║${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

echo -e "${BLUE}==> Phase 1: Cleanup${NC}"
echo ""

# Remove all our backup files
echo "Removing config backups..."
rm -f /etc/nvidia-container-runtime/config.toml.backup*
rm -f /etc/docker/daemon.json.backup*
echo -e "${GREEN}✓ Removed config backups${NC}"

# Remove our test scripts (keep the final ones)
echo "Cleaning up test scripts..."
rm -f /home/jpmschweitzer/Projects/tower-of-joy/scripts/fix-gpu-docker.py
rm -f /home/jpmschweitzer/Projects/tower-of-joy/scripts/diagnose-gpu-error.py
rm -f /home/jpmschweitzer/Projects/tower-of-joy/scripts/fix-gpu-csv-mode.py
rm -f /home/jpmschweitzer/Projects/tower-of-joy/scripts/try-nvidia-docker2.sh
echo -e "${GREEN}✓ Cleaned up test scripts${NC}"

echo ""
echo -e "${BLUE}==> Phase 2: Check Available Versions${NC}"
echo ""

# Check what versions are available
echo "Checking available nvidia-container-toolkit versions..."
apt-cache policy nvidia-container-toolkit

echo ""
read -p "Do you see version 1.17.x or 1.16.x in the list above? (y/n) " -n 1 -r
echo ""

if [[ $REPLY =~ ^[Yy]$ ]]; then
    echo ""
    echo -e "${BLUE}==> Phase 3: Downgrade to Compatible Version${NC}"
    echo ""

    # Get the version from user
    echo "Available versions from list above:"
    apt-cache madison nvidia-container-toolkit | head -5
    echo ""
    read -p "Enter the version to install (e.g., 1.17.2-1): " VERSION

    echo ""
    echo "Removing nvidia-container-toolkit 1.18..."
    apt-get remove -y nvidia-container-toolkit || true

    echo ""
    echo "Installing nvidia-container-toolkit version $VERSION..."
    apt-get install -y nvidia-container-toolkit=$VERSION

    echo ""
    echo "Holding package to prevent auto-upgrade..."
    apt-mark hold nvidia-container-toolkit

    echo -e "${GREEN}✓ Installed nvidia-container-toolkit $VERSION${NC}"

else
    echo ""
    echo -e "${YELLOW}⚠ Version 1.17/1.16 not available. Configuring 1.18 for legacy mode...${NC}"
    echo ""

    # Configure 1.18 properly for legacy mode
    CONFIG_FILE="/etc/nvidia-container-runtime/config.toml"

    # Backup original
    cp $CONFIG_FILE ${CONFIG_FILE}.original

    # Update config for legacy mode with proper settings
    cat > $CONFIG_FILE << 'EOF'
disable-require = false
supported-driver-capabilities = "compat32,compute,display,graphics,ngx,utility,video"

[nvidia-container-cli]
environment = []
ldcache = "/etc/ld.so.cache"
ldconfig = "@/sbin/ldconfig.real"
load-kmods = true

[nvidia-container-runtime]
log-level = "info"
mode = "legacy"
runtimes = ["runc", "crun"]

[nvidia-container-runtime.modes.legacy]
cuda-compat-mode = "ldconfig"

[nvidia-container-runtime-hook]
path = "nvidia-container-runtime-hook"
skip-mode-detection = false

[nvidia-ctk]
path = "nvidia-ctk"
EOF

    echo -e "${GREEN}✓ Configured nvidia-container-runtime for legacy mode${NC}"
fi

echo ""
echo -e "${BLUE}==> Phase 4: Configure Docker${NC}"
echo ""

# Create minimal Docker daemon.json
DAEMON_JSON="/etc/docker/daemon.json"
cp $DAEMON_JSON ${DAEMON_JSON}.original 2>/dev/null || true

cat > $DAEMON_JSON << 'EOF'
{
    "runtimes": {
        "nvidia": {
            "path": "nvidia-container-runtime",
            "args": []
        }
    }
}
EOF

echo -e "${GREEN}✓ Updated Docker daemon.json${NC}"

# Run nvidia-ctk configure
echo ""
echo "Running nvidia-ctk runtime configure..."
nvidia-ctk runtime configure --runtime=docker --config=$DAEMON_JSON

# Rebuild library cache
echo ""
echo "Rebuilding library cache..."
ldconfig

echo -e "${GREEN}✓ Configuration complete${NC}"

echo ""
echo -e "${BLUE}==> Phase 5: Restart Docker${NC}"
echo ""

systemctl restart docker
sleep 3

echo -e "${GREEN}✓ Docker restarted${NC}"

echo ""
echo -e "${BLUE}==> Phase 6: Test GPU Access${NC}"
echo ""

echo "Testing GPU access in Docker container..."
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 "Configuration saved to:"
    echo "  - /etc/docker/daemon.json"
    echo "  - /etc/nvidia-container-runtime/config.toml"
    echo ""
    echo "Originals backed up to:"
    echo "  - /etc/docker/daemon.json.original"
    echo "  - /etc/nvidia-container-runtime/config.toml.original"
    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 ""
    echo "Troubleshooting:"
    echo "1. Check Docker logs: sudo journalctl -u docker -n 50"
    echo "2. Verify nvidia-container-cli: sudo nvidia-container-cli info"
    echo "3. Check driver version: nvidia-smi"
    echo ""
    exit 1
fi