#!/usr/bin/env bash
set -euo pipefail

# elevate once, then never again
if (( EUID != 0 )); then
  exec sudo -p "Password: " "$0" "$@"
fi

PROG=$(basename "$0")

usage() {
  cat <<EOF
Usage: $PROG N|--reset
  N        Number of logical CPUs to keep online (integer >= 1)
  --reset  Bring all logical CPUs online
EOF
  exit 1
}

get_total_cpus() { nproc --all; }

set_cpu_online() {
  local cpu=$1 val=$2
  local f="/sys/devices/system/cpu/cpu${cpu}/online"
  [[ -e $f ]] || return 0
  local cur
  cur=$(cat "$f" 2>/dev/null || echo "")
  [[ $cur == "$val" ]] && return 0
  printf '%s\n' "$val" >"$f" 2>/dev/null || \
    echo "Warning: failed to set cpu${cpu} online=$val" >&2
}

reset_cpus() {
  local total=$1
  for ((i = 0; i < total; i++)); do
    set_cpu_online "$i" 1
  done
}

apply_target() {
  local target=$1 total=$2
  for ((i = 0; i < total; i++)); do
    if ((i < target)); then
      set_cpu_online "$i" 1
    elif ((i > 0)); then
      set_cpu_online "$i" 0
    fi
  done
}

list_online() {
  local total=$1
  local online=()
  for ((i = 0; i < total; i++)); do
    local f="/sys/devices/system/cpu/cpu${i}/online"
    if [[ -e $f ]]; then
      [[ $(cat "$f") == 1 ]] && online+=("$i")
    else
      online+=("$i")
    fi
  done
  echo "Online CPUs: ${online[*]}"
}

main() {
  local total
  total=$(get_total_cpus)

  if [[ $# -ne 1 ]]; then
    usage
  fi

  case $1 in
    --reset)
      reset_cpus "$total"
      echo "Done."
      ;;
    [0-9]*)
      local target=$1
      if ((target <= 0)); then
        echo "Warning: minimum of 1 CPU required; using 1." >&2
        target=1
      fi
      if ((target >= total)); then
        reset_cpus "$total"
        echo "Done."
      else
        apply_target "$target" "$total"
        list_online "$total"
      fi
      ;;
    *)
      usage
      ;;
  esac
}

main "$@"