#!/usr/bin/bash
#
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: LicenseRef-NVIDIA-Proprietary
#
# dcb/QoS tool wrapper — real hardware version
# Prefers dcb (iproute2), falls back to mlnx_qos (OFED) for trust-mode
source "$(dirname "$0")/_lib.sh"
trace "$@"

OP="$1"; shift
TARGET="$1"; shift
BDF=$(bdf_from_target "$TARGET")
IFACE=$(iface_from_bdf "$BDF")

# YANG-enum int -> tool-string decode table.
#
# libdms converts YANG enum string values to their YANG-declared numeric
# ordinal before passing to the backend. dcb's `app trust set` accepts the
# YANG-name strings (pcp/dscp), not the ints. Without this decode the SET
# silently no-ops — until tool error-reporting is honest, which it now is.
#
# YANG vocabulary diverges slightly from dcb's:
#   YANG `port` (value 0)     -> dcb has no equivalent; map to `pcp` (closest
#                                operational match, preserves prior behavior)
#   YANG `pcp` (value 1)      -> dcb `pcp`
#   YANG `dscp` (value 2)     -> dcb `dscp`
#   YANG `dscp-pcp` (value 3) -> dcb has no equivalent; map to `dscp` (DSCP
#                                with PCP fallback collapses to DSCP-trust)
#
# The 0->pcp and 3->dscp mappings are the documented divergences. All others
# follow the discipline (YANG name = tool string).
declare -A YANG_ENUM=(
    ["trust-mode:0"]="pcp"
    ["trust-mode:1"]="pcp"
    ["trust-mode:2"]="dscp"
    ["trust-mode:3"]="dscp"
    ["trust-mode:port"]="pcp"
    ["trust-mode:pcp"]="pcp"
    ["trust-mode:dscp"]="dscp"
    ["trust-mode:dscp-pcp"]="dscp"
)

if [[ -z "$IFACE" ]]; then
    echo "Error: no interface for $BDF" >&2
    exit 1
fi

trust_get() {
    local v
    v=$(dcb app trust show dev "$IFACE" 2>/dev/null | grep -oP 'trust \K\S+')
    if [[ -z "$v" ]] && command -v mlnx_qos >/dev/null 2>&1; then
        v=$(mlnx_qos -i "$IFACE" 2>/dev/null | grep -oP 'Priority trust state:\s*\K\S+')
    fi
    echo "${v:-pcp}"
}

trust_set() {
    local val="$1"
    # Decode YANG ordinal/name -> dcb-accepted string via YANG_ENUM table
    # (defined at top of file). Falls through to $val if no entry matches.
    val="${YANG_ENUM[trust-mode:$val]:-$val}"
    local err
    if err=$(dcb app trust set dev "$IFACE" "$val" 2>&1); then
        return 0
    fi
    if command -v mlnx_qos >/dev/null 2>&1; then
        local err2
        if err2=$(mlnx_qos -i "$IFACE" --trust="$val" 2>&1); then
            return 0
        fi
        echo "Error: trust-mode SET $val on $IFACE failed (dcb: $err; mlnx_qos: $err2)" >&2
        return 1
    fi
    echo "Error: trust-mode SET $val on $IFACE failed (dcb: $err; mlnx_qos: not installed)" >&2
    return 1
}

pfc_get() {
    local RAW BITMAP=""
    RAW=$(dcb pfc show dev "$IFACE" 2>/dev/null | grep -oP 'prio-pfc\s+\K.*')
    if [[ -n "$RAW" ]]; then
        for i in 0 1 2 3 4 5 6 7; do
            local VAL=$(echo "$RAW" | grep -oP "${i}:\K\S+")
            [[ "$VAL" == "on" ]] && BITMAP="${BITMAP}1," || BITMAP="${BITMAP}0,"
        done
        BITMAP="${BITMAP%,}"
    elif command -v mlnx_qos >/dev/null 2>&1; then
        local LINE=$(mlnx_qos -i "$IFACE" 2>/dev/null | grep -A1 'enabled' | tail -1)
        BITMAP=$(echo "$LINE" | tr -s ' \t' ',' | sed 's/^,//;s/,$//')
    else
        BITMAP="0,0,0,0,0,0,0,0"
    fi
    # Convert bitmap to priority list: "0,0,0,1,0,0,0,0" → "3"
    local priorities=""
    IFS=',' read -ra bits <<< "$BITMAP"
    for i in "${!bits[@]}"; do
        [[ "${bits[$i]}" == "1" ]] && priorities="${priorities}${priorities:+,}${i}"
    done
    echo "${priorities:-}"
}

pfc_set() {
    local val="$1"
    local dcb_args=""
    for i in 0 1 2 3 4 5 6 7; do
        if echo ",$val," | grep -q ",${i},"; then
            dcb_args="${dcb_args} ${i}:on"
        else
            dcb_args="${dcb_args} ${i}:off"
        fi
    done
    local err
    if err=$(dcb pfc set dev "$IFACE" prio-pfc $dcb_args 2>&1); then
        return 0
    fi
    if command -v mlnx_qos >/dev/null 2>&1; then
        local bitmap=""
        for i in 0 1 2 3 4 5 6 7; do
            echo ",$val," | grep -q ",${i}," && bitmap="${bitmap}1," || bitmap="${bitmap}0,"
        done
        local err2
        if err2=$(mlnx_qos -i "$IFACE" --pfc="${bitmap%,}" 2>&1); then
            return 0
        fi
        echo "Error: pfc SET $val on $IFACE failed (dcb: $err; mlnx_qos: $err2)" >&2
        return 1
    fi
    echo "Error: pfc SET $val on $IFACE failed (dcb: $err; mlnx_qos: not installed)" >&2
    return 1
}

prio_tc_get() {
    # dcb ets show dev <IFACE> prio-tc returns "prio-tc 0:0 1:0 2:0 ..."
    # Return as comma-separated 8-tuple matching the leaf-list shape.
    local raw
    raw=$(dcb ets show dev "$IFACE" prio-tc 2>/dev/null | grep -oP 'prio-tc\s+\K.*')
    if [[ -z "$raw" ]]; then
        echo "0,1,2,3,4,5,6,7"  # kernel identity default
        return
    fi
    local out=""
    for i in 0 1 2 3 4 5 6 7; do
        local v=$(echo "$raw" | tr ' ' '\n' | grep -oP "^${i}:\K\S+" | head -1)
        out="${out}${out:+,}${v:-0}"
    done
    echo "$out"
}

prio_tc_set() {
    local val="$1"
    IFS=',' read -ra arr <<< "$val"
    local args=""
    local i=0
    for v in "${arr[@]}"; do
        args="${args} ${i}:${v}"
        i=$(( i + 1 ))
    done
    local err
    # Some iproute2/dcb versions validate the existing tc-bw sum even when
    # changing only prio-tc. Keep ETS bandwidth at the zero/default map used
    # by the driver so prio-tc updates do not fail on stale 100% placeholders.
    if ! err=$(dcb ets set dev "$IFACE" prio-tc $args \
        tc-bw 0:0 1:0 2:0 3:0 4:0 5:0 6:0 7:0 2>&1); then
        echo "Error: dcb ets set prio-tc $val on $IFACE failed: $err" >&2
        return 1
    fi
    return 0
}

# Resolve PCI BDF -> infiniband device for ToS unmap (different sysfs root)
rdma_from_bdf() {
    local bdf="$1"
    local ibdir="/sys/bus/pci/devices/${bdf}/infiniband"
    if [[ -d "$ibdir" ]]; then
        ls "$ibdir" 2>/dev/null | head -1
    fi
}

# tos-unmap: write -1 to /sys/class/infiniband/<rdma>/tc/<tos>/traffic_class
# Input: comma-separated ToS code-points (e.g. "96" or "96,40").
# Works for PF and VF BDFs alike — sysfs structure is the same.
tos_unmap_set() {
    local val="$1"
    local rdma=$(rdma_from_bdf "$BDF")
    if [[ -z "$rdma" ]]; then
        echo "Error: no infiniband dev for BDF $BDF" >&2
        return 1
    fi
    IFS=',' read -ra arr <<< "$val"
    local rc=0
    for tos in "${arr[@]}"; do
        local path="/sys/class/infiniband/${rdma}/tc/${tos}/traffic_class"
        if [[ ! -w "$path" ]]; then
            echo "Error: tos-unmap path not writable: $path" >&2
            rc=1
            continue
        fi
        if ! echo -1 > "$path" 2>/dev/null; then
            echo "Error: tos-unmap write -1 to $path failed" >&2
            rc=1
        fi
    done
    return $rc
}

tos_unmap_get() {
    local rdma=$(rdma_from_bdf "$BDF")
    if [[ -z "$rdma" ]]; then
        echo ""
        return
    fi
    # Walk /sys/class/infiniband/<rdma>/tc/*/traffic_class; emit ToS values
    # whose mapping is currently -1 (unmapped).
    local out=""
    for tcfile in /sys/class/infiniband/${rdma}/tc/*/traffic_class; do
        [[ -f "$tcfile" ]] || continue
        local tos=$(echo "$tcfile" | grep -oE '/tc/[0-9]+/' | tr -d '/' | sed 's/tc//')
        local v=$(cat "$tcfile" 2>/dev/null)
        if [[ "$v" == "-1" ]]; then
            out="${out}${out:+,}${tos}"
        fi
    done
    echo "$out"
}

if [[ "$OP" == "GET" ]]; then
    for PARAM in "$@"; do
        case "$PARAM" in
            pfc-enabled-priorities) echo "enabled-priorities=$(pfc_get)" ;;
            trust-mode) echo "trust-mode=$(trust_get)" ;;
            prio-tc-map) echo "prio-tc-map=$(prio_tc_get)" ;;
            tos-unmap) echo "tos-unmap=$(tos_unmap_get)" ;;
        esac
    done
elif [[ "$OP" == "SET" ]]; then
    # Error policy: a SET that the kernel rejects must surface as a non-zero
    # exit. Previous version unconditionally echoed "key=value" after every
    # helper call, masking failures (e.g. dcb pfc set on an unconfigured iface)
    # as success. Helpers now return non-zero on real failure; we fold that
    # into the wrapper's exit code and only echo "key=value" when the helper
    # actually applied the value.
    rc=0
    for PAIR in "$@"; do
        PARAM="${PAIR%%=*}"
        VALUE="${PAIR#*=}"
        param_rc=0
        case "$PARAM" in
            pfc-enabled-priorities) pfc_set "$VALUE" || param_rc=$? ;;
            trust-mode) trust_set "$VALUE" || param_rc=$? ;;
            prio-tc-map) prio_tc_set "$VALUE" || param_rc=$? ;;
            tos-unmap) tos_unmap_set "$VALUE" || param_rc=$? ;;
        esac
        if [[ $param_rc -eq 0 ]]; then
            echo "${PARAM}=${VALUE}"
        else
            rc=$param_rc
        fi
    done
    exit $rc
fi
