#!/bin/bash
# Установка драйвера NVIDIA и CUDA на Ubuntu 24.04 для работы с нейросетями (RTX 40xx).
# Запуск: sudo bash setup-nvidia-cuda-ubuntu24.sh
# После выполнения нужна перезагрузка: sudo reboot

set -e

if [ "$(id -u)" -ne 0 ]; then
  echo "Запустите скрипт с sudo."
  exit 1
fi

echo "=== 1/6 Отключение драйвера Nouveau ==="
cat > /etc/modprobe.d/blacklist-nvidia-nouveau.conf << 'EOF'
blacklist nouveau
options nouveau modeset=0
EOF
update-initramfs -u 2>/dev/null || true

echo "=== 2/6 Обновление системы и установка зависимостей ==="
export DEBIAN_FRONTEND=noninteractive
apt-get update -qq
apt-get install -y build-essential wget

echo "=== 3/6 Установка драйвера NVIDIA (рекомендуемый для железа) ==="
# Для RTX 4060 Ti подойдёт драйвер 550+; ubuntu-drivers выберет подходящий
apt-get install -y ubuntu-drivers-common
DRIVER=$(ubuntu-drivers list 2>/dev/null | grep -m1 "nvidia-driver-" || echo "nvidia-driver-560")
if apt-cache show "$DRIVER" &>/dev/null; then
  apt-get install -y "$DRIVER"
else
  apt-get install -y nvidia-driver-560 || apt-get install -y nvidia-driver-550 || ubuntu-drivers autoinstall
fi

echo "=== 4/6 Добавление репозитория CUDA и установка CUDA Toolkit 12 ==="
KEYRING_DEB="/tmp/cuda-keyring.deb"
wget -q -O "$KEYRING_DEB" "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb" || {
  echo "Ошибка загрузки cuda-keyring. Проверьте сеть."
  exit 1
}
dpkg -i "$KEYRING_DEB"
rm -f "$KEYRING_DEB"
apt-get update -qq
apt-get install -y cuda-toolkit-12-6 || apt-get install -y cuda-toolkit-12-5 || apt-get install -y cuda

echo "=== 5/6 Настройка окружения CUDA (PATH и библиотеки) ==="
# Ubuntu/NVIDIA repo ставит в /usr/local/cuda-12.6, симлинк /usr/local/cuda создаётся пакетом cuda
for CUDA_ROOT in /usr/local/cuda /usr/local/cuda-12.6 /usr/local/cuda-12.5; do
  if [ -d "$CUDA_ROOT" ]; then
    cat > /etc/profile.d/cuda.sh << EOF
# CUDA for neural networks
export PATH=$CUDA_ROOT/bin:\$PATH
export LD_LIBRARY_PATH=$CUDA_ROOT/lib64:\${LD_LIBRARY_PATH:+:\$LD_LIBRARY_PATH}
EOF
    chmod 644 /etc/profile.d/cuda.sh
    echo "Файл /etc/profile.d/cuda.sh создан (CUDA_ROOT=$CUDA_ROOT)."
    break
  fi
done
if [ ! -f /etc/profile.d/cuda.sh ]; then
  echo "Предупреждение: каталог CUDA не найден. После установки cuda-toolkit выполните: sudo bash -c 'echo \"export PATH=/usr/local/cuda/bin:\\\$PATH\" > /etc/profile.d/cuda.sh'"
fi

echo "=== 6/6 Готово ==="
echo ""
echo "Драйвер NVIDIA и CUDA Toolkit установлены."
echo "ОБЯЗАТЕЛЬНО перезагрузите сервер для загрузки драйвера:"
echo "  sudo reboot"
echo ""
echo "После перезагрузки проверьте:"
echo "  nvidia-smi"
echo "  source /etc/profile.d/cuda.sh && nvcc --version"
echo "  python3 -c 'import torch; print(torch.cuda.is_available())'  # если ставите PyTorch"