#!/usr/bin/bash

# Parse command line arguments
OUTPUT_FORMAT="yaml"
while [[ $# -gt 0 ]]; do
    case $1 in
        --json)
            OUTPUT_FORMAT="json"
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [--json]"
            echo "  --json    Output in JSON format (default is YAML-like format)"
            exit 0
            ;;
        *)
            echo "Unknown option: $1" >&2
            echo "Use --help for usage information" >&2
            exit 1
            ;;
    esac
done

BOOTIMG_LOCATION=/lib/firmware/mellanox/boot/default.bfb

get_version()
{
	if [ -e /etc/debian_version ]; then
		dpkg --list $1 | grep -w "$1" | awk '{print $2,$3}'
	else
		if (rpm -q --quiet $1); then
			rpm -q --queryformat="[%{NAME}-%{VERSION}-%{RELEASE}]" $1
		fi
	fi
}

if [ -e "$BOOTIMG_LOCATION" ]; then
	BUILD_ATF=$(strings $BOOTIMG_LOCATION | grep -m 1 "(\(release\|debug\))")
	BUILD_UEFI=$(strings -e l $BOOTIMG_LOCATION | grep "BlueField" |\
                    cut -d':' -f 2)
	BOOTIMAGE_VER=$(get_version mlxbf-bootimages)
	BUILD_BSP=$(echo "$BOOTIMAGE_VER" | sed -e 's/mlxbf-bootimages-//')

	if [ -x "$(command -v bfver)" ]; then
		BFVER_VAL=$(bfver)
	fi
	if [ -n "$BFVER_VAL" ]; then
		BUILD_ATF=$(echo "$BFVER_VAL" | grep ATF | awk '{ print $NF }' | head -n 1)
                BUILD_UEFI=$(echo "$BFVER_VAL" | grep UEFI | awk '{ print $NF }' | head -n 1)
                BUILD_BSP=$(echo "$BFVER_VAL" | grep BSP | awk '{ print $NF }' | head -n 1)
        fi

fi


OFED=`ofed_info -s 2> /dev/null | tr -d ':' | cut -d '-' -f2-`
if [ ! -n "$OFED" ]; then
	OFED="in-box"
fi

DEV_TYPE=$(mlxfwmanager | grep "Device Type:" | awk '{print $3}')

if [ "$DEV_TYPE" == "BlueField" ]; then
	NIC_FW=`/opt/mellanox/mlnx-fw-updater/firmware/mlxfwmanager_sriov_dis_aarch64_41682 --list 2> /dev/null | head -3 | tail -1 | awk '{print $4}'`
elif [ "$DEV_TYPE" == "BlueField2" ]; then 
    NIC_FW=`/opt/mellanox/mlnx-fw-updater/firmware/mlxfwmanager_sriov_dis_aarch64_41686 --list 2> /dev/null | head -3 | tail -1 | awk '{print $4}'`
elif [ "$DEV_TYPE" == "BlueField3" ]; then
    NIC_FW=`/opt/mellanox/mlnx-fw-updater/firmware/mlxfwmanager_sriov_dis_aarch64_41692 --list 2> /dev/null | head -3 | tail -1 | awk '{print $4}'`
fi

MLX_REGEX=$(get_version mlx-regex 2> /dev/null)

get_version_and_release()
{
	if [ -e /etc/debian_version ]; then
		dpkg --list $1 | grep -w "$1" | awk '{print $3}'
	else
		if (rpm -q --quiet $1); then
			rpm -q --queryformat="[%{VERSION}-%{RELEASE}]" $1
		fi
	fi
}

print_ofed()
{
	if [ -e /etc/debian_version ]; then
		ofed_info | sed -n '/^-------------------$/ { :a; n; p; ba; }' | awk 'NF {print "- " $2, $3}'
	else
		ofed_info | sed -n '/^-------------------$/ { :a; n; p; ba; }' | xargs rpm -q --queryformat="[- %{NAME} %{VERSION}-%{RELEASE}]\n"
	fi
}

get_bmc_fw()
{
	if [ -e /lib/firmware/mellanox/bmc/bf2-bmc-fw.version ]; then
		cat /lib/firmware/mellanox/bmc/bf2-bmc-fw.version
	elif [ -e /lib/firmware/mellanox/bmc/bf3-bmc-fw.version ]; then
		cat /lib/firmware/mellanox/bmc/bf3-bmc-fw.version
	fi
}

get_cec_fw()
{
	if [ -e /lib/firmware/mellanox/cec/bf2-cec-fw.version ]; then
		cat /lib/firmware/mellanox/cec/bf2-cec-fw.version
	elif [ -e /lib/firmware/mellanox/cec/bf3-cec-fw.version ]; then
		cat /lib/firmware/mellanox/cec/bf3-cec-fw.version
	fi
}

# JSON utility functions
json_escape() {
    echo "$1" | sed 's/\\/\\\\/g; s/"/\\"/g; s/	/\\t/g' | tr -d '\n'
}

json_add_item() {
    local key="$1"
    local value="$2"
    local is_last="$3"
    
    if [ -n "$value" ] && [ "$value" != " " ]; then
        if [ "$is_last" = "true" ]; then
            echo "    \"$key\": \"$(json_escape "$value")\""
        else
            echo "    \"$key\": \"$(json_escape "$value")\","
        fi
    elif [ "$is_last" = "true" ]; then
        echo "    \"$key\": null"
    else
        echo "    \"$key\": null,"
    fi
}

json_add_array() {
    local key="$1"
    local items="$2"
    local is_last="$3"
    
    if [ "$is_last" = "true" ]; then
        echo "    \"$key\": ["
    else
        echo "    \"$key\": ["
    fi
    
    if [ -n "$items" ]; then
        echo "$items" | while IFS= read -r line; do
            if [ -n "$line" ]; then
                # Remove leading "- " from YAML format
                clean_line=$(echo "$line" | sed 's/^- //')
                echo "      \"$(json_escape "$clean_line")\","
            fi
        done | sed '$s/,$//'  # Remove trailing comma from last item
    fi
    
    if [ "$is_last" = "true" ]; then
        echo "    ]"
    else
        echo "    ],"
    fi
}

# Collect all version information
DPDK_VERSION=$(/opt/mellanox/dpdk/bin/dpdk-testpmd -v 2>&1 | grep "RTE Version:" | cut -d ':' -f 3)
KERNEL_VERSION=$(uname -r)
MFT_VERSION=$(get_version_and_release mft)
MSTFLINT_VERSION=$(get_version_and_release mstflint)
BMC_FW=$(get_bmc_fw)
CEC_FW=$(get_cec_fw)

# Collect storage packages
if [ "$DEV_TYPE" != "BlueField2" ]; then
    STORAGE_PACKAGES=$(
        version=$(get_version virtio-net-controller)
        [ -n "$version" ] && echo "- $version"
    )
else
    
    STORAGE_PACKAGES=$(
        version=$(get_version mlnx-libsnap)
        [ -n "$version" ] && echo "- $version"
        version=$(get_version mlnx-snap)
        [ -n "$version" ] && echo "- $version"
        version=$(get_version spdk)
        [ -n "$version" ] && echo "- $version"
        version=$(get_version virtio-net-controller)
        [ -n "$version" ] && echo "- $version"
    )
fi

# Collect DOCA packages
if [ -e /etc/debian_version ]; then
    DOCA_PACKAGES=$(
        for doca in $(dpkg --list | grep -E 'doca|rxp|dpa-compiler' | awk '{print $2}' | sort -n); do
            version=$(get_version $doca)
            [ -n "$version" ] && echo "- $version"
        done
        version=$(get_version collectx-clxapi)
        [ -n "$version" ] && echo "- collectx-clxapi: $version"
    )
else
    DOCA_PACKAGES=$(
        for doca in $(rpm -qa | grep -E 'doca|rxp|dpa-compiler' | sort -n); do
            version=$(get_version $doca)
            [ -n "$version" ] && echo "- $version"
        done
        version=$(get_version collectx-clxapi)
        [ -n "$version" ] && echo "- collectx-clxapi: $version"
    )
fi

# Collect FlexIO packages
if [ -e /etc/debian_version ]; then
    FLEXIO_PACKAGES=$(
        for flexio in $(dpkg --list | grep -E 'dpacc|flexio|dpaeumgmt|dpa-gdbserver|dpa-stats' | awk '{print $2}' | sort -n); do
            version=$(get_version $flexio)
            [ -n "$version" ] && echo "- $version"
        done
    )
else
    FLEXIO_PACKAGES=$(
        for flexio in $(rpm -qa | grep -E 'dpacc|flexio|dpaeumgmt|dpa-gdbserver|dpa-stats' | sort -n); do
            version=$(get_version $flexio)
            [ -n "$version" ] && echo "- $version"
        done
    )
fi

# Collect SoC Platform packages
if [ -e /etc/debian_version ]; then
    SOC_PACKAGES=$(
        for package in $(dpkg --list | grep -E 'mlxbf-gige|sdhci-of-dwcmshc|tmfifo|tmfifo|gpio-mlxbf|pinctrl-mlxbf3|i2c-mlxbf|mlx-OpenIPMI|ipmb-dev-int|mlxbf-livefish|ipmb-host|mlxbf-p|pwr-mlxbf|mlx-trio|mmc-utils' | awk '{print $2}' | sort -n); do
            version=$(get_version $package)
            [ -n "$version" ] && echo "- $version"
        done
    )
else
    SOC_PACKAGES=$(
        for package in $(rpm -qa | grep -E 'mlxbf-gige|sdhci-of-dwcmshc|tmfifo|tmfifo|gpio-mlxbf|pinctrl-mlxbf3|i2c-mlxbf|mlx-OpenIPMI|ipmb-dev-int|mlxbf-livefish|ipmb-host|mlxbf-p|pwr-mlxbf|mlx-trio|mmc-utils' | sort -n); do
            version=$(get_version $package)
            [ -n "$version" ] && echo "- $version"
        done
    )
fi

# Collect OFED packages if not in-box
if [ "$OFED" != "in-box" ]; then
    OFED_PACKAGES=$(print_ofed)
fi

# Output in the requested format
if [ "$OUTPUT_FORMAT" = "json" ]; then
    echo "{"
    echo "  \"firmware\": {"
    json_add_item "ATF" "$BUILD_ATF" "false"
    json_add_item "UEFI" "$BUILD_UEFI" "false"
    json_add_item "BSP" "$BUILD_BSP" "false"
    json_add_item "NIC_Firmware" "$NIC_FW" "false"
    json_add_item "BMC_Firmware" "$BMC_FW" "false"
    json_add_item "CEC_Firmware" "$CEC_FW" "true"
    echo "  },"
    echo "  \"drivers\": {"
    json_add_item "mlnx_dpdk" "$DPDK_VERSION" "false"
    json_add_item "Kernel" "$KERNEL_VERSION" "true"
    echo "  },"
    echo "  \"tools\": {"
    json_add_item "MFT" "$MFT_VERSION" "false"
    if [ -n "$MLX_REGEX" ]; then
        json_add_item "mstflint" "$MSTFLINT_VERSION" "false"
        json_add_item "mlx_regex" "$MLX_REGEX" "true"
    else
        json_add_item "mstflint" "$MSTFLINT_VERSION" "true"
    fi
    echo "  },"
    json_add_array "storage" "$STORAGE_PACKAGES" "false"
    json_add_array "doca" "$DOCA_PACKAGES" "false"
    json_add_array "flexio" "$FLEXIO_PACKAGES" "false"
    if [ "$OFED" != "in-box" ]; then
        json_add_array "soc_platform" "$SOC_PACKAGES" "false"
        json_add_array "ofed" "$OFED_PACKAGES" "true"
    else
        json_add_array "soc_platform" "$SOC_PACKAGES" "true"
    fi
    echo "}"
else
    # Original YAML-like format
    cat << EOF

Firmware:
- ATF: $BUILD_ATF
- UEFI: $BUILD_UEFI
- BSP: $BUILD_BSP
- NIC Firmware: $NIC_FW
- BMC Firmware: $BMC_FW
- CEC Firmware: $CEC_FW

Drivers:
- mlnx-dpdk:$DPDK_VERSION
- Kernel: $KERNEL_VERSION

Tools:
- MFT: $MFT_VERSION
- mstflint: $MSTFLINT_VERSION
EOF
    if [ -n "$MLX_REGEX" ]; then
        echo "- mlx-regex: $MLX_REGEX"
    fi

    cat << EOF

Storage:
$STORAGE_PACKAGES

DOCA:
$DOCA_PACKAGES

FlexIO:
$FLEXIO_PACKAGES

SoC Platform:
$SOC_PACKAGES
EOF

    if [ "$OFED" != "in-box" ]; then
        cat << EOF

OFED:
$OFED_PACKAGES

EOF
    fi
fi

