#!/bin/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() {
  echo "  VDDC Curve state $1: ${2}, ${4}"
  VDDC_CURVE[${1}]="vc ${1} ${2%*M[Hh]z} ${4%*mV}"
}

function parse_states() {
  mapfile -t STATE_LINES < $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_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}
        ;;
      "VDDC: "*)
        echo "    VDDC voltage ${LINE##* }"
        MAX_VDDC=${LINE##* }
        MAX_VDDC=${MAX_VDDC%*mV}
        ;;
      [0-9]": "*)
        $state_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#* }
        ;;
      "#"*) ;;
      "") ;;
      *)
        let "LINE_NUMBER = ${LNNO} + 1"
        echo "  Unexpected value in ${1}:${LINE_NUMBER}"
        exit 2
        ;;
    esac
  done
  if [ "$1" == "${SYS_PP_OD_CLK}" ]; then
    POWER_LIMIT=$(cat $PWR_CAP_FILE)
    echo "  Curent power cap: ${POWER_LIMIT%*000000}W"
  fi
}

function set_custom_states() {
  if [ "${FORCE_LEVEL}" ]; then
    echo ${FORCE_LEVEL} > ${PWR_LEVEL}
  fi
  if [ "${FORCE_PROFILE}" ]; then
    echo ${FORCE_PROFILE} > ${PWR_PROFILE}
  fi
  for CSTATE in "${SCLK[@]}"; do
    echo ${CSTATE} > ${SYS_PP_OD_CLK}
  done
  for MSTATE in "${MCLK[@]}"; do
    echo ${MSTATE} > ${SYS_PP_OD_CLK}
  done
  for VDDC_CURVE_STATE in "${VDDC_CURVE[@]}"; do
    echo ${VDDC_CURVE_STATE} > ${SYS_PP_OD_CLK}
  done
  echo 'c' > ${SYS_PP_OD_CLK}
  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
}

check_ppfeaturemask

for USER_STATE_FILE in $(ls -1 ${USER_STATES_PATH}*.card*); do
  BACKUP_STATE_FILE=/tmp/${USER_STATE_FILE##*/}.initial
  SYS_PP_OD_CLK=/sys/class/drm/${USER_STATE_FILE##*.}/device/pp_od_clk_voltage
  PWR_LEVEL=/sys/class/drm/${USER_STATE_FILE##*.}/device/power_dpm_force_performance_level
  PWR_PROFILE=/sys/class/drm/${USER_STATE_FILE##*.}/device/pp_power_profile_mode
  SYS_DPM_SCLK=/sys/class/drm/${USER_STATE_FILE##*.}/device/pp_dpm_sclk
  SYS_DPM_MCLK=/sys/class/drm/${USER_STATE_FILE##*.}/device/pp_dpm_mclk
  HWMON_PATH=/sys/class/drm/${USER_STATE_FILE##*.}/device/hwmon
  HWMON_ID=$(ls -1 /sys/class/drm/${USER_STATE_FILE##*.}/device/hwmon)
  PWR_CAP_FILE=/sys/class/drm/${USER_STATE_FILE##*.}/device/hwmon/${HWMON_ID}/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"
  else
    echo "WARNING: ${SYS_PP_OD_CLK} does not exist, skipping!"
  fi
done