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

# ====================== 基础配置 ======================
I210_DRIVER="igb"  # I210原生驱动
VFIO_DRIVER="vfio-pci"  # 直通驱动
WORK_DIR="/tmp/netcard"
LOG_FILE="${WORK_DIR}/i210_restore.log"

# ====================== 颜色定义 ======================
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# ====================== 前置准备 ======================
mkdir -p "$WORK_DIR"
> "$LOG_FILE"  # 清空旧日志


# ====================== 日志函数 ======================
log() {
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] $1" >> "$LOG_FILE"
    echo -e "[$timestamp] $1"
}

# ====================== 参数校验 ======================
if [ $# -eq 0 ]; then
    echo -e "\n${YELLOW}Usage: $0 <I210 NIC BDF1> <I210 NIC BDF2> ...${NC}"
    echo -e "${YELLOW}Example: $0 1a:00.1 03:00.0${NC}"
    echo -e "${YELLOW}Tip: If no BDF is specified, all I210 NICs bound to VFIO will be detected automatically${NC}\n"
    # 自动检测模式：基于你测试有效的命令重构（分两步：先找所有I210 BDF，再检查驱动是否为VFIO）
    INPUT_BDF_LIST=()
    # 第一步：获取所有Intel I210网卡的BDF
    I210_ALL_BDF=($(lspci -nn | grep -i "Intel Corporation I210" | grep -oE '[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]'))
    # 第二步：遍历每个I210 BDF，检查驱动是否为vfio-pci
    for BDF in "${I210_ALL_BDF[@]}"; do
        DRIVER=$(lspci -k -s "$BDF" | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
        if [ "$DRIVER" = "$VFIO_DRIVER" ]; then
            INPUT_BDF_LIST+=("$BDF")
        fi
    done
    # 去重排序
    INPUT_BDF_LIST=($(echo "${INPUT_BDF_LIST[@]}" | tr ' ' '\n' | sort -u | tr '\n' ' '))
    
    if [ ${#INPUT_BDF_LIST[@]} -eq 0 ]; then
        log "${YELLOW}No I210 NICs bound to VFIO detected, no recovery needed${NC}"
        exit 0
    fi
    log "${YELLOW}Automatically detected I210 NIC BDFs to recover: ${INPUT_BDF_LIST[*]}${NC}"
else
    INPUT_BDF_LIST=("$@")
fi

# 加载IGB驱动（确保原生驱动可用）
if ! lsmod | grep -q $I210_DRIVER; then
    log "${YELLOW}Loading I210 native driver ($I210_DRIVER)...${NC}"
    modprobe $I210_DRIVER 2>/dev/null || {
        log "${RED}Failed to load IGB driver!${NC}"
        exit 1
    }
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
        log "${YELLOW}Skipped: Invalid BDF format ($BDF)${NC}"
        continue
    fi
    # 2. 存在性校验（系统中是否有该BDF）
    if ! lspci -s "$BDF" >/dev/null 2>&1; then
        log "${YELLOW}Skipped: BDF does not exist ($BDF)${NC}"
        continue
    fi
    # 3. 类型校验（是否为Intel I210网卡）
    if ! lspci -nn -s "$BDF" | grep -qi "Intel Corporation I210"; then
        log "${YELLOW}Skipped: Not an I210 NIC ($BDF)${NC}"
        continue
    fi
    # 所有校验通过，加入有效列表
    VALID_I210_BDF_LIST+=("$BDF")
done

# 无有效I210 BDF则退出
if [ ${#VALID_I210_BDF_LIST[@]} -eq 0 ]; then
    log "${YELLOW}No valid Intel I210 NIC BDFs need recovery, script exits normally${NC}"
    exit 0  # 退出码从1改为0
fi

# ====================== 核心：恢复I210网卡驱动 ======================
log "${GREEN}===== Start recovering I210 NIC drivers =====${NC}"
for BDF in "${VALID_I210_BDF_LIST[@]}"; do
    echo ""
    log "${GREEN}[NIC BDF: $BDF] Starting recovery...${NC}"
    BDF_FULL="0000:$BDF"
    NIC_NAME=""

    # 步骤1：获取当前驱动状态
    CURRENT_DRIVER=$(lspci -k -s $BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
    log "Current driver: $CURRENT_DRIVER"

    # 若已是IGB驱动，跳过
    if [ "$CURRENT_DRIVER" = "$I210_DRIVER" ]; then
        log "${YELLOW}No recovery needed: This NIC is already bound to IGB driver${NC}"
        continue
    fi

    # 步骤2：解析网卡名（用于恢复网络连接）
    if [ -d "/sys/bus/pci/devices/$BDF_FULL/net/" ]; then
        NIC_NAME=$(ls /sys/bus/pci/devices/$BDF_FULL/net/ | head -1)
    fi
    if [ -n "$NIC_NAME" ]; then
        log "NIC name: $NIC_NAME"
    fi

    # 步骤3：解绑VFIO-PCI驱动（直通驱动）
    log "Unbinding VFIO-PCI driver..."
    # 先移除VFIO的设备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 [ -n "$PCI_ID" ]; then
        VENDOR_ID=$(echo $PCI_ID | cut -d: -f1)
        DEVICE_ID=$(echo $PCI_ID | cut -d: -f2)
        echo $VENDOR_ID $DEVICE_ID > /sys/bus/pci/drivers/$VFIO_DRIVER/remove_id 2>/dev/null || true
    fi
    # 解绑驱动
    echo $BDF_FULL > /sys/bus/pci/devices/$BDF_FULL/driver/unbind 2>/dev/null || true
    sleep 1  # 等待解绑生效

    # 步骤4：绑定IGB原生驱动
    log "Binding IGB native driver..."
    echo $BDF_FULL > /sys/bus/pci/drivers/$I210_DRIVER/bind 2>/dev/null || {
        log "${RED}Failed to bind IGB driver! Trying to rescan PCI bus manually...${NC}"
        echo 1 > /sys/bus/pci/rescan 2>/dev/null
        sleep 1
        echo $BDF_FULL > /sys/bus/pci/drivers/$I210_DRIVER/bind 2>/dev/null || {
            log "${RED}Failed to bind after manual rescan, please check driver status${NC}"
            continue
        }
    }

    # 步骤5：验证驱动恢复结果
    sleep 1  # 等待驱动绑定生效
    FINAL_DRIVER=$(lspci -k -s $BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
    if [ "$FINAL_DRIVER" = "$I210_DRIVER" ]; then
        log "${GREEN}Success: Driver recovered to IGB${NC}"
    else
        log "${RED}Failed: Driver recovery failed (Current driver: $FINAL_DRIVER)${NC}"
        continue
    fi

    # 步骤6：恢复网络连接
    if [ -n "$NIC_NAME" ]; then
        log "Restoring $NIC_NAME network connection..."
        systemctl start NetworkManager 2>/dev/null || true
        nmcli connection up $NIC_NAME 2>/dev/null || log "${YELLOW}Tip: No available network connection for $NIC_NAME${NC}"
    fi
done

# ====================== 最终输出 ======================
echo ""
log "${GREEN}===== Recovery completed =====${NC}"
log "${GREEN}Processed BDF list: ${VALID_I210_BDF_LIST[*]}${NC}"
log "${GREEN}Recovery log file: $LOG_FILE${NC}"
log "${YELLOW}Tip: Execute lspci -k -s <BDF> to verify driver status${NC}"

# 列出最终驱动状态
echo -e "\n${GREEN}=== Final driver status ===${NC}"
for BDF in "${VALID_I210_BDF_LIST[@]}"; do
    DRIVER=$(lspci -k -s $BDF | grep "Kernel driver in use" | awk -F': ' '{print $2}' | xargs)
    echo -e "BDF $BDF: $DRIVER"
done
echo ""
