#!/usr/bin/env bash
# cpucap - clamp scaling_max_freq on every cpufreq policy, per-policy hw-range safe.
#
#   cpucap 1.2GHz        cap every policy at 1.2 GHz (clamped into its own hw range)
#   cpucap 900MHz
#   cpucap 1200000       raw kHz also accepted
#   cpucap --reset       restore each policy to its cpuinfo_max_freq
#   cpucap --show        print current state (freq caps, package power, temp)
#   cpucap --power 12    also set RAPL PL1+PL2 to 12 W (see note at bottom)
#   cpucap --no-turbo    pin to base clock (2400 P / 1800 E); --turbo to undo
#
# Anything that writes needs root; it re-runs itself under sudo on its own.

set -euo pipefail

POLICIES=(/sys/devices/system/cpu/cpufreq/policy*)
RAPL=/sys/class/powercap/intel-rapl:0

die() { printf 'cpucap: %s\n' "$*" >&2; exit 1; }

# Re-exec under sudo rather than dying with a hint: this lives on the user's
# PATH via home.packages, so plain `sudo cpucap` would not resolve. CPUCAP_SELF
# is set by the nix wrapper to its own $out/bin path; "$0" is the fallback for
# running the file straight out of a checkout. Hardcoded sudo path because only
# /run/wrappers/bin's copy is setuid - the one in the store is not.
need_root() {
    if [[ $EUID -ne 0 ]]; then
        exec /run/wrappers/bin/sudo "${CPUCAP_SELF:-$0}" "$@"
    fi
}

# Accept 1.2GHz / 1200MHz / 1200000kHz / 1200000 -> kHz integer
to_khz() {
    local v="${1,,}" n
    case "$v" in
        *ghz) n="${v%ghz}"; awk -v n="$n" 'BEGIN{printf "%d", n * 1000000}' ;;
        *mhz) n="${v%mhz}"; awk -v n="$n" 'BEGIN{printf "%d", n * 1000}' ;;
        *khz) n="${v%khz}"; awk -v n="$n" 'BEGIN{printf "%d", n}' ;;
        *[!0-9]*) die "cannot parse frequency '$1' (use e.g. 1.2GHz, 900MHz, or raw kHz)" ;;
        *)    echo "$v" ;;
    esac
}

fmt_mhz() { awk -v k="$1" 'BEGIN{printf "%.0f MHz", k / 1000}'; }

show() {
    local p hw_min hw_max cur_min cur_max
    printf '%-9s %-11s %-11s %s\n' POLICY CAP MIN 'HW RANGE'
    for p in "${POLICIES[@]}"; do
        hw_min=$(< "$p/cpuinfo_min_freq"); hw_max=$(< "$p/cpuinfo_max_freq")
        cur_min=$(< "$p/scaling_min_freq"); cur_max=$(< "$p/scaling_max_freq")
        printf '%-9s %-11s %-11s %s - %s\n' \
            "${p##*/}" "$(fmt_mhz "$cur_max")" "$(fmt_mhz "$cur_min")" \
            "$(fmt_mhz "$hw_min")" "$(fmt_mhz "$hw_max")"
    done

    if [[ -r /sys/devices/system/cpu/intel_pstate/no_turbo ]]; then
        printf '\nintel_pstate: status=%s no_turbo=%s max_perf_pct=%s\n' \
            "$(< /sys/devices/system/cpu/intel_pstate/status)" \
            "$(< /sys/devices/system/cpu/intel_pstate/no_turbo)" \
            "$(< /sys/devices/system/cpu/intel_pstate/max_perf_pct)"
    fi

    if [[ -r $RAPL/constraint_0_power_limit_uw ]]; then
        printf 'RAPL package: PL1=%sW PL2=%sW\n' \
            "$(( $(< "$RAPL/constraint_0_power_limit_uw") / 1000000 ))" \
            "$(( $(< "$RAPL/constraint_1_power_limit_uw") / 1000000 ))"
    fi

    local t
    for t in /sys/class/hwmon/hwmon*; do
        [[ $(< "$t/name") == coretemp ]] || continue
        [[ -r $t/temp1_input ]] && printf 'package temp: %s C\n' "$(( $(< "$t/temp1_input") / 1000 ))"
    done
}

set_cap() {
    local target=$1 p hw_min hw_max want cur_min got
    printf '%-9s %-11s %s\n' POLICY REQUESTED APPLIED
    for p in "${POLICIES[@]}"; do
        hw_min=$(< "$p/cpuinfo_min_freq")
        hw_max=$(< "$p/cpuinfo_max_freq")

        # Clamp into this policy's own hw range. P-cores and E-cores differ.
        want=$target
        (( want < hw_min )) && want=$hw_min
        (( want > hw_max )) && want=$hw_max

        # scaling_min must not exceed the new max, or the write is rejected.
        cur_min=$(< "$p/scaling_min_freq")
        (( cur_min > want )) && echo "$want" > "$p/scaling_min_freq"

        echo "$want" > "$p/scaling_max_freq"

        # intel_pstate rounds to the nearest achievable P-state, so read back.
        got=$(< "$p/scaling_max_freq")
        printf '%-9s %-11s %s%s\n' "${p##*/}" "$(fmt_mhz "$want")" "$(fmt_mhz "$got")" \
            "$( (( got != want )) && echo '  (rounded)' )"
    done
}

reset_cap() {
    local p
    for p in "${POLICIES[@]}"; do
        echo "$(< "$p/cpuinfo_min_freq")" > "$p/scaling_min_freq"
        echo "$(< "$p/cpuinfo_max_freq")" > "$p/scaling_max_freq"
    done
    echo "restored all policies to their hardware maximum"
}

set_power() {
    local watts=$1 uw max
    if [[ -z $watts || $watts == *[!0-9]* ]]; then
        die "--power takes whole watts, got '$watts'"
    fi
    uw=$((watts * 1000000))
    [[ -w $RAPL/constraint_0_power_limit_uw ]] || die "RAPL not writable at $RAPL"
    max=$(< "$RAPL/constraint_0_max_power_uw")
    (( max > 0 && uw > max )) && die "requested ${watts}W exceeds RAPL max $((max / 1000000))W"
    echo "$uw" > "$RAPL/constraint_0_power_limit_uw"   # PL1, sustained
    echo "$uw" > "$RAPL/constraint_1_power_limit_uw"   # PL2, burst
    printf 'RAPL package: PL1=PL2=%sW\n' \
        "$(( $(< "$RAPL/constraint_0_power_limit_uw") / 1000000 ))"
}

set_turbo() {
    local off=$1 f=/sys/devices/system/cpu/intel_pstate/no_turbo
    [[ -w $f ]] || die "no_turbo not writable (intel_pstate not in active mode?)"
    echo "$off" > "$f"
    # Turbo-off pins each policy to base_frequency; an existing lower cap still wins.
    if [[ $off == 1 ]]; then
        echo "turbo disabled - ceiling is now base clock, or your cap if it is lower"
    else
        echo "turbo enabled"
    fi
}

[[ ${#POLICIES[@]} -gt 0 && -d ${POLICIES[0]} ]] || die "no cpufreq policies found"

case "${1-}" in
    ''|--show|-s)  show ;;
    --reset|-r)    need_root "$@"; reset_cap; echo; show ;;
    --power|-p)    need_root "$@"; [[ -n ${2-} ]] || die "--power needs a wattage"
                   set_power "$2" ;;
    --no-turbo)    need_root "$@"; set_turbo 1; echo; show ;;
    --turbo)       need_root "$@"; set_turbo 0; echo; show ;;
    -h|--help)     sed -n '2,12p' "$0" | sed 's/^# \?//' ;;
    -*)            die "unknown option '$1' (see --help)" ;;
    *)             need_root "$@"; set_cap "$(to_khz "$1")"
                   if [[ ${2-} == --power || ${2-} == -p ]]; then set_power "$3"; fi
                   echo; show ;;
esac