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

# ====================== 基础配置（仅保留必要项） ======================
WORK_DIR="/tmp/netcard"
FINAL_XML="${WORK_DIR}/i210_all_passthrough.xml"
FIXED_NET_BUS="0x06"  # 网卡bus号选择0x06无特殊含义，仅作区分用
I210_DRIVER="igb"

# ====================== 参数校验 ======================
if [ $# -eq 0 ]; then
    echo "Usage: $0 <I210 NIC BDF1> <I210 NIC BDF2> ..."
    echo "Example: $0 1a:00.1 03:00.0"
    exit 1
fi
INPUT_BDF_LIST=("$@")

# ====================== 网卡Slot号分配核心函数 ======================
# 获取0x06 Bus下已使用的PCI slot（从已有XML文件中）
get_used_pci_slots_for_net() {
    local xml_file="$1"
    local target_bus="$FIXED_NET_BUS"
    
    if [ ! -f "$xml_file" ]; then
        echo ""
        return 0
    fi
    
    # 过滤出0x06 Bus下的slot号并转十进制
    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
}

# 获取0x06 Bus下一个可用的PCI slot（从0x01开始顺序分配）
get_next_available_slot_for_net() {
    local xml_file="$1"
    local start_slot=1
    
    local used_slots=($(get_used_pci_slots_for_net "$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 $I210_DRIVER; then
    modprobe $I210_DRIVER 2>/dev/null || true
fi

# ====================== 验证I210网卡BDF ======================
VALID_I210_BDF_LIST=()
for BDF in "${INPUT_BDF_LIST[@]}"; do
    # 1. 格式校验（xx:xx.x）
    if ! echo "$BDF" | grep -qE '^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$'; then
        echo "Skipped: Invalid BDF format ($BDF)"
        continue
    fi
    # 2. 存在性校验（系统中是否有该BDF）
    if ! lspci -s "$BDF" >/dev/null 2>&1; then
        echo "Skipped: BDF does not exist ($BDF)"
        continue
    fi
    # 3. 类型校验（是否为Intel I210网卡）
    if ! lspci -nn -s "$BDF" | grep -qi "Intel Corporation I210"; then
        echo "Skipped: Not an I210 NIC ($BDF)"
        continue
    fi
    # 所有校验通过，加入有效列表
    VALID_I210_BDF_LIST+=("$BDF")
done

# 无有效I210 BDF则退出
if [ ${#VALID_I210_BDF_LIST[@]} -eq 0 ]; then
    echo "Error: No valid Intel I210 NIC BDFs!"
    exit 1
fi

# ====================== 核心：处理I210网卡并生成XML ======================
# 初始化：获取0x06 Bus下第一个可用的Slot号
current_slot=$(get_next_available_slot_for_net "$FINAL_XML")
XML_FRAGMENTS=()

echo "===== Start processing I210 NICs ====="
for BDF in "${VALID_I210_BDF_LIST[@]}"; do
    echo ""
    echo "[NIC BDF: $BDF]"
    BDF_FULL="0000:$BDF"

    # 步骤1：绑定IGB驱动（I210原生驱动，确保网卡名可解析）
    echo $BDF_FULL > /sys/bus/pci/drivers/$I210_DRIVER/bind 2>/dev/null || true
    sleep 1  # 等待驱动绑定生效
    # 验证IGB驱动绑定结果
    DRIVER_IN_USE=$(lspci -k -s $BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
    if [ "$DRIVER_IN_USE" != "$I210_DRIVER" ]; then
        echo "Failed: Failed to bind IGB driver (Current driver: $DRIVER_IN_USE)"
        exit 1
    fi

    # 步骤2：解析网卡名（适配实际系统路径）
    NIC_NAME=""
    if [ -d "/sys/bus/pci/devices/$BDF_FULL/net/" ]; then
        NIC_NAME=$(ls /sys/bus/pci/devices/$BDF_FULL/net/ | head -1)
    fi
    if [ -z "$NIC_NAME" ]; then
        NIC_NAME="i210_${BDF//:/_}"
    fi
    echo "NIC name: $NIC_NAME"

    # 步骤3：解析厂商/设备ID
    PCI_INFO=$(lspci -nn | grep -E "[[:space:]]?$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 vendor/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"

    # 步骤4：解绑IGB，绑定VFIO-PCI
    # 先停止网卡网络连接
    if [ -n "$NIC_NAME" ]; then
        systemctl start NetworkManager 2>/dev/null || true
        nmcli connection down $NIC_NAME 2>/dev/null || echo "Tip: No active connection for $NIC_NAME"
        systemctl stop NetworkManager 2>/dev/null || true
    fi
    # 解绑IGB驱动
    echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
    # 绑定VFIO-PCI驱动
    echo $VENDOR_ID $DEVICE_ID > /sys/bus/pci/drivers/vfio-pci/new_id 2>/dev/null || true
    sleep 1  # 等待驱动绑定生效
    # 验证VFIO驱动绑定结果
    DRIVER_IN_USE=$(lspci -k -s $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 (Current driver: $DRIVER_IN_USE)"
        exit 1
    fi
    echo "Success: VFIO driver binding completed"

    # 步骤5：拆分BDF为bus/slot/function（XML配置必需）
    BDF_BUS=$(echo $BDF | cut -d: -f1 | sed 's/^/0x/')
    BDF_SLOT=$(echo $BDF | cut -d: -f2 | cut -d. -f1 | sed 's/^/0x/')
    BDF_FUNC=$(echo $BDF | cut -d. -f2 | sed 's/^/0x/')

    # 步骤6：生成该网卡的XML配置片段（修改：Bus固定0x06，Slot用分配的current_slot）
    XML_FRAGMENT="<!-- I210网卡（$NIC_NAME） BDF=$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_NET_BUS' slot='$current_slot' function='0x0'/>
    </hostdev>"
    XML_FRAGMENTS+=("$XML_FRAGMENT")

    # 递增Slot号（而非Bus号），准备下一个网卡
    next_slot_dec=$((16#${current_slot#0x} + 1))
    current_slot=$(printf "0x%02x" $next_slot_dec)
done

# ====================== 生成最终合并的XML文件 ======================
# 清空旧文件（避免残留）
> "$FINAL_XML"
# 写入极简头部注释
cat >> "$FINAL_XML" << EOF
<!-- I210网卡直通配置 -->
<!-- 有效BDF：${VALID_I210_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: All I210 NIC passthrough configurations completed"
echo "Final XML file: $FINAL_XML"
echo "Next steps: Paste the XML content into the <devices> tag of the VM configuration"