#!/bin/bash

# 提取 Port GUID 并格式化为 SRIOV-VF 配置
generate_sriov_vf_config() {
    GUIDS=($(ibstat | grep "Port GUID" | awk '{print $3}'))
    for i in "${!GUIDS[@]}"; do
        echo "SRIOV-VF=${GUIDS[$i]}"
    done
}

# 更新 Docker 配置文件
update_docker_config() {
    local GUIDS=("$@")
    local DAEMON_JSON="/etc/docker/daemon.json"
    local TMP_JSON="/tmp/daemon.json"

    if [ ! -f "$DAEMON_JSON" ]; then
        echo "$DAEMON_JSON 文件不存在"
        exit 1
    fi

    local NODE_GENERIC_RESOURCES=$(jq -c '.["node-generic-resources"]' "$DAEMON_JSON")

    if [ "$NODE_GENERIC_RESOURCES" == "null" ]; then
        NODE_GENERIC_RESOURCES="[]"
    fi

    for GUID in "${GUIDS[@]}"; do
        if [[ ! $NODE_GENERIC_RESOURCES == *"$GUID"* ]]; then
            NODE_GENERIC_RESOURCES=$(echo "$NODE_GENERIC_RESOURCES" | jq --arg vf "$GUID" '. += [$vf]')
        fi
    done

    jq '.["node-generic-resources"] = '"$NODE_GENERIC_RESOURCES" "$DAEMON_JSON" > "$TMP_JSON"
    mv "$TMP_JSON" "$DAEMON_JSON"
}

# 主函数
main() {
    if [[ $EUID -ne 0 ]]; then
        echo "此脚本必须以 root 用户运行"
        exit 1
    fi

    GUIDS=($(generate_sriov_vf_config))
    if [ ${#GUIDS[@]} -eq 0 ]; then
        echo "未找到 SR-IOV VF 设备"
        exit 1
    fi

    update_docker_config "${GUIDS[@]}"
    echo "成功更新 $DAEMON_JSON 文件"
    systemctl restart docker
}

main "$@"