#!/bin/bash
set -e

# 检查管理权限
if [ "$(id -u)" -ne 0 ]; then
    echo "请使用sudo运行此脚本"
    exit 1
fi

# 步骤1：卸载系统自带NVIDIA相关包
echo "正在卸载系统自带NVIDIA驱动..."
dpkg -l |grep nvidia |awk '{print $2}'|xargs apt autoremove -y 2>/dev/null
apt autoremove -y

# 步骤2：禁用nouveau驱动
echo "禁用nouveau驱动..."
cat > /etc/modprobe.d/blacklist-nouveau.conf << EOF
blacklist nouveau
options nouveau modeset=0
EOF

# 步骤3：更新initramfs
echo "更新initramfs..."
update-initramfs -u 2>/dev/null

# 步骤4：安装NVIDIA驱动
echo "正在安装NVIDIA驱动..."
INSTALLER=$(find ./ -maxdepth 1 -name "NVIDIA-Linux-*.run" 2>/dev/null | head -n1)

if [ -z "$INSTALLER" ]; then
    echo "错误：未找到NVIDIA驱动安装包"
    echo "请将.run文件放置在脚本当前目录下"
    exit 1
fi

chmod +x "$INSTALLER"
if ! "$INSTALLER" --silent --accept-license --no-x-check --no-questions -M open; then
    echo "驱动安装失败！请检查日志。"
    exit 1
fi

echo "--------------------------------------------------"
echo "NVIDIA驱动安装完成！建议执行以下操作："
echo "1. 运行 nvidia-smi 验证驱动状态"
echo "2. 重启系统确保所有更改生效"
echo "3. 配置nvidia-settings进行显示设置"
echo "稍后将自动返回图形模式..."

sleep 5
# 步骤5：返回图形模式
echo "正在返回图形模式..."
systemctl isolate graphical.target