#!/usr/bin/bash
set -Eeuo pipefail

usage() {
    printf '%s\n' \
        'usage: nabu-suspend-diagnostics prepare' \
        '       nabu-suspend-diagnostics snapshot [LABEL]' \
        '       nabu-suspend-diagnostics test freezer|devices|platform|processors|core' \
        '       nabu-suspend-diagnostics suspend [RTC_SECONDS]'
}

[[ $(id -u) -eq 0 ]] || { printf 'Run as root.\n' >&2; exit 1; }
action="${1:-}"
log_root=/var/log/nabu-suspend-diagnostics
trace_root=/sys/kernel/tracing
[[ -d $trace_root ]] || trace_root=/sys/kernel/debug/tracing

mount_debug_filesystems() {
    mountpoint -q /sys/kernel/debug || mount -t debugfs debugfs /sys/kernel/debug
    mountpoint -q /sys/fs/pstore || mount -t pstore pstore /sys/fs/pstore
}

write_if_present() {
    local value="$1" path="$2"
    [[ -w $path ]] && printf '%s\n' "$value" >"$path"
}

prepare() {
    mount_debug_filesystems
    write_if_present 1 /sys/power/pm_print_times
    write_if_present 1 /sys/power/pm_debug_messages
    write_if_present 1 /sys/kernel/debug/pstore/record_ftrace
    if [[ -d $trace_root/events/power/suspend_resume ]]; then
        write_if_present 1 "$trace_root/events/power/suspend_resume/enable"
    fi
    if [[ -d $trace_root/events/power/device_pm_callback_start ]]; then
        write_if_present 1 "$trace_root/events/power/device_pm_callback_start/enable"
        write_if_present 1 "$trace_root/events/power/device_pm_callback_end/enable"
    fi
    if [[ -d $trace_root/events/rpm ]]; then
        write_if_present 1 "$trace_root/events/rpm/enable"
    fi
    printf 'Suspend diagnostics armed. Persistent function tracing has runtime overhead.\n'
}

snapshot() {
    local label="${1:-manual}" stamp target
    stamp="$(date -u +%Y%m%dT%H%M%SZ)"
    target="$log_root/${stamp}-${label//[^A-Za-z0-9_.-]/_}"
    install -d -m0700 "$target"
    uname -a >"$target/uname.txt"
    cat /proc/cmdline >"$target/cmdline.txt"
    cat /sys/power/state >"$target/power-state.txt"
    cat /sys/power/mem_sleep >"$target/mem-sleep.txt" 2>/dev/null || :
    cat /sys/power/pm_test >"$target/pm-test.txt" 2>/dev/null || :
    cat /sys/power/suspend_stats/* >"$target/suspend-stats.txt" 2>/dev/null || :
    cat /sys/power/pm_wakeup_irq >"$target/pm-wakeup-irq.txt" 2>/dev/null || :
    cat /proc/interrupts >"$target/interrupts.txt"
    find /sys/class/wakeup -maxdepth 2 -type f -readable -print -exec sh -c \
        'printf " = "; cat "$1" 2>/dev/null || true' sh {} \; \
        >"$target/wakeup-sources.txt" 2>/dev/null || :
    for rp in /sys/class/remoteproc/remoteproc*; do
        [[ -d $rp ]] || continue
        printf '%s name=' "$rp"
        cat "$rp/name" 2>/dev/null || :
        printf ' state='
        cat "$rp/state" 2>/dev/null || :
    done >"$target/remoteproc.txt"
    dmesg --color=never >"$target/dmesg.txt"
    journalctl -b --no-pager >"$target/journal.txt"
    if [[ -r $trace_root/trace ]]; then
        cat "$trace_root/trace" >"$target/kernel-trace.txt"
    fi
    printf '%s\n' "$target"
}

case "$action" in
    prepare)
        prepare
        ;;
    snapshot)
        snapshot "${2:-manual}"
        ;;
    test)
        level="${2:-}"
        case "$level" in
            freezer|devices|platform|processors|core) ;;
            *) usage; exit 2 ;;
        esac
        prepare
        snapshot "before-pm-test-$level"
        printf '%s\n' "$level" >/sys/power/pm_test
        sync
        systemctl suspend
        printf 'none\n' >/sys/power/pm_test
        snapshot "after-pm-test-$level"
        ;;
    suspend)
        seconds="${2:-45}"
        [[ $seconds =~ ^[0-9]+$ && $seconds -ge 15 && $seconds -le 300 ]] || {
            printf 'RTC_SECONDS must be between 15 and 300.\n' >&2
            exit 2
        }
        rtc=/sys/class/rtc/rtc1/wakealarm
        [[ -w $rtc ]] || { printf 'PMIC rtc1 wakealarm is unavailable.\n' >&2; exit 1; }
        prepare
        snapshot before-real-suspend
        printf '0\n' >"$rtc"
        printf '+%s\n' "$seconds" >"$rtc"
        printf 'none\n' >/sys/power/pm_test
        sync
        systemctl suspend
        printf '0\n' >"$rtc" || :
        snapshot after-real-suspend
        ;;
    *)
        usage
        exit 2
        ;;
esac
