#!/usr/bin/bash

# 原始设备映射表
declare -A devices_id=(
    ["Z100"]="54b7"
    ["Z100L"]="55b7"
    ["K100"]="62b7"
    ["K100-AI"]="6210"
    ["K100-AI-ECO"]="6211"
    ["BW"]="6320"
    ["BW1100"]="6430"
)

# 构建反向映射表（设备ID → 设备名称）
declare -A devices
for name in "${!devices_id[@]}"; do
    id="${devices_id[$name]}"
    devices["$id"]="$name" 
done

get_dcu() {
    # 检查lspci命令是否存在
    if ! command -v lspci &> /dev/null; then
        echo "错误: lspci 命令未找到，请先安装 pciutils 包" >&2
        return 1
    fi

    # 获取设备ID列表
    mapfile -t dcu_list < <(lspci -nn | grep -i -E "display|co-processor" | awk -F'[][]' '{print $4}' | awk -F ":" '{print $2}')
    
    local index=0
    local dcu_num=0
    local total=${#dcu_list[@]}
    local detected_devices=()
    
    # 处理没有设备的情况
    if [ $total -eq 0 ]; then
        echo "未检测到任何DCU设备"
        DEVICE_NAME=""
        DEVICE_ID=""
        return 0
    fi

    echo "=== 检测到的DCU设备 ==="
    while [ $index -lt $total ]; do
        current_id="${dcu_list[$index]}"

        if [ -n "${devices[$current_id]}" ]; then
            echo "DCU #$dcu_num: ID ${current_id} → ${devices[$current_id]}"
            detected_devices+=("${devices[$current_id]}")
            ((dcu_num++))
        else
            echo "DCU #$dcu_num: 未知设备ID: $current_id" >&2
            detected_devices+=("unknown")
        fi

        ((index++))
    done
    
    # 统计信息
    echo "=========================="
    if [ $dcu_num -eq 0 ]; then
        echo "总计: 0张DCU设备"
    else
        # 使用关联数组统计每种设备的数量
        local -A device_count
        for dev in "${detected_devices[@]}"; do
            ((device_count[$dev]++))
        done
        
        # 输出统计信息
        echo "总计: $dcu_num张DCU设备"
        for dev_name in "${!device_count[@]}"; do
            if [ "$dev_name" != "unknown" ]; then
                echo "  - ${dev_name}: ${device_count[$dev_name]}张"
            fi
        done
        if [ ${device_count["unknown"]:-0} -gt 0 ]; then
            echo "  - 未知设备: ${device_count["unknown"]}张"
        fi
    fi
    
    # 设置全局变量（取第一个有效设备）
    if [ $dcu_num -gt 0 ]; then
        DEVICE_NAME="${detected_devices[0]}"
        DEVICE_ID="${dcu_list[0]}"
    else
        DEVICE_NAME=""
        DEVICE_ID=""
    fi
}

get_dcu