#!/bin/bash
set -e

# 容错陷阱：出错时提示行号并退出
trap 'echo "Script execution error! Location: $0 Line number: $LINENO"; exit 1' ERR


WORK_DIR="/tmp/gpucard"
FINAL_XML="${WORK_DIR}/gpu_all_passthrough.xml"  # 通用命名
FIXED_GPU_BUS="0x07"  # 保持Bus分配逻辑
GPU_VALID_MODEL_KEYWORD="NVIDIA Corporation.*GeForce RTX"
AUDIO_VALID_MODEL_KEYWORD="Audio device"
GPU_NATIVE_DRIVER="nouveau"  # 显卡原生驱动
AUDIO_NATIVE_DRIVER="snd_hda_intel"  # 声卡原生驱动

# ====================== 参数校验======================
if [ $# -ne 2 ]; then
    echo "Usage: $0 <NVIDIA RTX GPU main BDF> <GPU audio device BDF>"
    echo "Example: $0 01:00.0 01:00.1"
    exit 1
fi
GPU_BDF="$1"          # 显卡主设备BDF（如01:00.0）
AUDIO_BDF="$2"        # 显卡音频设备BDF（如01:00.1）
INPUT_BDF_LIST=("$GPU_BDF" "$AUDIO_BDF")

# ====================== 显卡Slot号分配核心函数 ======================
get_used_pci_slots_for_gpu() {
    local xml_file="$1"
    local target_bus="$FIXED_GPU_BUS"
    
    if [ ! -f "$xml_file" ]; then
        echo ""
        return 0
    fi
    

    grep -B1 -A5 "bus='$target_bus'" "$xml_file" 2>/dev/null | grep -o "slot='0x[0-9a-f]*'" | sed "s/slot='0x//" | sed "s/'//" | while read hex; do
        printf "%d\n" "0x$hex" 2>/dev/null || echo 0
    done | sort -n
}

get_next_available_slot_for_gpu() {
    local xml_file="$1"
    local start_slot=1
    
    local used_slots=($(get_used_pci_slots_for_gpu "$xml_file"))
    local next_slot=$start_slot
    
    for used_slot in "${used_slots[@]}"; do
        if [ "$used_slot" -eq "$next_slot" ] 2>/dev/null; then
            ((next_slot++))
        elif [ "$used_slot" -gt "$next_slot" ] 2>/dev/null; then
            break
        fi
    done
    
    printf "0x%02x" $next_slot
}

# ====================== 前置准备：替换网卡驱动 ======================
mkdir -p "$WORK_DIR"

# 加载VFIO驱动（和网卡脚本完全一致）
if ! lsmod | grep -q vfio_pci; then
    modprobe vfio 2>/dev/null || true
    modprobe vfio-pci 2>/dev/null || true
fi
if [ ! -d "/sys/bus/pci/drivers/vfio-pci/" ]; then
    echo "Error: VFIO-PCI driver load failed!"
    exit 1
fi

# 卸载显卡原生驱动（替换网卡的"加载I210驱动"逻辑）
if lsmod | grep -q $GPU_NATIVE_DRIVER; then
    rmmod $GPU_NATIVE_DRIVER 2>/dev/null || true
fi

# ====================== 验证NVIDIA RTX显卡+音频BDF ======================
VALID_GPU_BDF_LIST=()

# 校验显卡主设备BDF（适配所有RTX显卡）
if ! echo "$GPU_BDF" | grep -qE '^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$'; then
    echo "Error: Invalid format for GPU main BDF ($GPU_BDF)"
    exit 1
fi
if ! lspci -s "$GPU_BDF" >/dev/null 2>&1; then
    echo "Error: GPU main BDF does not exist ($GPU_BDF)"
    exit 1
fi
if ! lspci -nn -s "$GPU_BDF" | grep -qi "$GPU_VALID_MODEL_KEYWORD"; then
    echo "Error: $GPU_BDF is not the main device of NVIDIA GeForce RTX series GPU"
    exit 1
fi
VALID_GPU_BDF_LIST+=("$GPU_BDF")

# 校验显卡音频设备BDF
if ! echo "$AUDIO_BDF" | grep -qE '^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$'; then
    echo "Error: Invalid format for audio BDF ($AUDIO_BDF)"
    exit 1
fi
if ! lspci -s "$AUDIO_BDF" >/dev/null 2>&1; then
    echo "Error: Audio BDF does not exist ($AUDIO_BDF)"
    exit 1
fi
if ! lspci -nn -s "$AUDIO_BDF" | grep -qi "NVIDIA Corporation" || ! lspci -nn -s "$AUDIO_BDF" | grep -qi "$AUDIO_VALID_MODEL_KEYWORD"; then
    echo "Error: $AUDIO_BDF is not the audio device of NVIDIA GPU"
    exit 1
fi
VALID_GPU_BDF_LIST+=("$AUDIO_BDF")

# 无有效BDF则退出
if [ ${#VALID_GPU_BDF_LIST[@]} -eq 0 ]; then
    echo "Error: No valid BDFs for NVIDIA RTX GPU or audio device!"
    exit 1
fi

# ====================== 核心：处理NVIDIA RTX显卡+音频并生成XML ======================
# 初始化：获取0x07 Bus下第一个可用的Slot号（单卡固定为0x01）
current_slot=$(get_next_available_slot_for_gpu "$FINAL_XML")
XML_FRAGMENTS=()

echo "===== Start processing NVIDIA RTX GPU + audio device ====="
# 处理显卡主设备
echo ""
echo "[GPU main BDF: $GPU_BDF]"
BDF_FULL="0000:$GPU_BDF"

# 先获取当前驱动状态
CURRENT_GPU_DRIVER=$(lspci -k -s $GPU_BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
echo "Current driver status: ${CURRENT_GPU_DRIVER:-No driver}"
if [ -n "$CURRENT_GPU_DRIVER" ]; then
    # 有驱动：分原生/非原生处理
    if [ "$CURRENT_GPU_DRIVER" = "$GPU_NATIVE_DRIVER" ]; then
        echo "Detected native driver ($GPU_NATIVE_DRIVER), starting to unbind and unload module..."
        # 解绑驱动
        if [ -d "/sys/bus/pci/devices/$BDF_FULL/driver/" ]; then
            echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
        fi
        # 卸载nouveau模块（确保彻底移除）
        rmmod $GPU_NATIVE_DRIVER 2>/dev/null || true
    else
        echo "Detected non-native driver ($CURRENT_GPU_DRIVER), starting to unbind..."
        # 解绑当前驱动
        if [ -d "/sys/bus/pci/devices/$BDF_FULL/driver/" ]; then
            echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
        fi
    fi
else
    # 无驱动：跳过解绑
    echo "No driver in use, skipping unbinding step"
fi
sleep 1
# 步骤2：解析PCI ID
PCI_INFO=$(lspci -nn | grep -E "[[:space:]]?$GPU_BDF")
PCI_ID=$(echo "$PCI_INFO" | grep -oP '\[\K[0-9a-fA-F]{4}:[0-9a-fA-F]{4}(?=\])' | tr -d ' ')
if [ -z "$PCI_ID" ]; then
    echo "Failed: Failed to parse GPU main device ID"
    exit 1
fi
VENDOR_ID=$(echo $PCI_ID | cut -d: -f1)
DEVICE_ID=$(echo $PCI_ID | cut -d: -f2)
echo "Vendor ID:Device ID: $VENDOR_ID:$DEVICE_ID"
# 步骤3：绑定VFIO驱动
echo $VENDOR_ID $DEVICE_ID > /sys/bus/pci/drivers/vfio-pci/new_id 2>/dev/null || true
sleep 1
DRIVER_IN_USE=$(lspci -k -s $GPU_BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
if [ "$DRIVER_IN_USE" != "vfio-pci" ]; then
    echo "Failed: Failed to bind VFIO driver to GPU main device (Current driver: $DRIVER_IN_USE)"
    exit 1
fi
echo "Success: VFIO driver binding for GPU main device completed"
# 步骤4：拆分BDF
BDF_BUS=$(echo $GPU_BDF | cut -d: -f1 | sed 's/^/0x/')
BDF_SLOT=$(echo $GPU_BDF | cut -d: -f2 | cut -d. -f1 | sed 's/^/0x/')
BDF_FUNC=$(echo $GPU_BDF | cut -d. -f2 | sed 's/^/0x/')
# 步骤5：生成主设备XML
XML_FRAGMENT="<!-- NVIDIA RTX显卡主设备（BDF=$GPU_BDF） -->
<hostdev mode='subsystem' type='pci' managed='yes'>
  <driver name='vfio'/>
  <source>
    <address domain='0x0000' bus='$BDF_BUS' slot='$BDF_SLOT' function='$BDF_FUNC'/>
  </source>
  <address type='pci' domain='0x0000' bus='$FIXED_GPU_BUS' slot='$current_slot' function='0x0' multifunction='on'/>
</hostdev>"
XML_FRAGMENTS+=("$XML_FRAGMENT")

# 处理显卡音频设备（复用同一个slot，function=0x1）
echo ""
echo "[Audio BDF: $AUDIO_BDF]"
BDF_FULL="0000:$AUDIO_BDF"
# 先获取当前驱动状态
CURRENT_AUDIO_DRIVER=$(lspci -k -s $AUDIO_BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
echo "Current driver status: ${CURRENT_AUDIO_DRIVER:-No driver}"
if [ -n "$CURRENT_AUDIO_DRIVER" ]; then
    # 有驱动：分原生/非原生处理
    if [ "$CURRENT_AUDIO_DRIVER" = "$AUDIO_NATIVE_DRIVER" ]; then
        echo "Detected native driver ($AUDIO_NATIVE_DRIVER), starting to unbind..."
        # 解绑驱动
        if [ -d "/sys/bus/pci/devices/$BDF_FULL/driver/" ]; then
            echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
        fi
    else
        echo "Detected non-native driver ($CURRENT_AUDIO_DRIVER), starting to unbind..."
        # 解绑当前驱动
        if [ -d "/sys/bus/pci/devices/$BDF_FULL/driver/" ]; then
            echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
        fi
    fi
else
    # 无驱动：跳过解绑
    echo "No driver in use, skipping unbinding step"
fi
sleep 1
# 步骤2：解析PCI ID
PCI_INFO=$(lspci -nn | grep -E "[[:space:]]?$AUDIO_BDF")
PCI_ID=$(echo "$PCI_INFO" | grep -oP '\[\K[0-9a-fA-F]{4}:[0-9a-fA-F]{4}(?=\])' | tr -d ' ')
if [ -z "$PCI_ID" ]; then
    echo "Failed: Failed to parse audio device ID"
    exit 1
fi
VENDOR_ID=$(echo $PCI_ID | cut -d: -f1)
DEVICE_ID=$(echo $PCI_ID | cut -d: -f2)
echo "Vendor ID:Device ID: $VENDOR_ID:$DEVICE_ID"
# 步骤3：绑定VFIO驱动
echo $VENDOR_ID $DEVICE_ID > /sys/bus/pci/drivers/vfio-pci/new_id 2>/dev/null || true
sleep 1
DRIVER_IN_USE=$(lspci -k -s $AUDIO_BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
if [ "$DRIVER_IN_USE" != "vfio-pci" ]; then
    echo "Failed: Failed to bind VFIO driver to audio device (Current driver: $DRIVER_IN_USE)"
    exit 1
fi
echo "Success: VFIO driver binding for audio device completed"
# 步骤4：拆分BDF
BDF_BUS=$(echo $AUDIO_BDF | cut -d: -f1 | sed 's/^/0x/')
BDF_SLOT=$(echo $AUDIO_BDF | cut -d: -f2 | cut -d. -f1 | sed 's/^/0x/')
BDF_FUNC=$(echo $AUDIO_BDF | cut -d. -f2 | sed 's/^/0x/')
# 步骤5：生成音频设备XML
XML_FRAGMENT="<!-- NVIDIA RTX显卡音频设备（BDF=$AUDIO_BDF） -->
<hostdev mode='subsystem' type='pci' managed='yes'>
  <driver name='vfio'/>
  <source>
    <address domain='0x0000' bus='$BDF_BUS' slot='$BDF_SLOT' function='$BDF_FUNC'/>
  </source>
  <address type='pci' domain='0x0000' bus='$FIXED_GPU_BUS' slot='$current_slot' function='0x1'/>
</hostdev>"
XML_FRAGMENTS+=("$XML_FRAGMENT")

# ====================== 生成最终合并的XML文件 ======================
# 清空旧文件
> "$FINAL_XML"
# 写入极简头部注释
cat >> "$FINAL_XML" << EOF
<!-- NVIDIA RTX显卡+音频设备直通配置 -->
<!-- 有效BDF：${VALID_GPU_BDF_LIST[*]} -->

EOF
# 写入所有显卡的XML配置片段
for FRAGMENT in "${XML_FRAGMENTS[@]}"; do
    echo "$FRAGMENT" >> "$FINAL_XML"
    echo "" >> "$FINAL_XML"
    echo "" >> "$FINAL_XML"
done

# ====================== 最终输出 ======================
echo ""
echo "===== Processing completed ====="
echo "Success: NVIDIA RTX GPU + audio device passthrough configuration completed"
echo "Final XML file: $FINAL_XML"
echo "Next steps: Paste the XML content into the <devices> tag of the VM configuration"