#! /usr/bin/bash

# Global vars/settings
command_name=$(basename $0)
rapl_path="/sys/class/powercap/intel-rapl:0"
pl1_path="${rapl_path}/constraint_0_power_limit_uw"
pl2_path="${rapl_path}/constraint_1_power_limit_uw"
script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
gui_name="GPD TDP Manager"
gui_options=""

# Set up help text
help_text="${command_name}: GDP Win 3 TDP management script

This script checks and sets the TDP using intel_rapl

Usage:
  ${command_name} COMMAND [ARGUMENTS]

Command:
  check [ARGS]     Checks current TDP and prints it in watts
  c [ARGS]         Shortened version of check

  set WATTS [ARGS] Sets TDP to the requested number in watts
  s WATTS [ARGS]   Shortened version of set
  
  gui              Runs graphical user interface

  help             Prints this help text
  COMMAND --help   Prints help for specified command"

check_help_text="${command_name} check [ARG]
Short command: c

Example:
  ${command_name} check --detail
  ${command_name} c -d

Arguments:
  --detail, -d     Prints PL2 as well as PL1
  --help           Prints this help text"

set_help_text="${command_name} set WATT [ARGUMENTS]
Short command: s

Example:
  ${command_name} set 10 --detail --same
  ${command_name} s 10 -d -s

Arguments:
  --detail, -d     Prints PL2 as well as PL1
  --same, -s       PL2 will be set to the same as PL1 rather than 2W higher
  --help           Prints this help text"
  
gui_help_text="${command_name} gui

No help is available for the GUI.
If it does not work, ensure the \"zenity\" command is installed and available"

print_help () {
  case $1 in
    "check")
      echo "$check_help_text"
      ;;
    "set")
      echo "$set_help_text"
      ;;
    "gui")
      echo "$gui_help_text"
      ;;
    *)
      echo "$help_text"
      ;;
  esac
}

print_unknown () {
  echo "Unknown command or incorrect arguments.
  Try \"${command_name} help\" or \"${command_name} COMMAND --help\""
}
# End of help text

# Conversion tools
uw_to_w () {
  if ! [[ $1 =~ ^-?[0-9]+$ ]]; then
    exit
  else
    echo $(expr $1 / 1000000)
  fi
}

w_to_uw () {
  if ! [[ $1 =~ ^-?[0-9]+$ ]]; then
    exit
  else
    echo $(expr $1 '*' 1000000)
  fi
}

# Common checks/tasks
is_detailed () {
  if [[ "$@" == *"--detail"* ]] || [[ "$@" == *"-d"* ]]; then
    exit 0
  else
    exit 1
  fi
}

is_help () {
  if [[ "$@" == *"--help"* ]] || [[ "$@" == *"-h"* ]]; then
    exit 0
  else
    exit 1
  fi
}

is_number () {
  if ! [[ $1 =~ ^-?[0-9]+$ ]]; then
    exit 1
  else
    exit 0
  fi
}

gui_msg () {
  zenity --info --title="$gui_name" $gui_options --text="$1" --no-wrap
}

gui_read () {
  zenity --entry --title="$gui_name" $gui_options --text="$1:"
}

gui_ask () {
  if $(zenity --question --title="$gui_name" $gui_options --text="$1" --no-wrap); then
    exit 0
  else
    exit 1
  fi
}

set_pl () {
  if [ $1 -eq 1 ] || [ $1 -eq 2 ] || [ ! -z $2 ]; then
    if [ $1 -eq 1 ]; then
      pl_path=$pl1_path
    elif [ $1 -eq 2 ]; then
      pl_path=$pl2_path
    fi
    if [ "$EUID" -eq 0 ]; then
      echo $(w_to_uw $2) > $pl_path
    elif [ $gui_used == 1 ]; then
      sudo -A bash -c "echo $(w_to_uw $2) > ${pl_path}"
    else
      sudo bash -c "echo $(w_to_uw $2) > ${pl_path}"
    fi
  else
    exit 1
  fi
}

get_pl () {
  if [ $1 -eq 1 ]; then
    cat $pl1_path
  elif [ $1 -eq 2 ]; then
    cat $pl2_path
  fi
}

# Retrieves current TDP and prints it
check_tdp () {
  if [ -z "$1" ] || ( is_detailed $@ ); then
    local pl1=$(uw_to_w $(get_pl 1))
    local pl2=$(uw_to_w $(get_pl 2))

    echo "PL1 is ${pl1}W" # Placeholder
    if ( is_detailed $@ ); then
      echo "PL2 is ${pl2}W" # Placeholder
    fi

  elif ( is_help $@ ); then
    print_help "check"

  else
    print_unknown
  fi
}

# Sets PL1 to number provided as first argument, and PL2 2W higher
set_tdp () {
  if [ -z "$1" ]; then
    echo "Please specify wattage"
    exit
  elif ( is_help $@ ); then
    print_help "set"
    exit
  elif ! (is_number $1); then
    echo "TDP is not a number or argument unknown!)"
    print_unknown
    exit
  fi

  if [ $1 -lt 5 ] || [ $1 -gt 30 ]; then
    echo "TDP too high or low, should be between 5W and 30W"
    echo "This is a sanity limit to prevent you from throttling to a near unusable state"
  else
    #PL1
    local watts=$1
    set_pl 1 $watts

    #PL2
    local watts2=$(expr $watts + 2)
    if [[ "$@" == *"--same"* ]] || [[ "$@" == *"-s"* ]]; then
      watts2=$watts
    fi
    set_pl 2 $watts2


    echo "PL1 is now ${watts}W (Long-term)"
    if ( is_detailed $@ ); then
      echo "PL2 is now ${watts2}W (Short-term)"
    fi
  fi
}

# Basic GUI using Zenity
gui_handler() {
  if [ -z $(which zenity 2>/dev/null) ]; then
    echo "Zenity is not available, GUI will not work until it is installed"
  elif (is_help $@); then
    print_help "gui"
  else
    export SUDO_ASKPASS="${script_dir}/gpd-sudo-prompt"
    export gui_used=1
    while : ; do
      # Contains some ugly hacks to "widen" the entries for the small GPD screen, might tweak later
      gui_action=$(zenity --list --title="${gui_name}" --width=400 --height=300 $gui_options --text="Chose action to perform:" --hide-column=1 --column="" --column="Action"\
        check "`printf "\n  Check current TDP\n "`" \
        set "`printf "\n  Set new TDP (requires sudo password)\n "`" \
        exit "`printf "\n  Exit TDP Manager\n "`")
      
      case $gui_action in
      
        "check")
        gui_msg "$($command_name check --detail)"
        ;;
      
        "set")
        gui_tdp=$(gui_read "Please enter TDP")
        
        if ! [ -z $gui_tdp ]; then
          if $(gui_ask "Should PL2 be set to the same wattage?\nIf unsure, answer No."); then
            gui_msg "$(tdp set $gui_tdp --same --detail)"
          else
            gui_msg "$(tdp set $gui_tdp --detail)"
          fi
          gui_tdp=""
        fi
        ;;
      
        *)
        break
        ;;
      esac
    done
  fi
}

# Command handler
case $1 in

  "check" | "c")
    check_tdp "${@:2}"
    ;;

  "set" | "s")
    set_tdp "${@:2}"
    ;;

  "gui") # Pass to GUI handler
    gui_handler "${@:2}"
    ;;

  "help")
    print_help
    ;;

  *)
    print_unknown
    ;;
esac