#!/usr/bin/env bash
#
# power.sh - Universal CPU Power & Multi-Core Temperature Monitor

INTERVAL=1.0 # Refresh rate in seconds

cleanup() {
    tput cnorm 2>/dev/null
    echo -e "\nExiting monitor."
    exit 0
}

trap cleanup SIGINT SIGTERM

if [[ $EUID -ne 0 ]]; then
    echo "Error: This script must be run as root to access CPU RAPL counters."
    exit 1
fi

if ! command -v perf &>/dev/null; then
    echo "Error: 'perf' command not found. Install with: apt install linux-perf"
    exit 1
fi

VENDOR_ID=$(grep -m 1 'vendor_id' /proc/cpuinfo | awk '{print $3}')

# --- Detect RAPL Power Events ---
EVENTS=()
PERF_EVENTS_LIST=$(perf list 2>/dev/null)

if echo "$PERF_EVENTS_LIST" | grep -q "power/energy-pkg/"; then
    EVENTS+=("-e" "power/energy-pkg/")
fi

if echo "$PERF_EVENTS_LIST" | grep -q "power_core/energy-core/"; then
    EVENTS+=("-e" "power_core/energy-core/")
elif echo "$PERF_EVENTS_LIST" | grep -q "power/energy-cores/"; then
    EVENTS+=("-e" "power/energy-cores/")
fi

if echo "$PERF_EVENTS_LIST" | grep -q "power/energy-ram/"; then
    EVENTS+=("-e" "power/energy-ram/")
fi

if [[ ${#EVENTS[@]} -eq 0 ]]; then
    echo "Error: No compatible RAPL energy events found on this system."
    exit 1
fi

# --- Map Thermal Sensors ---
# Arrays to hold ordered (label, file_path) pairs
CORE_LABELS=()
CORE_PATHS=()
PKG_LABEL=""
PKG_PATH=""

for hwmon in /sys/class/hwmon/hwmon*; do
    if [[ -f "$hwmon/name" ]]; then
        DRIVER_NAME=$(cat "$hwmon/name" 2>/dev/null)

        # Intel Sensors (coretemp)
        if [[ "$DRIVER_NAME" == "coretemp" ]]; then
            # Sort label files numerically (temp1_label, temp2_label...)
            for label_file in $(ls -v "$hwmon"/temp*_label 2>/dev/null); do
                if [[ -f "$label_file" ]]; then
                    LABEL=$(cat "$label_file" 2>/dev/null)
                    INPUT_FILE="${label_file%_label}_input"

                    if [[ "$LABEL" =~ Package ]]; then
                        PKG_LABEL="$LABEL"
                        PKG_PATH="$INPUT_FILE"
                    elif [[ "$LABEL" =~ Core ]]; then
                        CORE_LABELS+=("$LABEL")
                        CORE_PATHS+=("$INPUT_FILE")
                    fi
                fi
            done
            break

        # AMD Sensors (k10temp / zenpower)
        elif [[ "$DRIVER_NAME" == "k10temp" || "$DRIVER_NAME" == "zenpower" ]]; then
            for label_file in "$hwmon"/temp*_label; do
                if [[ -f "$label_file" ]]; then
                    LABEL=$(cat "$label_file" 2>/dev/null)
                    INPUT_FILE="${label_file%_label}_input"
                    if [[ "$LABEL" == "Tctl" ]]; then
                        PKG_LABEL="Tctl"
                        PKG_PATH="$INPUT_FILE"
                    elif [[ "$LABEL" == "Tccd1" ]]; then
                        CORE_LABELS+=("Tccd1")
                        CORE_PATHS+=("$INPUT_FILE")
                    fi
                fi
            done
            break
        fi
    fi
done

# Read temperature in Celsius
read_temp() {
    local sensor_path="$1"
    if [[ -n "$sensor_path" && -f "$sensor_path" ]]; then
        local raw
        raw=$(cat "$sensor_path" 2>/dev/null || echo 0)
        awk -v raw="$raw" 'BEGIN { printf "%.1f", raw / 1000 }'
    else
        echo "N/A"
    fi
}

# Colorize temperature values
format_temp() {
    local temp="$1"
    if [[ "$temp" == "N/A" ]]; then
        echo -e "\e[1;30mN/A\e[0m"
        return
    fi

    if awk "BEGIN {exit !($temp >= 80)}"; then
        printf "\e[1;31m%5.1f°C\e[0m" "$temp" # Red (>= 80C)
    elif awk "BEGIN {exit !($temp >= 65 && $temp < 80)}"; then
        printf "\e[1;33m%5.1f°C\e[0m" "$temp" # Yellow (65C - 80C)
    else
        printf "\e[1;32m%5.1f°C\e[0m" "$temp" # Green (< 65C)
    fi
}

tput civis 2>/dev/null
clear

while true; do
    tput cup 0 0

    # Sample RAPL energy counters
    PERF_OUT=$(perf stat "${EVENTS[@]}" -a sleep "$INTERVAL" 2>&1)

    # Extract Power Readings
    PKG_WATTS=$(echo "$PERF_OUT" | awk '/energy-pkg/ {print $1}' | tr -d ',')
    CORE_WATTS=$(echo "$PERF_OUT" | awk '/energy-core/ {print $1}' | tr -d ',')
    if [[ -z "$CORE_WATTS" ]]; then
        CORE_WATTS=$(echo "$PERF_OUT" | awk '/energy-cores/ {print $1}' | tr -d ',')
    fi
    DRAM_WATTS=$(echo "$PERF_OUT" | awk '/energy-ram/ {print $1}' | tr -d ',')

    PKG_WATTS=${PKG_WATTS:-"0.00"}
    CORE_WATTS=${CORE_WATTS:-"0.00"}

    # --- UI Rendering ---
    echo "=================================================="
    echo "        CPU Power & Temperature Monitor           "
    echo "=================================================="
    echo " Hostname      : $(hostname)"
    echo " CPU Vendor    : $VENDOR_ID"
    echo " Refresh Rate  : ${INTERVAL}s"
    echo "--------------------------------------------------"

    # Render Package / Tctl Temperature
    if [[ -n "$PKG_PATH" ]]; then
        PKG_TEMP=$(read_temp "$PKG_PATH")
        printf " Overall Package Temp (%-12s): %b\n" "$PKG_LABEL" "$(format_temp "$PKG_TEMP")"
    fi

    # Render Per-Core Temperatures (Formatted in a 2-column grid)
    if [[ ${#CORE_PATHS[@]} -gt 0 ]]; then
        echo "--------------------------------------------------"
        echo " Per-Core Temperatures:"

        col=0
        for i in "${!CORE_PATHS[@]}"; do
            c_label="${CORE_LABELS[$i]}"
            c_temp=$(read_temp "${CORE_PATHS[$i]}")

            printf "   %-8s: %b" "$c_label" "$(format_temp "$c_temp")"

            ((col++))
            if [[ $col -eq 2 ]]; then
                echo ""
                col=0
            else
                echo -n "   |"
            fi
        done
        [[ $col -ne 0 ]] && echo ""
    fi

    echo "--------------------------------------------------"
    printf " Total Package Power : \e[1;32m%7.2f Watts\e[0m\n" "$PKG_WATTS"
    printf " Cores Only Power   : \e[1;36m%7.2f Watts\e[0m\n" "$CORE_WATTS"
    if [[ -n "$DRAM_WATTS" ]]; then
        printf " DRAM Memory Power  : \e[1;35m%7.2f Watts\e[0m\n" "$DRAM_WATTS"
    fi
    echo "=================================================="
    echo "Press [CTRL+C] to exit."
done