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

CONFIG=/etc/t2-cpu-control.conf
STATE_DIR=/run/t2-cpu-control
DEFAULTS=$STATE_DIR/firmware-defaults
PERSISTENT_STATE_DIR=/var/lib/t2-cpu-control
THERMALD_STATE=$PERSISTENT_STATE_DIR/thermald-state
NO_TURBO=/sys/devices/system/cpu/intel_pstate/no_turbo
CPUFREQ=/sys/devices/system/cpu/cpufreq

fail() { printf 't2-cpu-control: %s\n' "$*" >&2; exit 1; }
zone() {
    local z
    for z in /sys/class/powercap/intel-rapl:*; do
        [[ -r "$z/name" ]] || continue
        [[ $(<"$z/name") == package-* ]] && { printf '%s' "$z"; return; }
    done
    fail "Intel package RAPL is unavailable"
}
constraint() {
    local z=$1 name=$2 f
    for f in "$z"/constraint_*_name; do
        [[ -r "$f" && $(<"$f") == "$name" ]] && { printf '%s' "${f%_name}"; return; }
    done
    fail "RAPL constraint '$name' is unavailable"
}
capture_defaults() {
    local z long short no_turbo=0 target tcc_offset max_freq_khz
    mkdir -p "$STATE_DIR"
    if [[ -e $DEFAULTS ]]; then
        if ! grep -q '^tcc_offset=' "$DEFAULTS"; then
            modprobe msr
            target=$((16#$(rdmsr -p 0 0x1a2)))
            printf 'tcc_offset=%s\ntcc_msr=%x\n' "$(((target >> 24) & 0x3f))" "$target" >>"$DEFAULTS"
        fi
        if ! grep -q '^max_freq_khz=' "$DEFAULTS"; then
            max_freq_khz=$(current_max_freq) || max_freq_khz=
            [[ -z $max_freq_khz ]] || printf 'max_freq_khz=%s\n' "$max_freq_khz" >>"$DEFAULTS"
        fi
        return
    fi
    z=$(zone); long=$(constraint "$z" long_term); short=$(constraint "$z" short_term)
    [[ -r $NO_TURBO ]] && no_turbo=$(<"$NO_TURBO")
    modprobe msr
    target=$((16#$(rdmsr -p 0 0x1a2)))
    tcc_offset=$(((target >> 24) & 0x3f))
    max_freq_khz=$(current_max_freq) || max_freq_khz=
    printf 'pl1_uw=%s\npl2_uw=%s\nno_turbo=%s\ntcc_offset=%s\ntcc_msr=%x\nmax_freq_khz=%s\n' \
        "$(<"${long}_power_limit_uw")" "$(<"${short}_power_limit_uw")" "$no_turbo" "$tcc_offset" "$target" "$max_freq_khz" >"$DEFAULTS"
}
write_limits() {
    local pl1=$1 pl2=$2 z long short
    [[ $pl1 =~ ^[0-9]+$ && $pl2 =~ ^[0-9]+$ ]] || fail "limits must be integer microwatts"
    ((pl1 >= 1000000 && pl2 >= pl1)) || fail "invalid PL1/PL2 limits"
    z=$(zone); long=$(constraint "$z" long_term); short=$(constraint "$z" short_term)
    printf '%s\n' "$pl2" >"${short}_power_limit_uw"
    printf '%s\n' "$pl1" >"${long}_power_limit_uw"
}
write_no_turbo() {
    local disabled=$1
    [[ $disabled == 0 || $disabled == 1 ]] || fail "no_turbo must be 0 or 1"
    [[ -w $NO_TURBO ]] || fail "intel_pstate turbo control is unavailable"
    printf '%s\n' "$disabled" >"$NO_TURBO"
    [[ $(<"$NO_TURBO") == "$disabled" ]] || fail "intel_pstate turbo state verification failed"
}
current_max_freq() {
    local f value= candidate
    for f in "$CPUFREQ"/policy*/scaling_max_freq; do
        [[ -r $f ]] || continue
        candidate=$(<"$f")
        [[ -n $value && $candidate -ge $value ]] || value=$candidate
    done
    [[ -n $value ]] || return 1
    printf '%s' "$value"
}
normalize_max_freq() {
    local value=$1
    [[ $value =~ ^[0-9]+$ ]] || fail "maximum CPU frequency must be integer kHz"
    printf '%s' "$(( ((value + 50000) / 100000) * 100000 ))"
}
write_max_freq() {
    local requested=$1 f minimum maximum readback index=0 policy_count=0
    local -a paths=() previous=()
    [[ $requested =~ ^[0-9]+$ ]] || fail "maximum CPU frequency must be integer kHz"
    ((requested % 100000 == 0)) || fail "maximum CPU frequency must use 100 MHz steps"
    for f in "$CPUFREQ"/policy*/scaling_max_freq; do
        [[ -e $f ]] || continue
        policy_count=$((policy_count + 1))
        [[ -r $f && -w $f ]] || fail "not all CPU frequency policies are writable"
        minimum=$(<"${f%/scaling_max_freq}/cpuinfo_min_freq")
        maximum=$(<"${f%/scaling_max_freq}/cpuinfo_max_freq")
        ((requested >= minimum && requested <= maximum)) || fail "maximum CPU frequency is outside policy range"
        paths+=("$f")
        previous+=("$(<"$f")")
    done
    ((policy_count > 0 && ${#paths[@]} == policy_count)) || fail "writable CPU frequency limits are unavailable"
    for f in "${paths[@]}"; do
        if ! printf '%s\n' "$requested" >"$f"; then
            for ((index=0; index<${#paths[@]}; index++)); do
                printf '%s\n' "${previous[index]}" >"${paths[index]}" 2>/dev/null || true
            done
            fail "cannot set maximum CPU frequency"
        fi
    done
    for f in "${paths[@]}"; do
        readback=$(<"$f")
        [[ $readback == "$requested" ]] || fail "maximum CPU frequency verification failed"
    done
}
write_tcc_offset() {
    local offset=$1 cpu core value new_value readback
    [[ $offset =~ ^[0-9]+$ ]] && ((offset <= 63)) || fail "TCC offset must be between 0 and 63 C"
    modprobe msr
    for cpu in /dev/cpu/[0-9]*; do
        [[ -d $cpu ]] || continue
        core=${cpu##*/}
        value=$((16#$(rdmsr -p "$core" 0x1a2))) || fail "cannot read MSR_TEMPERATURE_TARGET on CPU $core"
        new_value=$(((value & ~(0x3f << 24)) | (offset << 24)))
        printf -v new_value '0x%x' "$new_value"
        wrmsr -p "$core" 0x1a2 "$new_value" || fail "cannot set CPU thermal target on CPU $core"
        readback=$((16#$(rdmsr -p "$core" 0x1a2))) || fail "cannot verify CPU thermal target on CPU $core"
        ((((readback >> 24) & 0x3f) == offset)) || \
            fail "CPU thermal target verification failed on CPU $core"
        (((readback & ~(0x3f << 24)) == (value & ~(0x3f << 24)))) || \
            fail "MSR_TEMPERATURE_TARGET fields changed unexpectedly on CPU $core"
    done
}
write_tcc_target() {
    local target_c=$1 value tjmax offset
    [[ $target_c =~ ^[0-9]+$ ]] || fail "CPU thermal target must be an integer temperature"
    modprobe msr
    value=$((16#$(rdmsr -p 0 0x1a2))) || fail "cannot read MSR_TEMPERATURE_TARGET"
    tjmax=$(((value >> 16) & 0xff))
    ((target_c >= 60 && target_c <= tjmax && tjmax - target_c <= 63)) || \
        fail "CPU thermal target must be between $((tjmax - 63 > 60 ? tjmax - 63 : 60)) and $tjmax C"
    offset=$((tjmax - target_c))
    write_tcc_offset "$offset"
}
tcc_writable() {
    local target platform_info
    modprobe msr
    target=$((16#$(rdmsr -p 0 0x1a2))) || return 1
    platform_info=$((16#$(rdmsr -p 0 0xce))) || return 1
    (( ((target >> 31) & 1) == 0 && ((platform_info >> 30) & 1) == 1 ))
}
current_writable_tcc_target() {
    if tcc_writable; then
        current_tcc_target
    else
        printf '0'
    fi
}
write_tcc_msr() {
    local value=$1 cpu core readback
    [[ $value =~ ^[0-9a-fA-F]+$ ]] || fail "invalid saved MSR_TEMPERATURE_TARGET value"
    modprobe msr
    for cpu in /dev/cpu/[0-9]*; do
        [[ -d $cpu ]] || continue
        core=${cpu##*/}
        wrmsr -p "$core" 0x1a2 "0x$value" || fail "cannot restore MSR_TEMPERATURE_TARGET on CPU $core"
        readback=$(rdmsr -p "$core" 0x1a2) || fail "cannot verify MSR_TEMPERATURE_TARGET on CPU $core"
        [[ ${readback,,} == ${value,,} ]] || fail "MSR_TEMPERATURE_TARGET restore verification failed on CPU $core"
    done
}
current_tcc_target() {
    local value tjmax offset
    modprobe msr
    value=$((16#$(rdmsr -p 0 0x1a2))) || fail "cannot read MSR_TEMPERATURE_TARGET"
    tjmax=$(((value >> 16) & 0xff)); offset=$(((value >> 24) & 0x3f))
    printf '%s' "$((tjmax - offset))"
}
save_config() {
    local pl1=$1 pl2=$2 enabled=$3 no_turbo=$4 tcc_target_c=$5 max_freq_khz=$6
    printf 'pl1_uw=%s\npl2_uw=%s\nenabled=%s\nno_turbo=%s\ntcc_target_c=%s\nmax_freq_khz=%s\n' \
        "$pl1" "$pl2" "$enabled" "$no_turbo" "$tcc_target_c" "$max_freq_khz" >"$CONFIG"
    chmod 0644 "$CONFIG"
}
capture_thermald_state() {
    local enabled=0 active=0
    [[ -e "$THERMALD_STATE" ]] && return
    if systemctl is-enabled --quiet thermald.service; then
        enabled=1
    fi
    if systemctl is-active --quiet thermald.service; then
        active=1
    fi
    mkdir -p "$PERSISTENT_STATE_DIR"
    printf 'enabled=%s\nactive=%s\n' "$enabled" "$active" >"$THERMALD_STATE"
}
take_rapl_ownership() {
    local persist=$1 enabled=0 active=0
    capture_thermald_state
    # shellcheck disable=SC1090
    source "$THERMALD_STATE"
    if systemctl is-active --quiet thermald.service; then
        systemctl stop thermald.service
    fi
    if [[ $persist == 1 ]]; then
        if systemctl is-enabled --quiet thermald.service; then
            systemctl disable thermald.service
        fi
    elif [[ $enabled == 1 ]] && ! systemctl is-enabled --quiet thermald.service; then
        systemctl enable thermald.service
    fi
}
restore_thermald_state() {
    local enabled=0 active=0
    [[ -r "$THERMALD_STATE" ]] || return
    # shellcheck disable=SC1090
    source "$THERMALD_STATE"
    if [[ $enabled == 1 ]]; then
        systemctl enable thermald.service
    elif systemctl is-enabled --quiet thermald.service; then
        systemctl disable thermald.service
    fi
    if [[ $active == 1 ]]; then
        systemctl start thermald.service
    elif systemctl is-active --quiet thermald.service; then
        systemctl stop thermald.service
    fi
    rm -f "$THERMALD_STATE"
}
apply_saved() {
    local pl1_uw= pl2_uw= enabled=0 no_turbo=0 tcc_target_c= max_freq_khz= z long short active_pl1 active_pl2
    capture_defaults
    [[ -r "$CONFIG" ]] || return 0
    # shellcheck disable=SC1090
    source "$CONFIG"
    # The thermal target is an explicit persistent setting.  The enabled flag
    # only controls ownership and reapplication of the RAPL-related settings.
    [[ -z $tcc_target_c || $tcc_target_c == 0 ]] || write_tcc_target "$tcc_target_c"
    [[ $enabled == 1 ]] || return 0
    if [[ -n $max_freq_khz && $max_freq_khz != 0 ]]; then
        max_freq_khz=$(normalize_max_freq "$max_freq_khz")
    fi
    write_limits "$pl1_uw" "$pl2_uw"
    write_no_turbo "$no_turbo"
    [[ -z $max_freq_khz || $max_freq_khz == 0 ]] || write_max_freq "$max_freq_khz"
    z=$(zone); long=$(constraint "$z" long_term); short=$(constraint "$z" short_term)
    active_pl1=$(<"${long}_power_limit_uw")
    active_pl2=$(<"${short}_power_limit_uw")
    save_config "$pl1_uw" "$pl2_uw" "$enabled" "$no_turbo" "${tcc_target_c:-0}" "${max_freq_khz:-0}"
    printf 't2-cpu-control: requested PL1=%s PL2=%s no_turbo=%s TCC=%s max_freq=%s; read back PL1=%s PL2=%s no_turbo=%s TCC=%s max_freq=%s\n' \
        "$pl1_uw" "$pl2_uw" "$no_turbo" "${tcc_target_c:-unchanged}" "${max_freq_khz:-unchanged}" "$active_pl1" "$active_pl2" "$(<"$NO_TURBO")" "$(current_tcc_target)" "$(current_max_freq)"
}
fans_max() {
    local d f
    mkdir -p "$STATE_DIR"
    if systemctl is-active --quiet t2-fancontrol.service; then
        echo 1 >"$STATE_DIR/fan-service-was-active"
        systemctl stop t2-fancontrol.service
    fi
    for d in /sys/class/hwmon/hwmon*; do
        [[ -r "$d/name" ]] || continue
        case $(<"$d/name") in t2smc|macsmc) ;; *) continue ;; esac
        for f in "$d"/fan*_max; do
            [[ -r "$f" ]] || continue
            printf '%s\n' "$(<"$f")" >"${f%_max}_target"
        done
    done
}
fans_restore() {
    local d f
    for d in /sys/class/hwmon/hwmon*; do
        [[ -r "$d/name" ]] || continue
        case $(<"$d/name") in t2smc|macsmc) ;; *) continue ;; esac
        for f in "$d"/fan*_target; do [[ -e "$f" ]] && echo 0 >"$f"; done
    done
    if [[ -e "$STATE_DIR/fan-service-was-active" ]]; then
        systemctl start t2-fancontrol.service
        rm -f "$STATE_DIR/fan-service-was-active"
    fi
}

set_prochot_override() {
    local override=$1 cpu core value new_value readback expected
    [[ $override == 0 || $override == 1 ]] || fail "PROCHOT override state must be 0 or 1"
    modprobe msr
    expected=$((override == 0))
    for cpu in /dev/cpu/[0-9]*; do
        [[ -d $cpu ]] || continue
        core=${cpu##*/}
        value=$(rdmsr -p "$core" 0x1fc) || fail "cannot read MSR_POWER_CTL on CPU $core"
        value=$((16#$value))
        if [[ $override == 1 ]]; then
            new_value=$((value & ~1))
        else
            new_value=$((value | 1))
        fi
        printf -v new_value '%x' "$new_value"
        wrmsr -p "$core" 0x1fc "$new_value" || fail "cannot write MSR_POWER_CTL on CPU $core"
        readback=$(rdmsr -p "$core" 0x1fc) || fail "cannot verify MSR_POWER_CTL on CPU $core"
        readback=$((16#$readback))
        (( (readback & 1) == expected )) || fail "PROCHOT override verification failed on CPU $core"
    done
}

((EUID == 0)) || fail "root privileges are required"
case ${1:-} in
    apply)
        [[ $# == 7 || $# == 8 ]] || fail "usage: $0 apply <pl1_uw> <pl2_uw> <persist:0|1> <no_turbo:0|1> <tcc_target_c> <max_freq_khz> [prochot_override]"
        [[ $4 == 0 || $4 == 1 ]] || fail "persistence must be 0 or 1"
        capture_defaults
        take_rapl_ownership "$4"
        write_limits "$2" "$3"
        write_no_turbo "$5"
        [[ $6 == 0 ]] || write_tcc_target "$6"
        [[ $7 == 0 ]] || write_max_freq "$7"
        [[ $# == 7 ]] || set_prochot_override "$8"
        save_config "$2" "$3" "$4" "$5" "$6" "$7"
        ;;
    apply-saved) apply_saved ;;
    restore)
        [[ -r "$DEFAULTS" ]] || fail "firmware defaults were not captured during boot; reboot first"
        # shellcheck disable=SC1090
        source "$DEFAULTS"
        write_limits "$pl1_uw" "$pl2_uw"
        write_no_turbo "${no_turbo:-0}"
        [[ -z ${max_freq_khz:-} || ${max_freq_khz:-0} == 0 ]] || write_max_freq "$max_freq_khz"
        if tcc_writable; then
            if [[ -n ${tcc_msr:-} ]]; then
                write_tcc_msr "$tcc_msr"
            else
                write_tcc_offset "${tcc_offset:-0}"
            fi
        fi
        save_config "$pl1_uw" "$pl2_uw" 0 "${no_turbo:-0}" "$(current_writable_tcc_target)" "$(current_max_freq)"
        restore_thermald_state
        ;;
    disable-persistence)
        save_config 0 0 0 "$(<"$NO_TURBO")" "$(current_writable_tcc_target)" "$(current_max_freq)"
        restore_thermald_state
        ;;
    enable-persistence)
        capture_defaults
        z=$(zone); long=$(constraint "$z" long_term); short=$(constraint "$z" short_term)
        take_rapl_ownership 1
        save_config "$(<"${long}_power_limit_uw")" "$(<"${short}_power_limit_uw")" 1 \
            "$(<"$NO_TURBO")" "$(current_writable_tcc_target)" "$(current_max_freq)"
        ;;
    fans-max) fans_max ;;
    fans-restore) fans_restore ;;
    prochot-override)
        [[ $# == 2 ]] || fail "usage: $0 prochot-override <enabled:0|1>"
        set_prochot_override "$2"
        ;;
    *) fail "unknown command" ;;
esac
