#!/bin/bash
# GPU Passthrough Verification Script
# Checks GPU accessibility in Docker and in GPU-enabled containers

set -e

echo "=== GPU Passthrough Check ==="
echo ""

# Check if nvidia-smi is available on host
echo "[1/4] Checking NVIDIA driver on host..."
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
    echo "✅ NVIDIA driver detected"
else
    echo "❌ nvidia-smi not found on host"
    exit 1
fi
echo ""

# Check if NVIDIA Container Toolkit is installed
echo "[2/4] Checking NVIDIA Container Toolkit..."
if docker run --rm --gpus all nvidia/cuda:11.4.0-base-ubuntu20.04 nvidia-smi &> /dev/null; then
    echo "✅ NVIDIA Container Toolkit working"
else
    echo "❌ NVIDIA Container Toolkit not working"
    echo "   Install with: sudo apt install nvidia-container-toolkit"
    exit 1
fi
echo ""

# Check GPU in Ollama container (if running)
echo "[3/4] Checking GPU in Ollama container..."
if docker ps --format '{{.Names}}' | grep -q "^ollama$"; then
    if docker exec ollama nvidia-smi &> /dev/null; then
        echo "✅ Ollama has GPU access"
        docker exec ollama nvidia-smi --query-gpu=name,utilization.gpu,memory.used --format=csv,noheader
    else
        echo "⚠️  Ollama running but no GPU access"
    fi
else
    echo "⚠️  Ollama container not running"
fi
echo ""

# Check GPU in Jellyfin container (if running)
echo "[4/4] Checking GPU in Jellyfin container..."
if docker ps --format '{{.Names}}' | grep -q "^jellyfin$"; then
    if docker exec jellyfin nvidia-smi &> /dev/null; then
        echo "✅ Jellyfin has GPU access"
        docker exec jellyfin nvidia-smi --query-gpu=name,utilization.gpu,memory.used --format=csv,noheader
    else
        echo "⚠️  Jellyfin running but no GPU access"
    fi
else
    echo "⚠️  Jellyfin container not running"
fi
echo ""

echo "=== GPU Check Complete ==="