#!/bin/bash

# === Универсальный установщик драйверов NVIDIA для Ubuntu 24.04+ ===
# Автоматически определяет GPU, устанавливает драйверы и настраивает Docker

# --- Проверка прав ---
if [ "$(id -u)" -ne 0 ]; then
    echo "🚫 Запустите скрипт с sudo!" >&2
    exit 1
fi

# --- Определение GPU ---
detect_gpu() {
    echo "🔍 Поиск NVIDIA GPU..."
    if ! lspci | grep -i "nvidia" >/dev/null; then
        echo "❌ NVIDIA GPU не обнаружена!"
        echo "Проверьте:"
        echo "1. Физическое подключение карты"
        echo "2. Поддержку PCIe в BIOS"
        exit 1
    fi

    GPU_MODEL=$(lspci | grep -i "nvidia" | head -n 1 | cut -d ':' -f 4 | sed 's/^[ \t]*//;s/[ \t]*$//')
    GPU_ID=$(lspci -n | grep -i "10de" | cut -d ' ' -f 3 | cut -d ':' -f 2)

    case "$GPU_ID" in
        "2230"|"2231") DRIVER="470" ;;  # Tesla T4
        "1db6")        DRIVER="525" ;;  # Tesla A10
        "20b5")        DRIVER="535" ;;  # RTX A6000
        *)             DRIVER="550" ;;  # По умолчанию (новые карты)
    esac

    echo "✅ Обнаружена видеокарта: $GPU_MODEL (ID: 10de:$GPU_ID)"
    echo "🛠️  Рекомендуемый драйвер: $DRIVER"
}

# --- Установка драйверов ---
install_drivers() {
    echo "🧹 Очистка старых драйверов..."
    apt-get purge -y *nvidia* *cuda* *cudnn* 2>/dev/null
    apt-get autoremove -y
    rm -rf /usr/local/cuda*

    echo "📦 Установка зависимостей..."
    apt-get update
    apt-get install -y --no-install-recommends \
        ubuntu-drivers-common \
        linux-headers-$(uname -r) \
        build-essential

    echo "⚙️ Установка драйвера версии $DRIVER..."
    apt-get install -y nvidia-driver-$DRIVER nvidia-dkms-$DRIVER
}

# --- Проверка драйвера ---
verify_driver() {
    if ! nvidia-smi &>/dev/null; then
        echo "❌ Ошибка: nvidia-smi не работает!"
        echo "Возможные решения:"
        echo "1. Отключите Secure Boot в BIOS"
        echo "2. Выполните: sudo apt-get install --reinstall nvidia-dkms-$DRIVER"
        exit 1
    fi
}

# --- Настройка Docker ---
setup_docker() {
    echo "🐳 Настройка Docker для NVIDIA GPU..."
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo 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' | \
        sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list >/dev/null

    apt-get update
    apt-get install -y nvidia-container-toolkit
    systemctl restart docker

    echo "🔌 Проверка Docker + GPU..."
    if ! docker run --rm --gpus all nvidia/cuda:12.3.1-base nvidia-smi &>/dev/null; then
        echo "⚠️  Docker не видит GPU! Попробуйте:"
        echo "1. sudo systemctl restart docker"
        echo "2. Добавьте пользователя в группу docker: sudo usermod -aG docker $USER"
    fi
}

# --- Главный процесс ---
clear
echo "🖥️  Универсальный установщик NVIDIA GPU для Ubuntu"
echo "▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔"

detect_gpu
install_drivers
verify_driver
setup_docker

# --- Итоговая информация ---
echo ""
echo "✅ Установка завершена!"
echo "▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔"
echo "Информация о системе:"
nvidia-smi
echo ""
echo "Для использования GPU в контейнерах:"
echo "  docker run --gpus all ваш-образ"
echo ""
echo "Рекомендуемые базовые образы:"
echo "  • nvidia/cuda:12.3.1-cudnn8-runtime (для PyTorch/TensorFlow)"
echo "  • nvcr.io/nvidia/pytorch:23.10-py3 (официальный PyTorch)"
echo "▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔▔"

# Перезагрузка (опционально)
read -p "Перезагрузить систему сейчас? [y/N]: " reboot_choice
if [[ $reboot_choice =~ ^[Yy]$ ]]; then
    shutdown -r now
fi