#!/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"
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"]="cc-per-plane"
    ["cc_probe_mp_mode"]="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
    out=$(mlxreg -d "$BDF" --get --reg_name ROCE_ACCL 2>/dev/null)
    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
        echo "Error: mlxreg ROCE_ACCL adp-retx SET failed: $(cat "$err" 2>/dev/null)" >&2
        exit 1
    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.
for _arg in "$@"; do
    if adp_retx_is_field "${_arg%%=*}"; then
        case "$OP" in
            GET) adp_retx_get "$@" ;;
            SET) adp_retx_set "$@" ;;
            *) echo "Error: unsupported op '$OP' for adp-retx fields" >&2; exit 1 ;;
        esac
        exit $?
    fi
done

if [[ "$OP" == "GET" ]]; 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|\
            cc_probe_mp_mode|cc-per-plane-supported)
                NEEDS_ACCL=1 ;;
        esac
    done

    if [[ $NEEDS_ACCL -eq 1 && -z "$ACCL_OUTPUT" ]]; then
        ACCL_OUTPUT=$(mlxreg -d "$BDF" --get --reg_name ROCE_ACCL 2>/dev/null)
    fi

    for PARAM in "$@"; do
        LN="${LEAF[$PARAM]:-$PARAM}"
        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)
                VAL=$(echo "$ACCL_OUTPUT" | grep -E "^\s*${PARAM}\s" | awk '{print $NF}' | sed 's/0x//')
                [[ -n "$VAL" ]] && VAL=$((16#${VAL}))
                [[ "$VAL" == "1" ]] && echo "${LN}=true" || echo "${LN}=false" ;;
            cc_probe_mp_mode)
                # Enum leaf: 0 = default, 1 = discard-on-plane-mismatch.
                VAL=$(echo "$ACCL_OUTPUT" | grep -E "^\s*${PARAM}\s" | awk '{print $NF}' | sed 's/0x//')
                [[ -n "$VAL" ]] && VAL=$((16#${VAL}))
                [[ "$VAL" == "1" ]] && echo "${LN}=discard-on-plane-mismatch" || echo "${LN}=default" ;;
            cc-per-plane-supported)
                SEL=$(echo "$ACCL_OUTPUT" | grep -E "^\s*cc_per_plane_en_field_select\s" | awk '{print $NF}' | 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"]="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
            # Enum leaf: accept numeric or string form. The agentless JSON path
            # stringifies a numeric enum, so "1"/"0" arrive alongside the names.
            case "$VALUE" in
                1|discard-on-plane-mismatch) V=1; VALUE_OUT[$PARAM]=discard-on-plane-mismatch ;;
                0|default)                   V=0; VALUE_OUT[$PARAM]=default ;;
                *) echo "Error: bad cc_probe_mp_mode '$VALUE' (expected default/0 or discard-on-plane-mismatch/1)" >&2; exit 1 ;;
            esac
        else
            [[ "$VALUE" == "true" ]] && V=1 || V=0
        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
        ERR=$(cat "$ERR_FILE" 2>/dev/null)
        echo "Error: mlxreg ROCE_ACCL SET failed: ${ERR:-unknown}" >&2
        exit 1
    fi
    for PAIR in "$@"; do
        PARAM="${PAIR%%=*}"
        VALUE="${PAIR#*=}"
        LN="${LEAF[$PARAM]:-$PARAM}"
        echo "${LN}=${VALUE_OUT[$PARAM]:-$VALUE}"
    done
fi
