#!/bin/bash

echo "🔍 开始检测服务器 GPU 配置..."
echo ""

# 1. 检查 NVIDIA 驱动
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "1️⃣  检查 NVIDIA 驱动"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
if command -v nvidia-smi &> /dev/null; then
    nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv
    echo ""
    echo "✅ NVIDIA 驱动已安装"
else
    echo "❌ 未检测到 NVIDIA 驱动"
    echo "   请先安装驱动: https://www.nvidia.com/Download/index.aspx"
    exit 1
fi

# 2. 检查 CUDA
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "2️⃣  检查 CUDA 版本"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
if command -v nvcc &> /dev/null; then
    nvcc --version | grep release
    echo "✅ CUDA 工具包已安装"
else
    echo "⚠️  未检测到 CUDA 工具包（可选）"
    echo "   nvidia-smi 显示的 CUDA 版本可以直接使用"
fi

# 3. 检查 Python 和 PyTorch
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "3️⃣  检查 PyTorch GPU 支持"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
if command -v python3 &> /dev/null; then
    python3 << 'EOF'
import sys

try:
    import torch
    print(f"PyTorch 版本: {torch.__version__}")
    print(f"CUDA 可用: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"CUDA 版本: {torch.version.cuda}")
        print(f"cuDNN 版本: {torch.backends.cudnn.version()}")
        print(f"GPU 数量: {torch.cuda.device_count()}")
        print(f"GPU 名称: {torch.cuda.get_device_name(0)}")
        print(f"GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
        print("")
        print("✅ PyTorch GPU 支持已启用")
    else:
        print("")
        print("❌ PyTorch 未启用 GPU 支持")
        print("   需要安装 CUDA 版本的 PyTorch:")
        print("   pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118")
        sys.exit(1)
        
except ImportError:
    print("❌ PyTorch 未安装")
    print("   请先安装: pip install torch torchvision")
    sys.exit(1)
EOF

    if [ $? -eq 0 ]; then
        echo ""
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo "🎉 GPU 配置检测完成！"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo ""
        echo "📊 性能预估:"
        echo "   • 推理速度: 0.1-0.3 秒/次"
        echo "   • 比 CPU 快 3-5 倍"
        echo "   • 用户体验: ⭐⭐⭐⭐⭐"
        echo ""
        echo "🚀 可以开始部署 Florence-2 API 了！"
    else
        echo ""
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo "⚠️  需要安装 CUDA 版本的 PyTorch"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo ""
        echo "请运行以下命令:"
        echo ""
        echo "  # 查看 CUDA 版本"
        echo "  nvidia-smi | grep 'CUDA Version'"
        echo ""
        echo "  # 如果是 CUDA 11.8"
        echo "  pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118"
        echo ""
        echo "  # 如果是 CUDA 12.1"
        echo "  pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121"
        echo ""
    fi
    
else
    echo "❌ Python3 未安装"
fi