#!/bin/bash

# 检查参数数量
if [ "$#" -ne 2 ]; then
    echo "【错误】参数不正确！"
    echo "用法: toggle_gpu <PCI物理地址> <to_vm | to_host>"
    echo "示例: toggle_gpu 0000:41:00.0 to_vm"
    echo "查看物理地址请使用 'lspci -Dnn | grep -i nvidia'"
    # 选定你要直通的那张卡，记下它的 VGA 核心地址（如 0000:41:00.0）和对应的 Audio 音频控制器地址（如 0000:41:00.1）。两者必须同时隔离。
    exit 1
fi

VGA_ADDR=$1
# 智能推算同组的音频控制器地址 (将最后的 0 替换为 1)
AUDIO_ADDR="${VGA_ADDR::-1}1"
ACTION=$2

if [ "$ACTION" == "to_vm" ]; then
    echo "🔍 正在检查是否有进程占用 NVIDIA 显卡..."

    # 查找是否有进程在使用 /dev/nvidia* 设备
    # 将标准错误重定向到 null 以避免权限警告，提取结果
    IN_USE=$(lsof /dev/nvidia* 2>/dev/null)

    SERVICES_STOPPED="no"

    if [ -n "$IN_USE" ]; then
        echo -e "\n⚠️  警告：发现有进程正在使用显卡设备！"
        # 打印占用进程的前几行，方便查看
        echo "$IN_USE" | awk 'NR<=6 {print $1, $2, $9}'
        echo -e "...\n"

        read -p "❓ 是否强行停止 vast.ai、Docker 等服务进行清场？(输入 y 继续，其他取消): " choice
        case "$choice" in
          y|Y )
            echo "⏳ 正在停止相关服务..."
            systemctl stop vastai 2>/dev/null
            systemctl stop docker 2>/dev/null
            systemctl stop containerd 2>/dev/null
            systemctl stop nvidia-persistenced 2>/dev/null
            sleep 2 # 等待进程完全释放资源
            SERVICES_STOPPED="yes"
            ;;
          * )
            echo "❌ 操作已取消。为了系统安全，未进行显卡隔离。"
            exit 1
            ;;
        esac
    else
        echo "✅ 当前显卡通道空闲，可以安全隔离。"
    fi

    echo "⏳ 正在将显卡 $VGA_ADDR 及伴生音频 $AUDIO_ADDR 隔离给虚拟机..."
    driverctl set-override $VGA_ADDR vfio-pci
    driverctl set-override $AUDIO_ADDR vfio-pci
    echo "✅ 硬件隔离成功！"

    # 如果刚才停掉了服务，现在将它们拉起，让剩下的卡继续赚钱
    if [ "$SERVICES_STOPPED" == "yes" ]; then
        echo "⏳ 正在恢复 vast.ai 和 Docker 等业务服务..."
        systemctl start nvidia-persistenced 2>/dev/null
        systemctl start containerd 2>/dev/null
        systemctl start docker 2>/dev/null
        systemctl start vastai 2>/dev/null
        echo "✅ 服务已恢复正常！剩下的宿主机显卡已重新上线。"
    fi

elif [ "$ACTION" == "to_host" ]; then
    echo "⏳ 正在解除隔离，将显卡归还给宿主机..."
    driverctl unset-override $VGA_ADDR
    driverctl unset-override $AUDIO_ADDR
    echo "✅ 归还成功！NVIDIA 原生驱动已重新接管该卡。"

    # 归还显卡后，顺便询问是否重启 vast.ai 让它立刻识别到新增加的算力
    echo ""
    read -p "❓ 是否重启 vast.ai/Docker 以确保平台立刻识别到这张归还的显卡？(y/N): " choice
    case "$choice" in
      y|Y )
        echo "⏳ 正在重启平台服务..."
        systemctl restart docker 2>/dev/null
        systemctl restart vastai 2>/dev/null
        echo "✅ 服务重启完毕！所有显卡满血复活投入出租池。"
        ;;
      * )
        echo "✅ 操作完成。宿主机已认出显卡，你可能需要稍后手动重启相关业务服务。"
        ;;
    esac

else
    echo "❌ 错误：未知的操作指令 '$ACTION'。仅支持 'to_vm' 或 'to_host'"
    exit 1
fi

#验证隔离结果,确认目标显卡是否已经被 vfio-pci 成功接管
lspci -nnk -s $VGA_ADDR