#!/bin/bash

set -eu

if [[ $EUID -ne 0 ]];
then
    exec pkexec --disable-internal-agent "$0" "$@"
fi

error() {
   echo "ERROR at ${1}:${2}" >&2
   exit 1
}
trap 'error "${BASH_SOURCE}" "${LINENO}"' ERR

# Gett list of `rtw89` directories in `/sys/kernel/debug`.  We store in an array
# because there can be more then one if someone has plugged more than one dongle in
mapfile -t RTW89_DEBUGFS_DIRS < <(find /sys/kernel/debug/ieee80211 -name rtw89 -type d)

if [[ "${#RTW89_DEBUGFS_DIRS[@]}" == "0" ]]; then
    echo "No realtek dongles found, or driver too old!" >&2
    exit 1
fi

declare -A BITS
declare -A STATE

function parse_disable_dm() {
    local disable_dm_path="${1}"
    BITS=()
    STATE=()

    # Parse feature file
    while IFS= read -r line; do
        if [[ $line =~ ^\[([0-9]+)\][[:space:]]+([A-Z0-9_]+):[[:space:]]+([OX])$ ]]; then
            bit="${BASH_REMATCH[1]}"
            name="${BASH_REMATCH[2]}"
            val="${BASH_REMATCH[3]}"

            if [[ "$val" == "X" ]]; then
                val="1"
            fi

            BITS["$name"]="$bit"
            STATE["$name"]="$val"
        fi
    done < "${disable_dm_path}"
}

function print_mask() {
    # Build bitmask from file
    local mask=0
    for name in "${!BITS[@]}"; do
        if [[ ${STATE[$name]} == "1" ]]; then
            # I'm using `2 ** $X` here instead of `2 << $X` because vscode's heredoc highlighting triggers otherwise :(
            mask=$(( mask | ( 2 ** BITS[$name] ) ))
        fi
    done
    printf "0x%x\n" "$mask"
}

function boolify() {
    case "$1" in
        1|on|true|enable|enabled)
            echo -n "1"
            ;;
        0|off|false|disable|disabled)
            echo -n "0"
            ;;
        *)
            echo "Unable to convert '${1}' into a boolean!" >&2
            exit 1
            ;;
    esac
}

# Super confusing to deal in 0/1 due to this being a "disable" file.  Convert from `disabled` -> `1`.
function boolify_disable() {
    case "$(boolify "$1")" in
        1)
            echo -n "0"
            ;;
        0)
            echo -n "1"
            ;;
    esac
}

function set_feature() {
    local name="$1"
    local val="$2"

    if [[ ! -v STATE[$name] ]]; then
        echo "Unknown feature: $name" >&2
        return 1
    fi

    STATE[$name]="$(boolify_disable "$val")"
}

function set_pd_thresh() {
    local val="$1"
    echo "$1" | static_pd_th
}


function log2_floor()
{
    local x="${1}"
    local count="0"
    while [ "${x}" -gt 0 ]; do
        count=$((count + 1))
        x=$((x / 2))
    done
    echo -n "${count}"
}

function build_edca()
{
    TXOP="${1}"
    CWMAX="${2}"
    CWMIN="${3}"
    AIFSN="${4}"
    
    # cwmin/cwmax are, as in the 802.11 standard, locked to values of the form 2^n - 1
    # which are efficiently encoded as `int(log2(cwmax)) + 1`
    # txop is how long we pad our channel reservation, rounded up to the nearest 32us
    # aifs is 16 + 9*aifsn, which is what gets stored
    # Final bit layout is:
    #   int((txop - 1)/32) + 1 | int(log2(cwmax)) + 1 | int(log2(cwmin)) + 1 | 16 + 9*aifsn
    printf "%04X%1X%1X%02X\n" "$(( ( "${TXOP}" - 1 ) / 32 + 1))" "$(log2_floor "${CWMAX}")" "$(log2_floor "${CWMIN}")" "$((16 + 9 * "${AIFSN}"))"
}

function print_led_color()
{
    local LED_DIR="${1}"
    if [[ ! -d "${LED_DIR}" ]]; then
        return
    fi

    local BRIGHTNESS="$(cat "${LED_DIR}/brightness")"
    if [[ "${BRIGHTNESS}" == "0" ]]; then
        echo "LED: off"
    else
        INTENSITIES=( $(cat "${LED_DIR}/multi_intensity") )
        NAMES=( white red green blue )
        printf "LED: "
        for idx in 1 2 3 4; do
            if [[ "${INTENSITIES[${idx}]}" == "1" ]]; then
                printf "[${NAMES[${idx}]}] "
            else
                printf "${NAMES[${idx}]} "
            fi
        done
        printf "\n"
    fi
}


function print_usage() {
    local prog_name="$(basename "$0")"

    cat >&2 <<-EOF
Usage:

    $prog_name list                           - show all parameters for all connected rtw89 phy's
    $prog_name set-low-latency-mode [on|off]  - Turn on/off low-latency mode
    $prog_name reset                          - Shorthand for 'set-low-latency-mode off'

To manually set different parameters that are part of low latency mode, use:

    $prog_name set-amsdu [on|off]             - default is 'on'
    $prog_name set-feature HW_SCAN [on|off]   - default is 'on'
    $prog_name set-pd-thresh <val>            - default is '0'
    $prog_name set-powersave [on|off]         - default is 'on'
    $prog_name set-aggressive-edca [on|off]   - default is 'off'

NOTE: default values refer to when low-latency mode is disabled.
EOF
    exit 1
}

# For each dongle plugged in, set/unset the requested parameters
for DEBUGFS_DIR in "${RTW89_DEBUGFS_DIRS[@]}"; do
    PHY_NAME="$(basename "$(dirname "${DEBUGFS_DIR}" )" )"
    LED_DIR="/sys/class/leds/rtw89-${PHY_NAME}-multicolor"
    WLAN_NAME="$(basename "$(compgen -G "$(dirname "${DEBUGFS_DIR}" )/netdev:*")" | cut -d: -f2-)"

    # Load the current state
    parse_disable_dm "${DEBUGFS_DIR}/disable_dm"

    case "${1:-}" in
        set-low-latency-mode)
            if [[ -z "${2:-}" ]]; then
                print_usage
            fi
            case "$(boolify "${2}")" in
                1)
                    "$0" set-feature HW_SCAN off
                    "$0" set-amsdu off
                    "$0" set-mcs-rates 0-9
                    "$0" set-powersave off
                    "$0" set-pd-thresh -75
                    "$0" set-aggressive-edca off
                    ;;
                0)
                    "$0" reset
                    ;;
            esac
            ;;

        # Set aggregation of SDU's; realtek's driver has a bug with this aggregation, disable it for low-latency mode.
        set-amsdu)
            if [[ -z "${2:-}" ]]; then
                print_usage
            fi

            case "$(boolify "${2}")" in
                1)
                    iw dev "${WLAN_NAME}" set tidconf tids 0xff amsdu on
                    ;;
                0)
                    iw dev "${WLAN_NAME}" set tidconf tids 0xff amsdu off
                    ;;
            esac
            ;;

        # Set MCS rates; by limiting MCS rates to lower values, we improve resilience while limiting throughput
        set-mcs-rates)
            # Usually something like "0-9", but empty means default
            MCS_RATES="${2:-}"

            if [[ -z "${MCS_RATES}" ]]; then
                # Reset to default, ignore failure
                iw dev "${WLAN_NAME}" set bitrates 2>/dev/null || true
            else
                # Apply the MCS rates to each band, and each spatial stream
                IW_ARGS=()
                for BAND in 5 6; do
                    IW_ARGS+=( "he-mcs-${BAND}" )
                    for SPATIAL_STREAM in 1 2; do
                        IW_ARGS+=( "${SPATIAL_STREAM}:${MCS_RATES}" )
                    done
                done
                iw dev "${WLAN_NAME}" set bitrates "${IW_ARGS[@]}"
            fi
            ;;

        # Disable power save features, both firmware-internal powersave and host-side powersave.
        set-powersave)
            if [[ -z "${2:-}" ]]; then
                print_usage
            fi

            case "$(boolify "${2}")" in
                1)
                    iw dev "${WLAN_NAME}" set power_save on
                    set_feature "INACTIVE_PS" on
                    ;;
                0)
                    iw dev "${WLAN_NAME}" set power_save off
                    set_feature "INACTIVE_PS" off
                    ;;
            esac
            print_mask > "${DEBUGFS_DIR}/disable_dm"
            ;;

        # Set a static packet detection threshold attempt to ignore distant interferers
        set-pd-thresh)
            if [[ -z "${2:-}" ]]; then
                print_usage
            fi

            if [[ "${2}" == 0 ]]; then
                # packet detection threshold of 0 means default
                echo "0" >"${DEBUGFS_DIR}/static_pd_th"
            else
                # Typical values we might set: -75, -55.  Dynamic range is from -102 ~ -40
                echo "1 ${2}" >"${DEBUGFS_DIR}/static_pd_th"
            fi
            ;;

        # Set aggressive EDCA (packet timing) values, so we attempt to retry faster
        set-aggressive-edca)
            # We set values for each QoS category (we treat all equally)
            REGISTER_ADDRS=( c300 c304 c308 c30c )

            case "$(boolify "${2}")" in
                1)
                    # Keep these roughly inline with what Steam manually sets on Windows
                    EDCA_PARAMS="$(build_edca 2000 7 3 1)"
                    for REG_ADDR in "${REGISTER_ADDRS[@]}"; do
                        echo "${REG_ADDR} ${EDCA_PARAMS} 4" >"${DEBUGFS_DIR}/write_reg"
                    done
                    ;;
                0)
                    # We don't actually do anything here; EDCA resets upon disassociation
                    ;;
            esac
            ;;

        # Helper to set firmware toggles
        set-feature)
            set_feature "$2" "$3"
            print_mask > "${DEBUGFS_DIR}/disable_dm"
            ;;

        list)
            echo "-> ${PHY_NAME} (${WLAN_NAME})"
            cat "${DEBUGFS_DIR}/disable_dm"
            cat "${DEBUGFS_DIR}/static_pd_th" | grep -E 'DIG:|OFDM'
            print_led_color "${LED_DIR}"
            ;;
        reset)
            echo "0x00" > "${DEBUGFS_DIR}/disable_dm"
            "$0" set-pd-thresh 0
            "$0" set-aggressive-edca disable
            "$0" set-powersave on
            "$0" set-amsdu on
            "$0" set-mcs-rates
            ;;
        *)
            print_usage
            ;;
    esac
done
