#!/bin/bash
# Continuously samples system state. Writes a detailed snapshot when
# anomaly thresholds are crossed, or when the GUI requests a manual capture.

set -u

LOG_DIR=/var/log/freeze-watcher
CONFIG=/etc/freeze-watcher.conf
TRIGGER_FILE="$LOG_DIR/.trigger"

# --- 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
OUTPUT_JSON=0
OUTPUT_HTML=0

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

mkdir -p "$LOG_DIR"

last_capture=0
captures_since_prune=0

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
}

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
}

# Reads a snapshot file and emits one or more diagnosis lines.
# Pattern-matches D-state wchans and PSI metrics against known freeze types.
classify_snapshot() {
    local file="$1"
    local diagnoses=()

    local dstate_section
    dstate_section=$(awk '/^## D-state processes/{flag=1; next} /^## /{flag=0} flag' "$file")
    local softirq
    softirq=$(awk '/^## softirq/{flag=1; next} /^## /{flag=0} flag' "$file")

    local psi_io_some psi_io_full psi_cpu_some psi_mem_some
    psi_io_some=$(awk '/^\[io\]/{f=1; next} f && /^some/{match($0,/avg10=[0-9.]+/); v=substr($0,RSTART+6,RLENGTH-6); split(v,a,"."); print a[1]+0; exit}' "$file")
    psi_io_full=$(awk '/^\[io\]/{f=1; next} f && /^full/{match($0,/avg10=[0-9.]+/); v=substr($0,RSTART+6,RLENGTH-6); split(v,a,"."); print a[1]+0; exit}' "$file")
    psi_cpu_some=$(awk '/^\[cpu\]/{f=1; next} f && /^some/{match($0,/avg10=[0-9.]+/); v=substr($0,RSTART+6,RLENGTH-6); split(v,a,"."); print a[1]+0; exit}' "$file")
    psi_mem_some=$(awk '/^\[memory\]/{f=1; next} f && /^some/{match($0,/avg10=[0-9.]+/); v=substr($0,RSTART+6,RLENGTH-6); split(v,a,"."); print a[1]+0; exit}' "$file")

    local dcount
    dcount=$(echo "$dstate_section" | grep -c "^D ")

    # btrfs lock contention
    if echo "$dstate_section" | grep -qE "btrfs_lock_root_node|btrfs_sync_log"; then
        if [ "${psi_io_some:-0}" -lt 10 ]; then
            diagnoses+=("btrfs metadata lock contention — wchans show btrfs_lock_root_node/sync_log with low PSI io. Typical of qgroups + many snapshots, or heavy fsync churn.")
        else
            diagnoses+=("btrfs lock contention with concurrent disk I/O")
        fi
    fi

    # btrfs metadata pressure (less severe)
    if echo "$dstate_section" | grep -q "btrfs_search_slot" && \
       ! echo "$dstate_section" | grep -q "btrfs_lock_root_node"; then
        diagnoses+=("btrfs metadata pressure — cleaner walking trees, fsyncs queueing")
    fi

    # USB stall
    if echo "$dstate_section" | grep -qE "usb_sg_wait|scsi_eh|sd_check_events|usb_hcd"; then
        diagnoses+=("USB storage stall — one external drive responding slowly")
    fi

    # Memory swap thrash
    if echo "$dstate_section" | grep -q "folio_wait_bit_common"; then
        if [ "${psi_mem_some:-0}" -gt 5 ]; then
            diagnoses+=("memory pressure / swap thrash — folio_wait + PSI mem ${psi_mem_some}%")
        else
            diagnoses+=("page-fault stall on swapped-out pages — apps faulting back from zram/swap (consider lower swappiness)")
        fi
    fi

    # CPU oversubscription
    if [ "${psi_cpu_some:-0}" -gt 30 ] && [ "$dcount" -lt 3 ]; then
        diagnoses+=("CPU oversubscription — PSI cpu ${psi_cpu_some}%, few D-state procs. Heavy compile, transcoding, or runaway process competing for cores.")
    fi

    # NIC IRQ saturation — use delta against the running baseline so old
    # cumulative skew doesn't trigger forever.
    # With RPS enabled, per-flow concentration on one CPU is normal up to ~85%.
    # Only flag when no RPS spreading is happening at all, or it's truly extreme.
    if [ -r "$LOG_DIR/.softirq-baseline" ] && [ -n "$(echo "$softirq" | grep '^ *NET_RX:')" ]; then
        local cur prev
        cur=$(echo "$softirq" | grep '^ *NET_RX:' | sed 's/^ *NET_RX://')
        prev=$(awk '/^ *NET_RX:/{sub(/^ *NET_RX:/, ""); print; exit}' "$LOG_DIR/.softirq-baseline")
        if [ -n "$prev" ]; then
            local max=0 total=0 i=1
            for v in $cur; do
                [[ "$v" =~ ^[0-9]+$ ]] || { i=$((i+1)); continue; }
                local pv
                pv=$(echo "$prev" | awk -v n=$i '{print $n}')
                [[ "$pv" =~ ^[0-9]+$ ]] || pv=0
                local d=$((v - pv))
                [ "$d" -lt 0 ] && d=0
                [ "$d" -gt "$max" ] && max=$d
                total=$((total + d))
                i=$((i+1))
            done
            if [ "$total" -gt 100000 ]; then
                local pct=$((max * 100 / total))
                # Check if RPS is configured (any non-zero, non-single-bit value)
                local rps_active=0
                for rps in /sys/class/net/*/queues/rx-0/rps_cpus; do
                    [ -r "$rps" ] || continue
                    local v
                    v=$(cat "$rps" 2>/dev/null | tr -d ',')
                    # Strip leading zeros, count hex bits set
                    [ "$v" != "0" ] && [ "$v" != "00" ] && [ "$v" != "000" ] \
                        && [ "$v" != "0000" ] && rps_active=1
                done
                if [ "$pct" -gt 90 ] && [ "$rps_active" = "0" ]; then
                    diagnoses+=("NIC softirq saturating one CPU (${pct}% of recent NET_RX, RPS disabled) — enable RPS, install irqbalance")
                elif [ "$pct" -gt 95 ]; then
                    diagnoses+=("NIC softirq extreme concentration (${pct}% of recent NET_RX) — even with RPS, single-flow saturation; consider enabling RFS or reducing peer connections")
                fi
            fi
        fi
    fi

    # Genuine I/O bottleneck
    if [ "${psi_io_full:-0}" -gt 30 ]; then
        diagnoses+=("genuine I/O bottleneck — PSI io full ${psi_io_full}%, device truly can't keep up")
    fi

    if [ ${#diagnoses[@]} -eq 0 ]; then
        echo "  • uncertain — no known signature matched; review wchans and metrics manually"
    else
        for d in "${diagnoses[@]}"; do
            echo "  • $d"
        done
    fi
}

capture_snapshot() {
    local reason="$1"
    local body="$LOG_DIR/.snap-body-$$.tmp"
    local final="$LOG_DIR/snap-$(date +%Y%m%d-%H%M%S).log"

    # CAPTURE D-STATE FIRST - blocked spikes are often <1s long, so we need
    # to grab process state before reading slower /proc/* files
    local dstate_immediate
    dstate_immediate=$(ps -eo state,pid,user,wchan:25,comm,args --no-headers 2>/dev/null | awk '$1=="D"' | head -30)
    local dstate_t1
    sleep 0.05
    dstate_t1=$(ps -eo state,pid,user,wchan:25,comm,args --no-headers 2>/dev/null | awk '$1=="D"' | head -30)
    local dstate_t2
    sleep 0.05
    dstate_t2=$(ps -eo state,pid,user,wchan:25,comm,args --no-headers 2>/dev/null | awk '$1=="D"' | head -30)

    {
        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)"
        # Use the snapshot taken before /proc reads to catch transient blockers.
        # Combine 3 rapid samples (~100ms apart) and dedupe by pid (column 2).
        printf '%s\n%s\n%s\n' "$dstate_immediate" "$dstate_t1" "$dstate_t2" \
            | awk 'NF >= 4 && $1=="D" && !seen[$2]++' | 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
    } > "$body" 2>&1

    # Build final file: header + diagnosis (computed from body) + body
    {
        echo "=== FREEZE EVENT TRIGGERED: $reason ==="
        echo "Time: $(date -Iseconds)"
        echo
        echo "## Diagnosis (auto-classified)"
        classify_snapshot "$body"
        echo
        cat "$body"
    } > "$final"

    rm -f "$body"
    echo "$(date -Iseconds) Captured snapshot to $final (reason: $reason)" >> "$LOG_DIR/events.log"

    # Optional sidecar formats
    local convert_args=""
    [ "${OUTPUT_JSON:-0}" = "1" ] && convert_args="$convert_args --json"
    [ "${OUTPUT_HTML:-0}" = "1" ] && convert_args="$convert_args --html"
    if [ -n "$convert_args" ] && command -v freeze-snapshot-convert >/dev/null 2>&1; then
        freeze-snapshot-convert $convert_args "$final" 2>/dev/null &
    fi

    notify_users "$reason"
}

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

# Initial softirq baseline so classifier can compute deltas
cp /proc/softirqs "$LOG_DIR/.softirq-baseline" 2>/dev/null
softirq_baseline_at=$(date +%s)

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

    # Manual trigger from GUI (file-based, no auth required for users in log group)
    if [ -f "$TRIGGER_FILE" ]; then
        manual_reason="manual"
        if [ -s "$TRIGGER_FILE" ]; then
            manual_reason=$(head -c 100 "$TRIGGER_FILE" | tr -d '\n')
        fi
        rm -f "$TRIGGER_FILE"
        capture_snapshot "$manual_reason" &
        last_capture=$now
        sleep "$SAMPLE_INTERVAL"
        continue
    fi

    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

    # Refresh softirq baseline every 5 minutes so classifier deltas reflect recent activity
    if [ $((now - softirq_baseline_at)) -ge 300 ]; then
        cp /proc/softirqs "$LOG_DIR/.softirq-baseline" 2>/dev/null
        softirq_baseline_at=$now
    fi

    sleep "$SAMPLE_INTERVAL"
done