#!/bin/bash

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

# ===================== 核心可配置变量 =====================
TARGET_XML_PATH="/etc/vmdeploy/Windows_11.xml"
TARGET_CARD_KEYWORDS=("d500" "Processing accelerators")
FIXED_VM_BUS="0x05"

# 全局变量：存储校验通过的目标PCI地址
TARGET_PCI_ADDR=""

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1" >&2
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1" >&2
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

# 检查必要工具
check_dependencies() {
    local deps=("lspci" "grep" "awk" "sed")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            log_error "Required tool missing: $dep"
            exit 1
        fi
    done
    log_info "Dependency check passed"
}

# 检查XML文件有效性
check_xml_valid() {
    local xml_file="$1"
    
    if [ ! -f "$xml_file" ]; then
        log_error "XML file does not exist: $xml_file"
        return 1
    fi
    
    if ! grep -q "<domain" "$xml_file" || ! grep -q "</domain>" "$xml_file"; then
        log_error "Invalid XML file format: $xml_file"
        return 1
    fi
    
    log_info "XML file validation passed: $xml_file"
    return 0
}

# 校验BDF格式
validate_bdf_format() {
    local bdf="$1"
    if ! echo "$bdf" | grep -qE '^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$'; then
        log_error "Invalid BDF format ($bdf), correct format is xx:xx.x (e.g. 03:00.0)"
        return 1
    fi
    # 校验BDF是否存在于系统中
    if ! lspci -s "$bdf" >/dev/null 2>&1; then
        log_error "BDF $bdf does not exist in the system!"
        return 1
    fi
    return 0
}

# 将BDF（xx:xx.x）转换为完整PCI地址（0000:xx:xx.x）
bdf_to_full_pci() {
    local bdf="$1"
    echo "0000:$bdf"
}

# 检查BDF对应的设备是否是目标AI卡
is_target_ai_card() {
    local bdf="$1"
    local full_pci=$(bdf_to_full_pci "$bdf")
    
    # 获取设备描述
    local device_desc=$(lspci -D -nn -s "$full_pci" 2>/dev/null | sed 's/^[0-9a-f:.]* //')
    if [ -z "$device_desc" ]; then
        log_error "Failed to get device description for BDF $bdf"
        return 1
    fi
    
    # 校验是否匹配所有关键词
    local match_count=0
    for keyword in "${TARGET_CARD_KEYWORDS[@]}"; do
        if echo "$device_desc" | grep -qi "$keyword"; then
            ((match_count++))
        fi
    done
    
    if [ "$match_count" -eq "${#TARGET_CARD_KEYWORDS[@]}" ]; then
        return 0
    else
        log_info "Device corresponding to BDF $bdf is not the target HW D500 AI card (Device description: $device_desc)"
        log_info "Target card keywords: ${TARGET_CARD_KEYWORDS[*]}"
        return 1
    fi
}

# 解析PCI设备信息（兼容完整PCI地址和BDF）
parse_pci_info() {
    local pci_addr="$1"
    
    # 如果是BDF格式（无域号），先补全域号
    if echo "$pci_addr" | grep -qE '^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]$'; then
        pci_addr=$(bdf_to_full_pci "$pci_addr")
    fi
    
    # 提取域、总线、设备、功能号
    local domain=$(echo "$pci_addr" | cut -d: -f1)
    local bus=$(echo "$pci_addr" | cut -d: -f2)
    local slot=$(echo "$pci_addr" | cut -d: -f3 | cut -d. -f1)
    local function=$(echo "$pci_addr" | cut -d. -f2)
    
    # 去除前导零，但保持至少一位
    domain=$(echo "$domain" | sed 's/^0*//')
    domain=${domain:-0}
    
    bus=$(echo "$bus" | sed 's/^0*//')
    bus=${bus:-0}
    
    slot=$(echo "$slot" | sed 's/^0*//')
    slot=${slot:-0}
    
    function=$(echo "$function" | sed 's/^0*//')
    function=${function:-0}
    
    # 获取设备名称
    local device_name=$(lspci -s "$pci_addr" 2>/dev/null | sed 's/^[0-9a-f:.]* //' || echo "Unknown Device")
    
    echo "$domain:$bus:$slot:$function:$device_name"
}

# 检查设备是否已在XML中配置
is_device_in_xml() {
    local xml_file="$1"
    local domain="$2"
    local bus="$3"
    local slot="$4"
    local function="$5"
    
    local pattern="domain='0x${domain}' bus='0x${bus}' slot='0x${slot}' function='0x${function}'"
    
    if grep -q "$pattern" "$xml_file" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

# ========== 核心优化：仅查找0x05 Bus下已占用的Slot号 ==========
# 获取0x05 Bus下已使用的PCI slot
get_used_pci_slots() {
    local xml_file="$1"
    local target_bus="$FIXED_VM_BUS"  # 仅查AI卡专用Bus（0x05）
    
    # 核心优化：先过滤出bus='$target_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
}

# 获取0x05 Bus下下一个可用的PCI slot（从0x01开始顺序分配）
get_next_available_slot() {
    local xml_file="$1"
    local start_slot=1  # 从0x01（十进制1）开始分配
    
    local used_slots=($(get_used_pci_slots "$xml_file"))
    local next_slot=$start_slot
    
    # 遍历已用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
}

# 生成hostdev XML片段（固定Bus为0x05，使用分配的Slot号）
generate_hostdev_xml() {
    local domain="$1"
    local bus="$2"
    local slot="$3"
    local function="$4"
    local device_name="$5"
    local vm_slot="$6"
    
    cat << EOF
    <hostdev mode='subsystem' type='pci' managed='yes'>
      <source>
        <address domain='0x$domain' bus='0x$bus' slot='0x$slot' function='0x$function'/>
      </source>
      <address type='pci' domain='0x0000' bus='$FIXED_VM_BUS' slot='$vm_slot' function='0x0'/>
    </hostdev>
EOF
}

# 在memballoon标签前插入配置
insert_gpu_config() {
    local xml_file="$1"
    local gpu_xml="$2"
    
    local temp_file="${xml_file}.tmp"
    
    awk -v gpu_xml="$gpu_xml" '
    BEGIN { inserted = 0 }
    /<memballoon/ && !inserted {
        print gpu_xml
        inserted = 1
    }
    { print }
    ' "$xml_file" > "$temp_file"
    
    if ! grep -q "hostdev mode='subsystem' type='pci'" "$temp_file" 2>/dev/null; then
        awk -v gpu_xml="$gpu_xml" '
        /<\/devices>/ {
            print gpu_xml
            print $0
            inserted = 1
            next
        }
        { print }
        ' "$xml_file" > "$temp_file"
    fi
    
    if grep -q "hostdev mode='subsystem' type='pci'" "$temp_file" 2>/dev/null; then
        mv "$temp_file" "$xml_file"
        log_info "Configuration inserted into XML file"
        return 0
    else
        rm -f "$temp_file" 2>/dev/null
        log_error "Failed to insert configuration"
        return 1
    fi
}

# 适配：根据传入的BDF查找目标卡
detect_target_card_by_bdf() {
    local bdf="$1"
    log_info "Verifying if the device corresponding to BDF $bdf is the target HW D500 AI card"
    
    # 校验BDF格式和存在性
    if ! validate_bdf_format "$bdf"; then
        return 1
    fi
    
    # 检查是否是目标AI卡
    if ! is_target_ai_card "$bdf"; then
        log_error "Device corresponding to BDF $bdf is not the target HW D500 AI card, terminating configuration process"
        return 1
    fi
    
    # 转换为完整PCI地址并设置全局变量
    local full_pci=$(bdf_to_full_pci "$bdf")
    local device_desc=$(lspci -D -nn -s "$full_pci" | sed 's/^[0-9a-f:.]* //')
    
    log_info "Target AI card found:"
    log_info "  BDF address: $bdf"
    log_info "  Full PCI address: $full_pci"
    log_info "  Device description: $device_desc"
    
    TARGET_PCI_ADDR="$full_pci"
    return 0
}

# 主函数
main() {
    # 仅接收1个参数：BDF地址
    local target_bdf="$1"
    
    # 校验参数数量：必须且仅能传入1个BDF参数
    if [ $# -ne 1 ]; then
        log_error "Script only supports receiving 1 BDF parameter!"
        usage
        exit 1
    fi
    
    # 校验BDF是否为空
    if [ -z "$target_bdf" ]; then
        log_error "Please specify the BDF address of the AI card! Usage: $0 <BDF>"
        exit 1
    fi
    
    log_info "Starting HW D500 card passthrough configuration:"
    log_info "Target BDF: $target_bdf"
    log_info "Fixed XML path: $TARGET_XML_PATH"
    log_info "AI card dedicated Bus number: $FIXED_VM_BUS"
    
    # 基础检查
    check_dependencies
    if ! check_xml_valid "$TARGET_XML_PATH"; then
        exit 1
    fi
    
    # 查找目标卡
    if ! detect_target_card_by_bdf "$target_bdf"; then
        exit 1
    fi
    local target_pci_addr="$TARGET_PCI_ADDR"
    
    # 解析目标卡PCI信息
    local pci_info=$(parse_pci_info "$target_pci_addr")
    IFS=':' read -r domain bus slot function device_name <<< "$pci_info"
    
    # 检查卡是否已配置
    if is_device_in_xml "$TARGET_XML_PATH" "$domain" "$bus" "$slot" "$function"; then
        log_warn "Target card $target_pci_addr is already configured in XML, skipping"
        exit 0
    fi
    
    # 获取0x05 Bus下可用Slot号
    local next_slot=$(get_next_available_slot "$TARGET_XML_PATH")
    log_info "Assigned Slot number for target card (under 0x05 Bus): $next_slot"
    
    # 生成直通XML片段
    local device_xml=$(generate_hostdev_xml "$domain" "$bus" "$slot" "$function" "$device_name" "$next_slot")
    
    # 插入到XML中
    log_info "Inserting target card passthrough configuration..."
    if insert_gpu_config "$TARGET_XML_PATH" "$device_xml"; then
        log_info "Target card passthrough configuration completed!"
        log_info "Added configuration fragment:"
        echo "$device_xml"
    else
        log_error "Target card passthrough configuration failed!"
        exit 1
    fi
}

# 脚本用法（仅显示BDF单参数）
usage() {
    echo "Usage: $0 <BDF>"
    echo
    echo "Function: Add PCI passthrough configuration for HW D500 card to fixed XML file (only BDF format parameter is supported)"
    echo
    echo "Parameters:"
    echo "  BDF           BDF address of AI card (format: xx:xx.x, e.g. 02:00.0)"
    echo
    echo "Examples:"
    echo "  $0 02:00.0"
}

# 入参判断（仅支持-h/--help或1个BDF参数）
if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
    usage
    exit 0
fi

if [ "$EUID" -ne 0 ]; then
    log_warn "It is recommended to run this script with root privileges"
fi

# 执行主逻辑
main "$1"
