#!/bin/bash
# Continuously samples system state. Writes a detailed snapshot when
# anomaly thresholds are crossed (D-states, PSI pressure, blocked procs).

set -u

LOG_DIR=/var/log/freeze-watcher
CONFIG=/etc/freeze-watcher.conf

# --- defaults (overridden by /etc/freeze-watcher.conf) -------------------
PSI_IO_THRESHOLD=40
PSI_CPU_THRESHOLD=50
DSTATE_THRESHOLD=5
BLOCKED_THRESHOLD=5
LOAD_THRESHOLD=20
COOLDOWN=30
SAMPLE_INTERVAL=1
MAX_SNAPSHOTS=200
MAX_SNAPSHOT_AGE_DAYS=30
NOTIFY_USERS=""
NOTIFY_URGENCY=normal
CAPTURE_KERNEL_STACKS=1
CAPTURE_BPFTRACE=0
CAPTURE_CGROUP_PSI=1

[ -f "$CONFIG" ] && . "$CONFIG"

mkdir -p "$LOG_DIR"

last_capture=0
captures_since_prune=0

# Send a desktop notification to each configured user.
# Daemon runs as root; we sudo into the user's session for DBus access.
notify_users() {
    [ -z "$NOTIFY_USERS" ] && return 0
    local body="$1"
    local IFS=,
    for user in $NOTIFY_USERS; do
        user=$(echo "$user" | tr -d '[:space:]')
        [ -z "$user" ] && continue
        local uid
        uid=$(id -u "$user" 2>/dev/null) || continue
        sudo -u "$user" \
            DBUS_SESSION_BUS_ADDRESS="unix:path=/run/user/$uid/bus" \
            notify-send -u "$NOTIFY_URGENCY" -a "freeze-watcher" \
                "Freeze captured" "$body" 2>/dev/null &
    done
}

# Delete old snapshots based on MAX_SNAPSHOTS and MAX_SNAPSHOT_AGE_DAYS.
prune_snapshots() {
    if [ "${MAX_SNAPSHOT_AGE_DAYS:-0}" -gt 0 ]; then
        find "$LOG_DIR" -maxdepth 1 -name 'snap-*.log' -type f \
            -mtime "+$MAX_SNAPSHOT_AGE_DAYS" -delete 2>/dev/null
    fi
    if [ "${MAX_SNAPSHOTS:-0}" -gt 0 ]; then
        local files
        files=$(ls -1t "$LOG_DIR"/snap-*.log 2>/dev/null | tail -n "+$((MAX_SNAPSHOTS+1))")
        [ -n "$files" ] && echo "$files" | xargs -r rm -f
    fi
}

capture_kernel_stacks() {
    echo "## Kernel stack traces of D-state procs"
    ps -eo pid,state --no-headers | awk '$2=="D"{print $1}' | while read -r pid; do
        local comm
        comm=$(cat "/proc/$pid/comm" 2>/dev/null) || continue
        echo "--- pid=$pid comm=$comm ---"
        cat "/proc/$pid/stack" 2>/dev/null | head -20
    done
}

capture_cgroup_psi() {
    echo "## Top cgroups by PSI io.some avg10"
    {
        for f in /sys/fs/cgroup/io.pressure \
                 /sys/fs/cgroup/*/io.pressure \
                 /sys/fs/cgroup/*/*/io.pressure; do
            [ -r "$f" ] || continue
            local v
            v=$(awk '/^some/{for(i=1;i<=NF;i++) if($i~/^avg10=/){gsub(/avg10=/,"",$i); print $i; exit}}' "$f" 2>/dev/null)
            [ -z "$v" ] && continue
            [ "$v" = "0.00" ] && continue
            printf "%s\t%s\n" "$v" "$(dirname "$f" | sed 's|/sys/fs/cgroup||')"
        done
    } 2>/dev/null | sort -rn | head -10
    echo
    echo "## Top cgroups by PSI cpu.some avg10"
    {
        for f in /sys/fs/cgroup/cpu.pressure \
                 /sys/fs/cgroup/*/cpu.pressure \
                 /sys/fs/cgroup/*/*/cpu.pressure; do
            [ -r "$f" ] || continue
            local v
            v=$(awk '/^some/{for(i=1;i<=NF;i++) if($i~/^avg10=/){gsub(/avg10=/,"",$i); print $i; exit}}' "$f" 2>/dev/null)
            [ -z "$v" ] && continue
            [ "$v" = "0.00" ] && continue
            printf "%s\t%s\n" "$v" "$(dirname "$f" | sed 's|/sys/fs/cgroup||')"
        done
    } 2>/dev/null | sort -rn | head -10
}

capture_bpftrace() {
    command -v bpftrace >/dev/null 2>&1 || { echo "bpftrace not installed"; return; }
    echo "## bpftrace: top kernel stacks for io_schedule waits over 2s"
    timeout 3 bpftrace -e '
        kprobe:io_schedule { @[kstack(5)] = count(); }
        interval:s:2 { exit(); }
    ' 2>/dev/null | tail -40
}

capture_snapshot() {
    local reason="$1"
    local file="$LOG_DIR/snap-$(date +%Y%m%d-%H%M%S).log"
    {
        echo "=== FREEZE EVENT TRIGGERED: $reason ==="
        echo "Time: $(date -Iseconds)"
        echo
        echo "## /proc/loadavg"
        cat /proc/loadavg
        echo
        echo "## /proc/stat (procs_running, procs_blocked)"
        grep -E "procs_running|procs_blocked|intr|ctxt" /proc/stat | head -5
        echo
        echo "## PSI cpu/io/memory"
        for k in cpu io memory; do
            echo "[$k]"; cat /proc/pressure/$k 2>/dev/null
        done
        echo
        if [ "${CAPTURE_CGROUP_PSI:-1}" = "1" ]; then
            capture_cgroup_psi
            echo
        fi
        echo "## D-state processes (with kernel wchan)"
        ps -eo state,pid,user,wchan:25,comm,args --no-headers | awk '$1=="D"' | head -30
        echo
        if [ "${CAPTURE_KERNEL_STACKS:-1}" = "1" ]; then
            capture_kernel_stacks
            echo
        fi
        echo "## Top 15 by CPU"
        ps -eo pid,user,pcpu,pmem,state,comm --no-headers --sort=-pcpu | head -15
        echo
        echo "## Top 15 by RSS memory"
        ps -eo pid,user,pcpu,pmem,rss,comm --no-headers --sort=-rss | head -15
        echo
        echo "## /proc/meminfo summary"
        grep -E "MemTotal|MemFree|MemAvailable|Dirty|Writeback|Slab|Cached|Buffers" /proc/meminfo
        echo
        echo "## iostat (1s snapshot)"
        iostat -xt 1 2 2>&1 | tail -50
        echo
        echo "## interrupts (top by total)"
        awk 'NR==1{next} {sum=0; for(i=2;i<=NF-2;i++) sum+=$i; if(sum>1000) print sum, $0}' /proc/interrupts | sort -rn | head -15
        echo
        echo "## softirq"
        cat /proc/softirqs
        echo
        if [ "${CAPTURE_BPFTRACE:-0}" = "1" ]; then
            capture_bpftrace
            echo
        fi
        echo "## Recent dmesg (last 50 lines)"
        dmesg -T | tail -50
        echo
        echo "## journal warnings/errors last 60s"
        journalctl --since "60 seconds ago" -p warning --no-pager 2>&1 | tail -40
    } > "$file" 2>&1
    echo "$(date -Iseconds) Captured snapshot to $file (reason: $reason)" >> "$LOG_DIR/events.log"
    notify_users "$reason"
}

echo "$(date -Iseconds) freeze-watcher started (pid $$)" >> "$LOG_DIR/events.log"
prune_snapshots

while true; do
    now=$(date +%s)

    psi_io=$(awk '/^some/{for(i=1;i<=NF;i++) if($i~/^avg10=/){gsub(/avg10=/,"",$i); print int($i+0.5); exit}}' /proc/pressure/io 2>/dev/null)
    psi_cpu=$(awk '/^some/{for(i=1;i<=NF;i++) if($i~/^avg10=/){gsub(/avg10=/,"",$i); print int($i+0.5); exit}}' /proc/pressure/cpu 2>/dev/null)
    dstate=$(ps -eo state --no-headers | grep -c "^D")
    blocked=$(awk '/^procs_blocked/{print $2}' /proc/stat)
    load1=$(awk '{print int($1+0.5)}' /proc/loadavg)

    triggered=""
    [ "${psi_io:-0}" -ge "$PSI_IO_THRESHOLD" ] && triggered="psi_io=$psi_io%"
    [ "${psi_cpu:-0}" -ge "$PSI_CPU_THRESHOLD" ] && triggered="${triggered:+$triggered, }psi_cpu=$psi_cpu%"
    [ "${dstate:-0}" -ge "$DSTATE_THRESHOLD" ] && triggered="${triggered:+$triggered, }dstate=$dstate"
    [ "${blocked:-0}" -ge "$BLOCKED_THRESHOLD" ] && triggered="${triggered:+$triggered, }blocked=$blocked"
    [ "${load1:-0}" -ge "$LOAD_THRESHOLD" ] && triggered="${triggered:+$triggered, }load1=$load1"

    if [ -n "$triggered" ] && [ $((now - last_capture)) -ge "$COOLDOWN" ]; then
        capture_snapshot "$triggered" &
        last_capture=$now
        captures_since_prune=$((captures_since_prune+1))
        if [ "$captures_since_prune" -ge 20 ]; then
            prune_snapshots &
            captures_since_prune=0
        fi
    fi

    sleep "$SAMPLE_INTERVAL"
done