#!/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
#
# mlxreg tool wrapper — real hardware version
# Handles register-based params (ROCE_ACCL, PIPG, PAOS, etc.)
source "$(dirname "$0")/_lib.sh"
source "$(dirname "$0")/_value_maps.sh"
trace "$@"

declare -A LEAF=(
    ["roce_adp_retrans_en"]="adaptive-retransmission"
    ["adaptive_routing_forced_en"]="adaptive-routing-force"
    ["selective_repeat_forced_en"]="selective-repeat"
    ["roce_slow_restart_en"]="slow-restart"
    ["roce_slow_restart_idle_en"]="slow-restart-idle"
    ["roce_tx_window_en"]="tx-window"
    ["cc_per_plane_en"]="cc-per-plane"
    ["cc_probe_mp_mode"]="cc-probe-mp-mode"
    ["roce_adp_retrans_field_select"]="adaptive-retransmission"
    ["adaptive_routing_forced_en_field_select"]="adaptive-routing-force"
    ["selective_repeat_forced_en_field_select"]="selective-repeat"
    ["roce_slow_restart_field_select"]="slow-restart"
    ["roce_slow_restart_idle_field_select"]="slow-restart-idle"
    ["roce_tx_window_field_select"]="tx-window"
    ["cc_per_plane_en_field_select"]="cc-per-plane"
    ["cc_probe_mp_mode_field_select"]="cc-probe-mp-mode"
    ["cc-per-plane-supported"]="cc-per-plane-supported"
    ["is-supported"]="is-supported"
    ["max-profiles"]="max-profiles"
    ["max-ranges-per-profile"]="max-ranges-per-profile"
    ["min-base-timeout"]="min-base-timeout"
)

# ─────────────────────────────────────────────────────────────────────────
# Adaptive-retransmission profiles (ROCE_ACCL register, adp_retx_* fields).
# Mirrors the gNMI reference in gnxi/gnoi/gnmi/commands.json: a profile is
# programmed with a single "mlxreg --set adp_retx_profile_id=<id>,<field>=<v>,…"
# call and read back field-by-field from ROCE_ACCL. Per-range fields carry a
# "_<range-id>" suffix (e.g. range_low_bound_0).
# ─────────────────────────────────────────────────────────────────────────

# True when $1 is an adp-retx register field (range fields keep their _<id> suffix).
adp_retx_is_field() {
    case "${1%_[0-9]*}" in
        adp_retx_profile_id|time_base|time_unit|range_num|timeout_init_low_bound|\
        timeout_init_range_size|start_range_index|qp_total_timeout|retx_total_timeout|\
        range_low_bound|range_size|timeout_retry_num|dec_mode|prev_range_index)
            return 0 ;;
        *) return 1 ;;
    esac
}

# Map a register field name to its YANG leaf name (range suffix stripped).
adp_retx_leaf_name() {
    case "${1%_[0-9]*}" in
        adp_retx_profile_id) echo "profile-id" ;;
        time_base) echo "time-base" ;;
        time_unit) echo "time-unit" ;;
        range_num) echo "range-num" ;;
        timeout_init_low_bound) echo "timeout-init-low-bound" ;;
        timeout_init_range_size) echo "timeout-init-range-size" ;;
        start_range_index) echo "start-range-index" ;;
        qp_total_timeout) echo "qp-total-timeout" ;;
        retx_total_timeout) echo "retx-total-timeout" ;;
        range_low_bound) echo "range-low-bound" ;;
        range_size) echo "range-size" ;;
        timeout_retry_num) echo "timeout-retry-num" ;;
        dec_mode) echo "dec-mode" ;;
        prev_range_index) echo "prev-range-index" ;;
        *) echo "${1%_[0-9]*}" ;;
    esac
}

# Encode a YANG value into the numeric form mlxreg expects (enums/bools).
adp_retx_encode() {
    local field="${1%_[0-9]*}" value="$2"
    case "$field" in
        time_unit)        case "$value" in usec|1) echo 1 ;; *) echo "$value" ;; esac ;;
        qp_total_timeout) case "$value" in true|True|TRUE|1) echo 1 ;; *) echo 0 ;; esac ;;
        dec_mode)
            case "$value" in
                div-by-4|0) echo 0 ;;
                div-by-2|1) echo 1 ;;
                low-bound|2) echo 2 ;;
                *) echo "$value" ;;
            esac ;;
        *) echo "$value" ;;
    esac
}

# Decode a numeric register value into the YANG value (enums/bools).
adp_retx_decode() {
    local field="${1%_[0-9]*}" value="$2"
    case "$field" in
        time_unit)        case "$value" in 1) echo usec ;; *) echo "$value" ;; esac ;;
        qp_total_timeout) [[ "$value" == "1" ]] && echo true || echo false ;;
        dec_mode)
            case "$value" in
                0) echo div-by-4 ;;
                1) echo div-by-2 ;;
                2) echo low-bound ;;
                *) echo "$value" ;;
            esac ;;
        *) echo "$value" ;;
    esac
}

# GET: read ROCE_ACCL once and emit "<yang-leaf>=<value>" per requested field.
adp_retx_get() {
    local out param raw val
    # A failed read must not be decoded as real register state (0-derived values).
    if ! out=$(mlxreg -d "$BDF" --get --reg_name ROCE_ACCL 2>/dev/null); then
        tool_fail 1 "mlxreg-wrapper: ROCE_ACCL read failed on ${BDF} (adp-retx GET)"
    fi
    for param in "$@"; do
        raw=$(echo "$out" | grep -Ew "$param" | head -1 | cut -d'|' -f2 | tr -d ' ' | sed 's/0x//')
        if [[ -n "$raw" ]]; then
            val=$((16#${raw}))
        else
            val=0
        fi
        echo "$(adp_retx_leaf_name "$param")=$(adp_retx_decode "$param" "$val")"
    done
}

# SET: program all supplied fields in a single ROCE_ACCL modify, echo readable KV.
adp_retx_set() {
    local pair param value fields="" line err
    local -a out_lines=()
    for pair in "$@"; do
        param="${pair%%=*}"
        value="${pair#*=}"
        fields="${fields},${param}=$(adp_retx_encode "$param" "$value")"
        out_lines+=("$(adp_retx_leaf_name "$param")=${value}")
    done
    fields="${fields#,}"
    err=$(mktemp /tmp/dms-mlxreg-adp-retx.XXXXXX) || exit 1
    trap 'rm -f "$err"' EXIT
    if ! mlxreg -d "$BDF" -y --set "$fields" --reg_name ROCE_ACCL 2>"$err"; then
        # Flatten stderr: tool_fail forwards only the first stdout line.
        ERR=$(tr '\n' ' ' <"$err" 2>/dev/null | sed -E 's/[[:space:]]+/ /g; s/^ +| +$//g')
        tool_fail 1 "mlxreg-wrapper: ROCE_ACCL adp-retx SET failed: ${ERR:-unknown}"
    fi
    for line in "${out_lines[@]}"; do
        echo "$line"
    done
}

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

# Route adaptive-retransmission profile fields to the dedicated handler. Every arg
# must belong to the adp-retx family; a mixed batch is rejected so non-adp params
# never leak into the adp-retx ROCE_ACCL path (fabricated GET values / bad SET).
adp_retx_seen=0
adp_retx_all=1
for _arg in "$@"; do
    if adp_retx_is_field "${_arg%%=*}"; then
        adp_retx_seen=1
    else
        adp_retx_all=0
    fi
done
if [[ $adp_retx_seen -eq 1 ]]; then
    [[ $adp_retx_all -eq 1 ]] || \
        tool_fail 1 "mlxreg-wrapper: adp-retx fields cannot be mixed with other ROCE_ACCL params in one request"
    case "$OP" in
        GET) adp_retx_get "$@" ;;
        SET) adp_retx_set "$@" ;;
        *) tool_fail 1 "mlxreg-wrapper: unsupported op '$OP' for adp-retx fields" ;;
    esac
    exit $?
fi

field_value() {
    local output="$1" field="$2"
    awk -v want="$field" '
        {
            line=$0
            sub(/^[[:space:]]+/, "", line)
            first=line
            sub(/[[:space:]|:].*$/, "", first)
            if (first == want) {
                n = split(line, a, /[[:space:]|:]+/)
                print a[n]
                exit
            }
        }' <<<"$output"
}

to_bool_bit() {
    case "$1" in
        true|True|1|0x1|0x00000001) echo 1 ;;
        false|False|0|0x0|0x00000000) echo 0 ;;
        *) return 1 ;;
    esac
}

to_enum_raw() {
    local param="$1" value="$2" mapped
    mapped="${FWD[$param:$value]:-$value}"
    case "$mapped" in
        0|1) echo "$mapped" ;;
        0x0|0x00000000) echo 0 ;;
        0x1|0x00000001) echo 1 ;;
        *) return 1 ;;
    esac
}

if [[ "$OP" == "GET" || "$OP" == "GET-CAPS" ]]; then
    # For ROCE_ACCL register: query all fields at once
    ACCL_OUTPUT=""
    NEEDS_ACCL=0
    for PARAM in "$@"; do
        case "$PARAM" in
            roce_adp_retrans_en|adaptive_routing_forced_en|selective_repeat_forced_en|\
            roce_slow_restart_en|roce_slow_restart_idle_en|roce_tx_window_en|cc_per_plane_en|\
            cc_probe_mp_mode|\
            roce_adp_retrans_field_select|adaptive_routing_forced_en_field_select|\
            selective_repeat_forced_en_field_select|roce_slow_restart_field_select|\
            roce_slow_restart_idle_field_select|roce_tx_window_field_select|\
            cc_per_plane_en_field_select|cc_probe_mp_mode_field_select|\
            cc-per-plane-supported)
                NEEDS_ACCL=1 ;;
        esac
    done

    if [[ $NEEDS_ACCL -eq 1 && -z "$ACCL_OUTPUT" ]]; then
        if ! ACCL_OUTPUT=$(mlxreg -d "$BDF" --get --reg_name ROCE_ACCL 2>/dev/null); then
            # A failed ROCE_ACCL read is a real error for a config GET, so surface it
            # instead of synthesizing false/default values. For GET-CAPS an unreadable
            # register means the features are unavailable, so fall through and let the
            # per-field logic report them as false (matches mlxreg-roce-accl-caps).
            [[ "$OP" == "GET-CAPS" ]] || \
                tool_fail 1 "mlxreg-wrapper: ROCE_ACCL read failed on ${BDF}"
            ACCL_OUTPUT=""
        fi
    fi

    for PARAM in "$@"; do
        LN="${LEAF[$PARAM]:-$PARAM}"
        case "$PARAM" in
            roce_adp_retrans_field_select|adaptive_routing_forced_en_field_select|\
            selective_repeat_forced_en_field_select|roce_slow_restart_field_select|\
            roce_slow_restart_idle_field_select|roce_tx_window_field_select|\
            cc_per_plane_en_field_select|cc_probe_mp_mode_field_select)
                VAL=$(field_value "$ACCL_OUTPUT" "$PARAM")
                [[ "$VAL" == "0x00000001" || "$VAL" == "0x1" || "$VAL" == "1" ]] && echo "${LN}=true" || echo "${LN}=false" ;;
            roce_adp_retrans_en|adaptive_routing_forced_en|selective_repeat_forced_en|\
            roce_slow_restart_en|roce_slow_restart_idle_en|roce_tx_window_en|cc_per_plane_en)
                VAL=$(field_value "$ACCL_OUTPUT" "$PARAM" | sed 's/0x//')
                [[ -n "$VAL" ]] && VAL=$((16#${VAL}))
                [[ "$VAL" == "1" ]] && echo "${LN}=true" || echo "${LN}=false" ;;
            cc_probe_mp_mode)
                RAW=$(field_value "$ACCL_OUTPUT" "$PARAM")
                if [[ -z "$RAW" ]]; then
                    # ROCE_ACCL read failed or field absent: fall back to the enum
                    # default rather than failing the whole batch invocation.
                    echo "${LN}=default"
                else
                    MAPPED="${REV[$PARAM:$RAW]:-}"
                    if [[ -z "$MAPPED" ]]; then
                        RAW_HEX="${RAW#0x}"
                        [[ -n "$RAW_HEX" ]] && RAW_DEC=$((16#$RAW_HEX)) || RAW_DEC=""
                        MAPPED="${REV[$PARAM:$RAW_DEC]:-}"
                    fi
                    if [[ -z "$MAPPED" ]]; then
                        tool_fail 1 "mlxreg-wrapper: reserved cc_probe_mp_mode value '${RAW}'"
                    fi
                    echo "${LN}=${MAPPED}"
                fi ;;
            cc-per-plane-supported)
                SEL=$(field_value "$ACCL_OUTPUT" "cc_per_plane_en_field_select" | sed 's/0x//')
                [[ -n "$SEL" ]] && SEL=$((16#${SEL}))
                [[ "$SEL" == "1" ]] && echo "${LN}=true" || echo "${LN}=false" ;;
            is-supported|max-profiles|max-ranges-per-profile|min-base-timeout)
                echo "${LN}=0" ;;
            *)
                echo "${LN}=unknown" ;;
        esac
    done
elif [[ "$OP" == "SET" ]]; then
    declare -A SELECT=(
        ["roce_adp_retrans_en"]="roce_adp_retrans_field_select"
        ["roce_tx_window_en"]="roce_tx_window_field_select"
        ["roce_slow_restart_en"]="roce_slow_restart_field_select"
        ["roce_slow_restart_idle_en"]="roce_slow_restart_idle_field_select"
        ["adaptive_routing_forced_en"]="adaptive_routing_forced_en_field_select"
        ["selective_repeat_forced_en"]="selective_repeat_forced_en_field_select"
        ["cc_per_plane_en"]="cc_per_plane_en_field_select"
        ["cc_probe_mp_mode"]="cc_probe_mp_mode_field_select"
    )
    SET_FIELDS=""
    declare -A VALUE_OUT=()
    for PAIR in "$@"; do
        PARAM="${PAIR%%=*}"
        VALUE="${PAIR#*=}"
        if [[ "$PARAM" == "cc_probe_mp_mode" ]]; then
            # to_enum_raw already accepts default/discard-on-plane-mismatch plus the
            # numeric/hex forms via the shared FWD map, so a failure here means the
            # value is genuinely invalid.
            if ! V=$(to_enum_raw "$PARAM" "$VALUE"); then
                tool_fail 1 "mlxreg-wrapper: invalid cc_probe_mp_mode value '$VALUE'"
            fi
            case "$V" in
                1) VALUE_OUT[$PARAM]=discard-on-plane-mismatch ;;
                0) VALUE_OUT[$PARAM]=default ;;
            esac
        elif ! V=$(to_bool_bit "$VALUE"); then
            tool_fail 1 "mlxreg-wrapper: invalid boolean value '$VALUE' for $PARAM"
        fi
        SET_FIELDS="${SET_FIELDS},${PARAM}=${V}"
        SEL="${SELECT[$PARAM]}"
        [[ -n "$SEL" ]] && SET_FIELDS="${SET_FIELDS},${SEL}=1"
    done
    SET_FIELDS="${SET_FIELDS#,}"
    ERR_FILE=$(mktemp /tmp/dms-mlxreg-wrapper.XXXXXX) || exit 1
    trap 'rm -f "$ERR_FILE"' EXIT

    if ! mlxreg -d "$BDF" -y --set "$SET_FIELDS" --reg_name ROCE_ACCL 2>"$ERR_FILE"; then
        # Flatten stderr: tool_fail forwards only the first stdout line.
        ERR=$(tr '\n' ' ' <"$ERR_FILE" 2>/dev/null | sed -E 's/[[:space:]]+/ /g; s/^ +| +$//g')
        tool_fail 1 "mlxreg-wrapper: ROCE_ACCL SET failed: ${ERR:-unknown}"
    fi
    for PAIR in "$@"; do
        PARAM="${PAIR%%=*}"
        VALUE="${PAIR#*=}"
        LN="${LEAF[$PARAM]:-$PARAM}"
        echo "${LN}=${VALUE_OUT[$PARAM]:-$VALUE}"
    done
fi
