#!/bin/bash

set -e
set -o pipefail

PROGNAME="$(basename "$0")"

TEMP=`getopt --shell sh \
        -o h --long help,ht,noht,normal,full,verbose \
        -n "$PROGNAME" -- "$@"`

eval set -- "$TEMP"

last_cpu=$(cd /sys/devices/system/cpu; ls -d cpu[0-9]* | sed -e 's/^cpu//' | sort -n | tail -1)
prev_last_cpu=$(($last_cpu - 1))

usage() {

cat <<EOF
Usage: $PROGNAME [OPTIONS...]
  Get (no options) or set processors mode
Options:
  -h|--help: this help
  --verbose: more information
  --normal: disable the last core (reserved for the hypervisor)
  --full: enable the last core (shared with the hypervisor)
  --ht: enable hyperthread processors
  --noht: disable hyperthread processors
EOF

}

cpuonline() {
	local nb="$1";
	if [ -z "$2" ]; then
		cat /sys/devices/system/cpu/cpu$nb/online
	else
		echo "$2" > /sys/devices/system/cpu/cpu$nb/online
	fi
}

varcases() {
	local nb="$1" ; shift

	while [ $nb != 0 ]; do
		shift
		nb=$(( $nb - 1 ))
	done
	printf "%s" "$1"
}

setup() {
	HT_ENABLED="$(cpuonline 1)"
	FULL_ENABLED="$(cpuonline $(($last_cpu - 1)))"
}

check() {
	for n in $(seq 2 2 $(($last_cpu - 2))); do
		if [ $(cpuonline $n) != 1 ]; then
			echo 1>&1 "Strange: CPU$n is disabled"
		fi
	done

	for n in $(seq 1 2 $(varcases $FULL_ENABLED $(($last_cpu - 2)) $last_cpu)); do
		if [ $(cpuonline $n) != $HT_ENABLED ]; then
			echo 1>&2 "Strange: HT CPU$n is $(varcases $HT_ENABLED enabled disabled)"
		fi
	done
}

setup
HT_ENABLE=$HT_ENABLED
FULL_ENABLE=$FULL_ENABLED
VERBOSE=0

check

DO=get
while true ; do
        case "$1" in
        -h|--help) usage ; exit 0 ;;
	--ht) HT_ENABLE=1 ; DO=set ;;
	--noht) HT_ENABLE=0 ; DO=set ;;
	--normal) FULL_ENABLE=0 ; DO=set ;;
	--full) FULL_ENABLE=1 ; DO=set ;;
	--verbose) VERBOSE=1 ;;
        --) shift ; break ;;
        *) echo "Unknown option/argument '$1'" ; exit 1 ;;
        esac
        shift;
done

do_get() {
	echo "HT enabled: $(varcases $HT_ENABLED no yes)"
	echo "CPU mode: $(varcases $FULL_ENABLED "normal (last core disabled)" "full (last code enabled)")"
}

do_set() {
	if [ $FULL_ENABLE != $FULL_ENABLED ]; then
		cpuonline $((last_cpu - 1)) $FULL_ENABLE 
		if [ $FULL_ENABLE$HT_ENABLE = 11 ]; then
			cpuonline $last_cpu 1
		else
			cpuonline $last_cpu 0
		fi
	fi
	if [ $HT_ENABLE != $HT_ENABLED ]; then
		for n in $(seq 1 2 $(varcases $FULL_ENABLE $((last_cpu - 2)) $last_cpu)); do
			cpuonline $n $HT_ENABLE
		done
	fi
	if [ $VERBOSE = 1 ]; then
		setup
		check
		do_get
	fi
}

do_$DO

