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

# vcgencmd get_throttled
# vcgencmd measure_temp

CONFIG_FILE="/etc/rpi-throttle-monitor.conf"
TEST_MODE=0
TEST_VALUES=()
TEST_SEND_NOTIFICATIONS=0
STOP_REASON=""
CLEANUP_DONE=0

write_service_log() {
    local msg="$1"
    local ts
    ts="$(date -Is)"

    if (( TEST_MODE == 0 )); then
        logger -t rpi-throttle-monitor "$msg"
    fi

    if [[ "${ENABLE_LOCAL_LOG}" == "1" ]]; then
        mkdir -p "$(dirname "$LOCAL_LOG_FILE")"
        printf '%s %s\n' "$ts" "$msg" >> "$LOCAL_LOG_FILE"

        if (( TEST_MODE == 1 )); then
            printf '%s %s\n' "$ts" "$msg"
        fi
    fi
}

if [[ -f "$CONFIG_FILE" ]]; then
    # shellcheck diable=SC1091
    source "$CONFIG_FILE"
fi

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --test)
                TEST_MODE=1
                shift
                while [[ $# -gt 0 && "$1" != --* ]]; do
                    TEST_VALUES+=("$1")
                    shift
                done
                ;;
            --send)
                TEST_SEND_NOTIFICATIONS=1
                ;;
            -h|--help)
                cat <<'EOF'
Usage:
  pi-throttle-monitor
  pi-throttle-monitor --test <hex> [<hex> ...]

Examples:
  pi-throttle-monitor --test 0x0
  pi-throttle-monitor --test 0x1 0x5 0x50005
  pi-throttle-monitor --test 50005
EOF
                exit 0
                ;;
            *)
                echo "Unknown argument: $1" >&2
                exit 2
                ;;
        esac
    done

    if (( TEST_MODE == 1 )) && (( ${#TEST_VALUES[@]} == 0 )); then
        echo "--test requires at least one hex value" >&2
        exit 2
    fi
}

: "${INTERVAL:=30}"
: "${STATE_FILE:=/var/lib/rpi-throttle-monitor/last_state}"
: "${SEND_HISTORICAL:=0}"
: "${ENABLE_LOCAL_LOG:=1}"
: "${ENABLE_WEBHOOK:=0}"
: "${ENABLE_EMAIL:=0}"

: "${LOCAL_LOG_FILE:=/var/log/rpi-throttle-monitor/rpi-throttle-monitor.log}"
: "${EMAIL_SUBJECT_PREFIX:=[Raspberry Pi Throttle Alert]}"
: "${MAIL_COMMAND:=/usr/sbin/sendmail}"

if [[ "${ENABLE_WEBHOOK}" == "1" ]]; then
    : "${WEBHOOK_URL:?WEBHOOK_URL must be set when ENABLE_WEBHOOK=1}"
fi

if [[ "${ENABLE_EMAIL}" == "1" ]]; then
    : "${EMAIL_TO:?EMAIL_TO must be set when ENABLE_EMAIL=1}"
    : "${EMAIL_FROM:?EMAIL_FROM must be set when ENABLE_EMAIL=1}"
fi

parse_args "$@"

HOSTNAME_VALUE="${HOST_LABEL:-$(hostname)}"

if (( TEST_MODE == 1 )); then
    STATE_FILE="./test/last_state"
    ENABLE_LOCAL_LOG=1
    LOCAL_LOG_FILE="./test/rpi-throttle-monitor.log"
fi

mkdir -p "$(dirname "$STATE_FILE")"
touch "$STATE_FILE"

if [[ "${ENABLE_LOCAL_LOG}" == "1" ]]; then
    mkdir -p "$(dirname "$LOCAL_LOG_FILE")"
    touch "$LOCAL_LOG_FILE"
fi

ts="$(date -Is)"

write_service_log "Starting Raspberry Pi throttle monitor"

if (( TEST_MODE == 1 )); then
    write_service_log "------------ TEST MODE ------------"
fi

decode_flags() {
    local value="$1"
    local msgs=()

    (( value & (1 << 0) )) && msgs+=("under_voltage_now")
    (( value & (1 << 1) )) && msgs+=("freq_capped_now")
    (( value & (1 << 2) )) && msgs+=("throttled_now")
    (( value & (1 << 16) )) && msgs+=("under_voltage_occurred")
    (( value & (1 << 17) )) && msgs+=("freq_capped_occurred")
    (( value & (1 << 18) )) && msgs+=("throttling_occurred")

    if [ ${#msgs[@]} -eq 0 ]; then
        printf 'ok'
    else
        local IFS=,
        printf '%s' "${msgs[*]}"
    fi
}

should_send() {
    local value="$1"

    local active_mask=$(( (1 << 0) | (1 << 1) | (1 << 2) ))
    local historical_mask=$(( (1 << 16) | (1 << 17) | (1 << 18) ))

    local active_now=$(( value & active_mask ))
    local historical=$(( value & historical_mask ))

    if (( active_now != 0 )); then
        return 0
    fi

    if [[ "$SEND_HISTORICAL" == "1" ]] && (( historical != 0 )); then
        return 0
    fi

    return 1
}

log_local() {
    local ts="$1"
    local hex="$2"
    local flags="$3"
    local temp="$4"

    #local msg="timestamp=${ts} host=${HOSTNAME_VALUE} throttled_hex=0x${hex} temp=${temp} flags=${flags}"
    local msg="throttled_hex=0x${hex}\ntemp=${temp}\nflags=${flags}\n"

    write_service_log "$msg"

    if [[ "$ENABLE_LOCAL_LOG" == "1" || "$TEST_MODE" == "1" ]]; then 
        msg="host=${HOSTNAME_VALUE}\n${msg}"
        write_service_log "$msg"
    fi
}

send_webhook() {
    local ts="$1"
    local hex="$2"
    local flags="$3"
    local temp="$4"

    local payload
    payload="$(cat <<EOF 
{
    "timestamp": "$ts",
    "hostname": "$HOSTNAME_VALUE",
    "throttled_hex": "0x$hex",
    "temperature": "$temp",
    "flags": "$flags"
} 
EOF
)"

    curl -fsS -X POST "$WEBHOOK_URL" \
        -H "Content-Type: application/json" \
        -d "$payload"

}

send_email() {
    local ts="$1"
    local hex="$2"
    local flags="$3"
    local temp="$4"

    local subject="${EMAIL_SUBJECT_PREFIX} ${HOSTNAME_VALUE} 0x${hex}"

    local body
    body="$(cat << EOF
Raspberry Pi throttling event detected

Host: ${HOSTNAME_VALUE}
Timestape: ${ts}
Throttled Hex: 0x${hex}
Temperature: ${temp}
Flags: ${flags}

This alert was generated by rpi-throttle-monitor.
EOF
)"
    {
        printf 'To: %s\n' "$EMAIL_TO"
        printf 'From: %s\n' "$EMAIL_FROM"
        printf 'Subject: %s\n' "$subject"
        printf 'Content-Type: text/plain; charset=UTF-8\n'
        printf '\n'
        printf '%s\n' "$body"
    } | "$MAIL_COMMAND" -t
}

normalize_hex() {
    local input="$1"

    input="${input#0x}"
    input="${input#0X}"

    if [[ ! "$input" =~ ^[0-9A-Fa-f]+$ ]]; then
        echo "invalid hex value: $1" >&2
        return 1
    fi

    printf '%s' "${input^^}"
}

process_hex_value() {
    local hex="$1"
    local temp="$2"
    local msg=""
    local value
    local current
    local flags
    local ts

    value=$((16#$hex))
    current="$(cat "$STATE_FILE" 2>/dev/null || true)"

    if [[ "$hex" != "$current" ]]; then
        echo "$hex" > "$STATE_FILE"

        flags="$(decode_flags "$value")"

        if should_send "$value"; then

            write_service_log "throttled_hex=0x${hex} temp=${temp} flags=${flags}"
    
            if [[ "$ENABLE_WEBHOOK" == "1" || ( "$TEST_MODE" == "1" && "$TEST_SEND_NOTIFICATIONS" == "1" ) ]]; then
                if ! send_webhook "$ts" "$hex" "$flags" "$temp"; then
                    write_service_log "webhook delivery failed"
                fi
            fi
    
            if [[ "$ENABLE_EMAIL" == "1" || ( "$TEST_MODE" == "1" && "$TEST_SEND_NOTIFICATIONS" == "1" ) ]]; then
                if ! send_email "$ts" "$hex" "$flags" "$temp"; then
                    write_service_log "email delivery failed"
                fi
            fi
        fi
    fi
}

run_test_mode() {
    local raw_hex
    local hex

    for raw_hex in "${TEST_VALUES[@]}"; do
        hex="$(normalize_hex "$raw_hex")" || exit 2
        process_hex_value "$hex" "97"
    done
}

cleanup() {
    [[ "$CLEANUP_DONE" == "1" ]] && return 0
    CLEANUP_DONE=1

    local ts
    ts="$(date -Is)"
    write_service_log "host=${HOSTNAME_VALUE} service_stopping reason=${STOP_REASON:-EXIT}"
    write_service_log "Stopping Raspberry Pi throttle monitor"
}

on_term() {
    STOP_REASON="SIGTERM"
    cleanup
    exit 0
}

on_int() {
    STOP_REASON="SIGINT"
    cleanup
    exit 0
}

on_exit() {
    STOP_REASON="${STOP_REASON:-EXIT}"
    cleanup
}

trap on_term TERM
trap on_int INT
trap on_exit EXIT

if (( TEST_MODE == 1 )); then
    run_test_mode
    exit 0
fi

while true; do
    if ! raw="$(vcgencmd get_throttled 2>/dev/null)"; then
        logger -t rpi-throttle-monitor "vcgencmd failed"
        sleep "$INTERVAL"
        continue
    fi

    if ! temp_raw="$(vcgencmd measure_temp 2>/dev/null)"; then
        logger -t rpi-throttle-monitor "vcgencmd failed"
        sleep "$INTERVAL"
        continue
    fi

    hex="${raw#throttled=0x}"
    temp="${temp_raw#temp=}"

    process_hex_value "$hex" "$temp"
    sleep "$INTERVAL"
done
    

