#!/bin/bash

# 配置项
VM_XML_SOURCE_PATH="/etc/vmdeploy/Windows_11.xml"  # VM main XML path
USB_PASSTHROUGH_XML="/tmp/usb_passthrough.xml"  # Temporary XML file
LOG_FILE="/tmp/usb_direct.log"

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

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

# 预处理校验 - 检查USB设备（Bus+Device）是否已直通（基于实际XML内容，非注释）
check_usb_device_existed() {
    local busnum=$1
    local devnum=$2
    # 匹配USB hostdev中source下的address bus和device（精准匹配实际设备地址）
    local existed=$(grep -A2 "<source>" "$VM_XML_SOURCE_PATH" | grep -E "<address bus='$busnum' device='$devnum'/>" | wc -l)
    if [ $existed -gt 0 ]; then
        return 0  # Exists, return 0
    else
        return 1  # Not exists, return 1
    fi
}

# 根据 VID:PID 解析 lsusb，返回真实的 Bus 和 Device 号
parse_usb_id_to_bus_dev() {
    local vid_pid=$1
    local vid=$(echo "$vid_pid" | cut -d: -f1)
    local pid=$(echo "$vid_pid" | cut -d: -f2)
    
    # 核心修改：添加 -i 参数开启大小写不敏感匹配，兼容16进制字母（a-f/A-F）
    local lsusb_result=$(lsusb | grep -iE "ID $vid:$pid" | head -n1)
    if [ -z "$lsusb_result" ]; then
        echo -e "${RED}Error: USB device with VID:PID = $vid:$pid not found${NC}"
        log "USB device with VID:PID = $vid:$pid not found, skip directly"
        return 1  # Return non-0 for failure
    fi
    
    # 提取 Bus/Device（严格格式解析）
    local busnum=$(echo "$lsusb_result" | awk '{print $2}')
    local devnum=$(echo "$lsusb_result" | awk '{print $4}' | sed 's/://')
    # 去掉前导0，避免被当作八进制
    busnum=$((10#$busnum))
    devnum=$((10#$devnum))
    if [ -z "$busnum" ] || [ -z "$devnum" ]; then
        echo -e "${RED}Error: Failed to parse Bus/Device number (VID:PID=$vid_pid)${NC}"
        log "Failed to parse Bus/Device, skip device $vid_pid"
        return 1
    fi
    
    echo "$busnum $devnum"
    return 0
}

# 查找空闲的USB端口
find_free_usb_port() {
    # 限制usb bus号为2，增加正则选项，避免端口号冲突
    local used_ports=($(grep -A1 "bus='2'" "$VM_XML_SOURCE_PATH" | grep -o "port='[0-9]\+'" | cut -d= -f2 | tr -d "'" | sort -n))
    local free_port=1
    while [[ " ${used_ports[@]} " =~ " $free_port " ]]; do
        free_port=$((free_port + 1))
    done
    echo "$free_port"
}

# 生成USB直通XML
generate_usb_xml() {
    local busnum=$1
    local devnum=$2
    local free_port=$3
    local vid_pid=$4
    cat > "$USB_PASSTHROUGH_XML" << EOF
<!-- USB passthrough config: VID:PID=$vid_pid → Bus=$busnum, Device=$devnum -->
<hostdev mode='subsystem' type='usb' managed='yes'>
  <source>
    <address bus='$busnum' device='$devnum'/>
  </source>
  <address type='usb' bus='2' port='$free_port'/>
</hostdev>
EOF
    log "Generated USB passthrough XML: VID:PID=$vid_pid → Bus=$busnum, Device=$devnum"
}

# 移除原有USB配置
remove_existing_usb_config() {
    local vid_pid=$1
    sed -i "/<!-- USB passthrough config: VID:PID=$vid_pid/,/<\/hostdev>/d" "$VM_XML_SOURCE_PATH"
    log "Removed existing USB config: VID:PID=$vid_pid"
}

# 插入USB XML到memballoon前
insert_usb_xml() {
    local vid_pid=$1
    local busnum=$2
    local devnum=$3
    local free_port=$(find_free_usb_port)
    
    # 备份原XML
    local backup_xml="${VM_XML_SOURCE_PATH}.bak_$(date +%Y%m%d_%H%M%S)"
    cp "$VM_XML_SOURCE_PATH" "$backup_xml"
    log "Backed up VM XML to: $backup_xml"
    echo -e "${YELLOW}Backed up original VM XML to: $backup_xml${NC}"

    # 移除原有配置
    remove_existing_usb_config "$vid_pid"

    # 生成XML
    generate_usb_xml "$busnum" "$devnum" "$free_port" "$vid_pid"

    # 插入到memballoon前（过滤空行+颜色字符）
    grep -v '^$' "$USB_PASSTHROUGH_XML" | sed -r 's/\x1B\[[0-9;]*m//g' > /tmp/tmp_usb_xml
    sed -i "/<memballoon model='virtio'>/e cat /tmp/tmp_usb_xml" "$VM_XML_SOURCE_PATH"

    # 清理临时文件
    rm -f /tmp/tmp_usb_xml

    # 验证插入结果
    if grep -q "USB passthrough config: VID:PID=$vid_pid" "$VM_XML_SOURCE_PATH"; then
        echo -e "${GREEN}Successfully inserted USB passthrough config: VID:PID=$vid_pid → Bus=$busnum, Device=$devnum${NC}"
        log "Successfully inserted USB passthrough config: VID:PID=$vid_pid"
        return 0
    else
        echo -e "${RED}Failed to insert USB passthrough config, restoring backup${NC}"
        cp "$backup_xml" "$VM_XML_SOURCE_PATH"
        log "Insert failed, restored backup: $backup_xml"
        return 1
    fi
}

# 主逻辑（设备不存在则完全跳过，无任何XML操作）
main() {
    if [ $# -eq 0 ]; then
        echo -e "${RED}Error: Please pass USB device VID:PID (e.g., 3535:6300)${NC}"
        log "No USB VID:PID parameters passed"
        exit 1
    fi

    if [ ! -f "$VM_XML_SOURCE_PATH" ]; then
        echo -e "${RED}Error: VM XML file $VM_XML_SOURCE_PATH does not exist${NC}"
        log "VM XML file does not exist"
        exit 1
    fi

    # 处理每个USB VID:PID参数
    for vid_pid in "$@"; do
        echo -e "${YELLOW}Processing USB device: VID:PID = $vid_pid${NC}"
        log "Start processing USB device: VID:PID = $vid_pid"
        
        # 解析Bus/Device，失败则直接跳过当前设备
        bus_dev=$(parse_usb_id_to_bus_dev "$vid_pid")
        if [ $? -ne 0 ] || [ -z "$bus_dev" ]; then
            echo -e "${YELLOW}Skipping invalid USB device: $vid_pid${NC}"
            log "Skipping invalid USB device: $vid_pid"
            # 完全跳过，不执行任何XML相关操作
            continue
        fi
        
        # 拆分Bus/Device
        busnum=$(echo "$bus_dev" | awk '{print $1}')
        devnum=$(echo "$bus_dev" | awk '{print $2}')
        
        # ========== 预处理校验 ==========
        if check_usb_device_existed "$busnum" "$devnum"; then
            echo -e "${YELLOW}This USB device (Bus=$busnum, Device=$devnum) is already passed through, skip duplicate config: $vid_pid${NC}"
            log "Device already exists (Bus=$busnum, Device=$devnum), skip duplicate config: $vid_pid"
            continue
        fi
        
        echo -e "${YELLOW}Parsed real device info: Bus=$busnum, Device=$devnum${NC}"
        log "Parsed real device info: VID:PID=$vid_pid → Bus=$busnum, Device=$devnum"
        
        # 插入XML（仅处理有效设备）
        if ! insert_usb_xml "$vid_pid" "$busnum" "$devnum"; then
            echo -e "${RED}Failed to process USB device $vid_pid${NC}"
            log "Failed to process USB device $vid_pid"
            continue
        fi
    done

    # 清理临时文件
    rm -f "$USB_PASSTHROUGH_XML"
    log "All USB devices processed (invalid devices skipped)"
    echo -e "${GREEN}=====================${NC}"
    echo -e "${GREEN}All valid USB passthrough configs are in effect${NC}"
    echo -e "${GREEN}=====================${NC}"
    exit 0
}

# 执行主逻辑
main "$@"
