#!/usr/bin/env bash

USER_STATES_PATH=${USER_STATES_PATH:-/etc/default/amdgpu-custom-state}
SYS_PPFMASK=/sys/module/amdgpu/parameters/ppfeaturemask
[ ! -r ${SYS_PPFMASK} ] && echo "Can't access ${SYS_PPFMASK}" && exit 2

if [ "$1" == "restore" ]; then
  RESTORE=true
fi

function check_ppfeaturemask() {
  CURRENT_HEX_MASK=$(printf '%#x' "$(( $(cat ${SYS_PPFMASK}) ))")
  # 0x4000 or 14th bit is one indicating if OverDrive has been enabled
  OVERDRIVE_MASK=$(printf '%#x' "$(( CURRENT_HEX_MASK & 0x4000 ))")
  [ "${OVERDRIVE_MASK}" == "0x4000" ] && return 0
  WANTED_MASK=$(printf '%#x' "$(( CURRENT_HEX_MASK | 0x4000 ))")
  echo -n "In order to set custom amdgpu power states, enable OverDrive by "
  echo -n "booting the machine with amdgpu.ppfeaturemask=${WANTED_MASK} "
  echo    "kernel option" && exit 2
}

function fill_sclks() {
  if [ -z "${3}" ]; then     # Vega20 and later ASICs
    echo "  SCLK state ${1}: ${2}"
    SCLK[${1}]="s ${1} ${2%*M[Hh]z}"
  else                       # Vega10 and previous ASICs
    echo "  SCLK state ${1}: ${2}, ${3}"
    SCLK[${1}]="s ${1} ${2%*M[Hh]z} ${3%*mV}"
  fi
}

function fill_mclks() {
  if [ -z "${3}" ]; then     # Vega20 and later ASICs
    echo "  MCLK state ${1}: ${2}"
    MCLK[${1}]="m ${1} ${2%*M[Hh]z}"
  else                       # Vega10 and previous ASICs
    echo "  MCLK state ${1}: ${2}, ${3}"
    MCLK[${1}]="m ${1} ${2%*M[Hh]z} ${3%*mV}"
  fi
}

function fill_vddccurve() {
  if [ $# -eq 3 ]; then
    echo "  VDDC Curve state $1: ${2} ${3}"
    VDDC_CURVE[${1}]="vc ${1} ${2%*M[Hh]z} ${3%*mV}"
  elif [ $# -eq 4 ]; then
    echo "  VDDC Curve state $1: ${2} @ ${4}"
    VDDC_CURVE[${1}]="vc ${1} ${2%*M[Hh]z} ${4%*mV}"
  else
    echo "ERROR: Can't parse '$*', did pp_od_clk_voltage API change again?"
    exit 2
  fi
}

function fill_vddgfx_offset() {
  echo "  VDD GFX Offset: $1"
  VDDGFX_OFFSET="vo ${1%*mV}"
}

function parse_states() {
  mapfile -t STATE_LINES < <(tr -d '\000' < "$1")
  for LNNO in "${!STATE_LINES[@]}"; do
    LINE="${STATE_LINES[$LNNO]}"
    case ${LINE} in
      "OD_SCLK:")
        state_fill_func=fill_sclks ;;
      "OD_MCLK:")
        state_fill_func=fill_mclks ;;
      "OD_VDDC_CURVE:")
        state_fill_func=fill_vddccurve ;;
      "VDDC_CURVE_SCLK["[012]"]: "*)
        ;; # Just ignoring these for now
      "VDDC_CURVE_VOLT["[012]"]: "*)
        ;; # Just ignoring these for now
      "OD_VDDGFX_OFFSET:")
        offset_fill_func=fill_vddgfx_offset ;;
      "OD_RANGE:")
        echo "  Maximum clocks & voltages:";;
      "SCLK: "*)
        echo "    SCLK clock ${LINE##* }"
        MAX_SCLK=${LINE##* }
        MAX_SCLK=${MAX_SCLK%*M[Hh]z}
        ;;
      "MCLK: "*)
        echo "    MCLK clock ${LINE##* }"
        MAX_MCLK=${LINE##* }
        MAX_MCLK=${MAX_MCLK%*M[Hh]z}
        ;;
      "VDDGFX_OFFSET: "*)
        echo "    VDDGFX offset range: ${LINE#*: }"
        OFFSET_RANGE=${LINE#*: }
        ;;
      "VDDC: "*)
        echo "    VDDC voltage ${LINE##* }"
        MAX_VDDC=${LINE##* }
        MAX_VDDC=${MAX_VDDC%*mV}
        ;;
      [0-9]": "*)
        $state_fill_func ${LINE%%:*} ${LINE#* }
        ;;
      [-0-9]*"mV")
        $offset_fill_func ${LINE%%:*} ${LINE#* }
        ;;
      "FORCE_SCLK: "[0-9]*)
        echo "  Force SCLK state to ${LINE#* }"
        FORCE_SCLK=${LINE#* }
        ;;
      "FORCE_MCLK: "[0-9]*)
        echo "  Force MCLK state to ${LINE#* }"
        FORCE_MCLK=${LINE#* }
        ;;
      "FORCE_POWER_CAP: "[0-9]*)
        MICROWATTS=${LINE#* }
        echo "  Force power cap to ${MICROWATTS%*000000}W"
        FORCE_POWER_CAP=${LINE#* }
        ;;
      "FORCE_PERF_LEVEL: "[a-z]*)
        echo "  Force performance level to ${LINE#* }"
        FORCE_LEVEL=${LINE#* }
        ;;
      "FORCE_POWER_PROFILE: "[0-9]*)
        echo "  Force power profile to ${LINE#* }"
        FORCE_PROFILE=${LINE#* }
        ;;
      "#"*) ;;
      "") ;;
      *)
        echo "  Unexpected value in ${1}:$((LNNO + 1))"
        exit 2
        ;;
    esac
  done
  if [ "$1" == "${SYS_PP_OD_CLK}" ]; then
    DETECTED_MCLK=("${MCLK[@]}")
    DETECTED_VDDC_CURVE=("${VDDC_CURVE[@]}")
    POWER_LIMIT=$(cat "$PWR_CAP_FILE")
    echo "  Curent power cap: ${POWER_LIMIT%*000000}W"
  fi
}

function set_custom_states() {
  if [ "${RESTORE}" == "true" ] && [ "${FORCE_LEVEL}" == "auto" ]; then
    echo manual > "${PWR_LEVEL}"
    sleep 0.25
    echo 0 > "${SYS_DPM_SCLK}" 2>&1
    echo 0 > "${SYS_DPM_MCLK}" 2>&1
    sleep 0.25
  fi
  for CSTATE in "${SCLK[@]}"; do
    echo "${CSTATE}" > "${SYS_PP_OD_CLK}"
  done
  for MSTATE in "${MCLK[@]}"; do
    OLD_MSTATE=${DETECTED_MCLK[$((${MSTATE:3:1}))]}
    if [ "${MSTATE}" != "${OLD_MSTATE}" ]; then
      echo "${MSTATE}" > "${SYS_PP_OD_CLK}" ||
        echo "ERROR: echo ${MSTATE} > ${SYS_PP_OD_CLK}"
    fi
  done
  for VDDC_CURVE_STATE in "${VDDC_CURVE[@]}"; do
    OLD_CURVE_STATE=${DETECTED_VDDC_CURVE[$((${VDDC_CURVE_STATE:3:1}))]}
    if [ "${VDDC_CURVE_STATE}" != "${OLD_CURVE_STATE}" ]; then
      echo "${VDDC_CURVE_STATE}" > "${SYS_PP_OD_CLK}" ||
        echo "ERROR: echo ${VDDC_CURVE_STATE} > ${SYS_PP_OD_CLK}"
    fi
  done
  if [ "${VDDGFX_OFFSET}" ]; then
    echo "${VDDGFX_OFFSET}" > "${SYS_PP_OD_CLK}"
  fi
  echo 'c' > "${SYS_PP_OD_CLK}"
  if [ "${FORCE_LEVEL}" ]; then
    echo ${FORCE_LEVEL} > ${PWR_LEVEL}
  fi
  if [ "${FORCE_PROFILE}" ]; then
    echo ${FORCE_PROFILE} > ${PWR_PROFILE}
  fi
  if [ "${FORCE_SCLK}" ]; then
    echo "${FORCE_SCLK}" > "${SYS_DPM_SCLK}"
  fi
  if [ "${FORCE_MCLK}" ]; then
    echo "${FORCE_MCLK}" > "${SYS_DPM_MCLK}"
  fi
  if [ "${FORCE_POWER_CAP}" ]; then
    echo "${FORCE_POWER_CAP}" > "${PWR_CAP_FILE}"
  fi
}

function backup_states() {
  if [ ! -r "${BACKUP_STATE_FILE}" ]; then
    cp "${SYS_PP_OD_CLK}" "${BACKUP_STATE_FILE}"
    if [ "${PWR_LEVEL}" == "manual" ]; then
      echo "FORCE_POWER_PROFILE: 0" >> "${BACKUP_STATE_FILE}"
    fi
    echo "FORCE_PERF_LEVEL: $(cat "${PWR_LEVEL}")" >> "${BACKUP_STATE_FILE}"
    echo "FORCE_POWER_CAP: $(cat "${PWR_CAP_FILE}")" >> "${BACKUP_STATE_FILE}"
    echo "Writen initial backup states to ${BACKUP_STATE_FILE}"
  else
    echo "Won't write initial state to ${BACKUP_STATE_FILE}, it already exists."
  fi
}

function restore_states() {
  if [ -f "${BACKUP_STATE_FILE}" ]; then
    echo "Restoring all states to the initial defaults from ${BACKUP_STATE_FILE}"
    USER_STATE_FILE=${BACKUP_STATE_FILE}
  else
    echo "Cant't access initial defaults at ${BACKUP_STATE_FILE}, aborting."
    exit 2
  fi
}

function clear_vars() {
    unset SCLK MCLK
    unset VDDC_CURVE VDDGFX_OFFSET
    unset sys_device sys_hwmons
}

check_ppfeaturemask

shopt -s extglob
for USER_STATE_FILE in ${USER_STATES_PATH}*.+(card?|card??|pci:????:??:??.?); do
  BACKUP_STATE_FILE=/tmp/${USER_STATE_FILE##*/}.initial
  if [[ "${USER_STATE_FILE}" == *"pci"* ]]; then
    STATE_PCIDEV=${USER_STATE_FILE#*.}
    for CARD in /sys/class/drm/card{?,??}; do
      LINK=$(readlink $CARD)
      if [ "${LINK}" ]; then
        CARD_PCDEV=${LINK%*/drm/card*}
        if [[ "${CARD_PCDEV##*/}" == "${STATE_PCIDEV#*:}" ]]; then
          sys_device="${CARD}"/device
        fi
      fi
    done
  else
    sys_device=/sys/class/drm/${USER_STATE_FILE##*.}/device
  fi
  sys_hwmons=( "${sys_device}"/hwmon/hwmon[0-9] )

  SYS_PP_OD_CLK=${sys_device}/pp_od_clk_voltage
  PWR_LEVEL=${sys_device}/power_dpm_force_performance_level
  PWR_PROFILE=${sys_device}/pp_power_profile_mode
  SYS_DPM_SCLK=${sys_device}/pp_dpm_sclk
  SYS_DPM_MCLK=${sys_device}/pp_dpm_mclk
  PWR_CAP_FILE=${sys_hwmons[0]}/power1_cap

  if [ -f "${SYS_PP_OD_CLK}" ]; then
    if [ "${RESTORE}" == "true" ]; then
      restore_states
    else
      backup_states
    fi
    echo "Detecting the state values at ${SYS_PP_OD_CLK}:"
    parse_states "${SYS_PP_OD_CLK}"
    echo "Verifying user state values at ${USER_STATE_FILE}:"
    parse_states "${USER_STATE_FILE}"
    echo "Committing custom states to ${SYS_PP_OD_CLK}:"
    set_custom_states
    echo "  Done"
    clear_vars
  else
    echo "WARNING: ${SYS_PP_OD_CLK} does not exist, skipping!"
  fi
done