#!/bin/bash
# -- ADAPTIVE STRESS-NG TEST (TMUX) -------------------------------------------
# Auto-sizes workers and memory to the host. Runs hard but stays under 100%
# so the kernel doesn't OOM-kill itself or wedge the box.
#
# stress-ng with --metrics-brief is silent during the run. To give the frontend
# (and CLI) live progress, a background writer injects a status line into the
# log every 5 seconds while stress-ng is running. Tmux still wraps everything
# so it survives SSH disconnects and shows a status bar on tmux attach.
#
# Duration: 20 minutes
# Log:      /tmp/stress_test_<MAC>_<TS>.log
# -----------------------------------------------------------------------------

DURATION=1200
DURATION_HUMAN="20m"
TMUX_SESSION="stress_test"

# -- Identify host via eth0 MAC -----------------------------------------------
get_eth0_mac() {
    local mac=""
    if [ -f /sys/class/net/eth0/address ]; then
        mac=$(cat /sys/class/net/eth0/address 2>/dev/null)
    fi
    if [ -z "$mac" ]; then
        for iface in /sys/class/net/*/address; do
            local name=$(basename "$(dirname "$iface")")
            [ "$name" = "lo" ] && continue
            mac=$(cat "$iface" 2>/dev/null)
            [ -n "$mac" ] && break
        done
    fi
    echo "${mac:-unknown}"
}

ETH0_MAC=$(get_eth0_mac)
MAC_SAFE=$(echo "$ETH0_MAC" | tr ':' '-')
HOSTNAME=$(uname -n)
TS=$(date +%Y%m%d-%H%M%S)
LOG="/tmp/stress_test_${MAC_SAFE}_${TS}.log"
INNER_LOG="/tmp/stress_test_inner.log"
STATUS_FILE="/tmp/stress_test_status.txt"
STATUS_PID_FILE="/tmp/stress_test_status.pid"
WATCHDOG_PID_FILE="/tmp/stress_test_watchdog.pid"
TMUX_CONFIG="/tmp/stress_test_tmux.conf"

START_EPOCH=$(date +%s)

# -- Install stress-ng if missing ---------------------------------------------
if ! command -v stress-ng &>/dev/null; then
    echo "stress-ng not found -- installing..."
    if command -v apt-get &>/dev/null; then
        DEBIAN_FRONTEND=noninteractive apt-get install -y stress-ng &>/dev/null
    elif command -v yum &>/dev/null; then
        yum install -y stress-ng &>/dev/null
    elif command -v dnf &>/dev/null; then
        dnf install -y stress-ng &>/dev/null
    fi
    if ! command -v stress-ng &>/dev/null; then
        echo "ERROR: failed to install stress-ng. Install manually and re-run."
        exit 1
    fi
fi

# -- Ensure tmux --------------------------------------------------------------
if ! command -v tmux &>/dev/null; then
    echo "tmux not found -- installing..."
    if command -v apt-get &>/dev/null; then
        DEBIAN_FRONTEND=noninteractive apt-get install -y tmux &>/dev/null
    elif command -v yum &>/dev/null; then
        yum install -y tmux &>/dev/null
    elif command -v dnf &>/dev/null; then
        dnf install -y tmux &>/dev/null
    fi
    if ! command -v tmux &>/dev/null; then
        echo "ERROR: failed to install tmux. Install manually and re-run."
        exit 1
    fi
fi

# -- Detect resources ---------------------------------------------------------
LOGICAL_CPUS=$(nproc)
PHYSICAL_CPUS=$(grep "physical id" /proc/cpuinfo | sort -u | wc -l)
[ "$PHYSICAL_CPUS" -lt 1 ] && PHYSICAL_CPUS=1
CPU_MODEL=$(grep "model name" /proc/cpuinfo | head -1 | cut -d: -f2 | xargs)

RAM_TOTAL_MB=$(awk '/MemTotal/ {printf "%d", $2/1024}' /proc/meminfo)
RAM_AVAIL_MB=$(awk '/MemAvailable/ {printf "%d", $2/1024}' /proc/meminfo)

# -- Adaptive sizing ----------------------------------------------------------
CPU_WORKERS=$LOGICAL_CPUS
MATRIX_WORKERS=$(( LOGICAL_CPUS / 2 ));   [ "$MATRIX_WORKERS" -lt 1 ] && MATRIX_WORKERS=1
CACHE_WORKERS=$(( LOGICAL_CPUS / 4 ));    [ "$CACHE_WORKERS" -lt 1 ]  && CACHE_WORKERS=1
VM_WORKERS=$(( RAM_TOTAL_MB / 4096 ))
[ "$VM_WORKERS" -lt 2 ]              && VM_WORKERS=2
[ "$VM_WORKERS" -gt "$LOGICAL_CPUS" ] && VM_WORKERS=$LOGICAL_CPUS
# Memory budget: 90% of TOTAL RAM, divided across VM workers.
# With --vm-keep --vm-populate, this memory stays resident for the full run.
# 90% leaves ~10% headroom for kernel + page cache + stressor overhead.
# If exit 137 (OOM kill) shows up at this level, drop to 80%.
VM_TOTAL_MB=$(( RAM_TOTAL_MB * 90 / 100 ))
VM_PER_WORKER_MB=$(( VM_TOTAL_MB / VM_WORKERS ))

# -- Status updater (writes one-liner to STATUS_FILE every second) ------------
update_status() {
    local end_epoch=$(( START_EPOCH + DURATION ))
    while true; do
        local now=$(date +%s)
        local elapsed=$(( now - START_EPOCH ))
        local remaining=$(( end_epoch - now ))
        local el_m=$(( elapsed / 60 ))
        local el_s=$(( elapsed % 60 ))
        local load=$(awk '{print $1}' /proc/loadavg 2>/dev/null)
        local mem_used=$(awk '/MemTotal/{t=$2}/MemAvailable/{a=$2} END{printf "%d", (t-a)/1024}' /proc/meminfo 2>/dev/null)
        local time_now=$(date '+%H:%M:%S')

        if [ $remaining -gt 0 ]; then
            local rm_m=$(( remaining / 60 ))
            local rm_s=$(( remaining % 60 ))
            printf '%s | %s | elapsed: %02d:%02d / remain: %02d:%02d | load: %s | mem used: %s MB | %s' \
                "$HOSTNAME" "$ETH0_MAC" "$el_m" "$el_s" "$rm_m" "$rm_s" "$load" "$mem_used" "$time_now" \
                > "$STATUS_FILE"
        else
            local over=$(( -remaining ))
            local ov_m=$(( over / 60 ))
            local ov_s=$(( over % 60 ))
            printf '%s | %s | elapsed: %02d:%02d / OVERRUN +%02d:%02d | load: %s | mem used: %s MB | %s' \
                "$HOSTNAME" "$ETH0_MAC" "$el_m" "$el_s" "$ov_m" "$ov_s" "$load" "$mem_used" "$time_now" \
                > "$STATUS_FILE"
        fi
        sleep 1
    done
}

# -- Tmux config --------------------------------------------------------------
cat > "$TMUX_CONFIG" << TMUXEOF
set -g status on
set -g status-interval 1
set -g status-position bottom
set -g status-justify left
set -g status-style 'bg=colour234,fg=colour255'
set -g status-left-length 40
set -g status-left  '#[bg=colour22,fg=colour255,bold] STRESS-NG #[bg=colour234,fg=colour22] '
set -g status-right-length 250
set -g status-right '#[fg=colour220]#(cat $STATUS_FILE 2>/dev/null) '
set -g window-status-current-style 'bg=colour22,fg=colour255'
TMUXEOF

# -- Inner script that runs inside tmux ---------------------------------------
# Critical bits:
#   - exec > >(tee -a INNER_LOG) so all output lands in the log file
#   - stdbuf -oL on stress-ng forces line-buffered output
#   - A background "progress writer" appends status lines to the log every
#     5 seconds while stress-ng is running, so the frontend (which tails
#     the log) sees live progress instead of 20 minutes of silence
INNER_SCRIPT=$(mktemp /tmp/stress_inner.XXXXXX.sh)
cat > "$INNER_SCRIPT" << INNEREOF
#!/bin/bash
exec > >(tee -a "$INNER_LOG") 2>&1

echo "------------------------------------------------------------------------"
echo "  ADAPTIVE STRESS-NG TEST"
echo "------------------------------------------------------------------------"
echo "  Hostname:          $HOSTNAME"
echo "  ETH0 MAC:          $ETH0_MAC"
echo "  Started:           \$(date '+%Y-%m-%d %H:%M:%S %Z')"
echo "  Duration:          $DURATION_HUMAN"
echo "  Log file:          $LOG"
echo "------------------------------------------------------------------------"
echo "  CPU model:         $CPU_MODEL"
echo "  Physical CPUs:     $PHYSICAL_CPUS"
echo "  Logical cores:     $LOGICAL_CPUS"
echo "  RAM total:         ${RAM_TOTAL_MB} MB"
echo "  RAM available:     ${RAM_AVAIL_MB} MB"
echo "------------------------------------------------------------------------"
echo "  Workload sizing (adaptive):"
echo "    CPU workers:       $CPU_WORKERS"
echo "    Matrix workers:    $MATRIX_WORKERS"
echo "    Cache workers:     $CACHE_WORKERS"
echo "    VM workers:        $VM_WORKERS"
echo "    VM mem/worker:     ${VM_PER_WORKER_MB} MB"
echo "    VM mem total:      ${VM_TOTAL_MB} MB (90% of RAM)"
echo "  stress-ng version:   \$(stress-ng --version 2>&1 | head -1)"
echo "------------------------------------------------------------------------"
echo ""

echo "[\$(date '+%H:%M:%S')] Pre-test sensors snapshot:"
if command -v sensors &>/dev/null; then
    sensors 2>/dev/null | grep -E "Core|Package|temp" | head -20
else
    echo "  (lm-sensors not installed -- skipping temp snapshot)"
fi
echo ""

echo "[\$(date '+%H:%M:%S')] Starting stress-ng for ${DURATION_HUMAN}..."
echo "------------------------------------------------------------------------"

# Background progress writer: appends a status line to the log every 5s.
(
    while true; do
        STATUS=\$(cat "$STATUS_FILE" 2>/dev/null)
        echo "[progress] \$STATUS"
        sleep 5
    done
) &
PROGRESS_PID=\$!

# stress-ng changes:
#   - dropped --verify  (doubles work for some stressors, extends runs)
#   - dropped --hdd     (IO starvation prevented other stressors from honoring
#                        timeout; we already test memory pressure via --vm)
#   - dropped --io      (same — sync stressor wedges on dirty page flush)
#   - kept timer-pacing CPU/matrix/cache/vm stressors which honor timeout
#     promptly because they're CPU-bound, not IO-blocked
stdbuf -oL -eL stress-ng \\
    --cpu        $CPU_WORKERS \\
    --cpu-method all \\
    --matrix     $MATRIX_WORKERS \\
    --cache      $CACHE_WORKERS \\
    --vm         $VM_WORKERS \\
    --vm-bytes   ${VM_PER_WORKER_MB}M \\
    --vm-method  all \\
    --vm-keep    \\
    --vm-populate \\
    --timeout    ${DURATION}s \\
    --metrics-brief \\
    --times \\
    --tz &
STRESS_PID=\$!

# Hard kill watchdog: if stress-ng (PID \$STRESS_PID specifically) hasn't
# exited DURATION+60s after launch, kill ONLY that PID. We do NOT use
# pkill -f stress-ng because that matches every stress-ng on the system,
# which is dangerous if this script is ever launched twice or this watchdog
# is orphaned beyond its parent's lifetime.
WATCHDOG_DEADLINE=$(( DURATION + 60 ))
(
    sleep \$WATCHDOG_DEADLINE
    if kill -0 \$STRESS_PID 2>/dev/null; then
        kill -9 \$STRESS_PID 2>/dev/null
        echo "[watchdog] stress-ng PID \$STRESS_PID exceeded \${WATCHDOG_DEADLINE}s -- sent SIGKILL"
    fi
) &
WATCHDOG_PID=\$!
echo "\$WATCHDOG_PID" > "$WATCHDOG_PID_FILE"

# Wait for stress-ng to finish on its own
wait \$STRESS_PID
STRESS_EXIT=\$?

# Stop watchdog and progress writer
kill \$WATCHDOG_PID 2>/dev/null
wait \$WATCHDOG_PID 2>/dev/null
rm -f "$WATCHDOG_PID_FILE"
kill \$PROGRESS_PID 2>/dev/null
wait \$PROGRESS_PID 2>/dev/null

echo "------------------------------------------------------------------------"
echo "[\$(date '+%H:%M:%S')] stress-ng exited with code \$STRESS_EXIT"
echo ""

echo "[\$(date '+%H:%M:%S')] Post-test sensors snapshot:"
if command -v sensors &>/dev/null; then
    sensors 2>/dev/null | grep -E "Core|Package|temp" | head -20
else
    echo "  (lm-sensors not installed -- skipping temp snapshot)"
fi
echo ""

echo "[\$(date '+%H:%M:%S')] Kernel ring buffer (filtered for issues):"
DMESG_OUT=\$(dmesg 2>/dev/null | tail -200 | grep -iE "oom|hung|panic|mce|throttl" | tail -30)
if [ -n "\$DMESG_OUT" ]; then
    echo "\$DMESG_OUT"
else
    echo "  (no concerning kernel messages found)"
fi
echo ""

echo "------------------------------------------------------------------------"
echo "  TEST COMPLETE"
echo "------------------------------------------------------------------------"
echo "  Hostname:          $HOSTNAME"
echo "  ETH0 MAC:          $ETH0_MAC"
echo "  Finished:          \$(date '+%Y-%m-%d %H:%M:%S %Z')"
echo "  Exit code:         \$STRESS_EXIT"
if [ "\$STRESS_EXIT" -eq 0 ]; then
    echo "  Result:            PASS -- all stressors completed without errors"
else
    echo "  Result:            FAIL -- stress-ng reported issues (exit \$STRESS_EXIT)"
fi
echo "  Log file:          $LOG"
echo "------------------------------------------------------------------------"

cp "$INNER_LOG" "$LOG" 2>/dev/null
sync
echo "__STRESS_TEST_DONE__" >> "$INNER_LOG"
INNEREOF
chmod +x "$INNER_SCRIPT"

# -- Launch ------------------------------------------------------------------
# Kill any stale state from a prior run that was interrupted/killed.
tmux kill-session -t "$TMUX_SESSION" 2>/dev/null

# Kill the previous status-updater via its PID file
if [ -f "$STATUS_PID_FILE" ]; then
    OLD_PID=$(cat "$STATUS_PID_FILE" 2>/dev/null)
    if [ -n "$OLD_PID" ] && kill -0 "$OLD_PID" 2>/dev/null; then
        kill -9 "$OLD_PID" 2>/dev/null
    fi
    rm -f "$STATUS_PID_FILE"
fi

# Kill the previous watchdog via its PID file. CRITICAL: an orphaned
# watchdog from a prior run will eventually wake up and SIGKILL stress-ng,
# even if that stress-ng is from a NEW run minutes/hours later.
if [ -f "$WATCHDOG_PID_FILE" ]; then
    OLD_PID=$(cat "$WATCHDOG_PID_FILE" 2>/dev/null)
    if [ -n "$OLD_PID" ] && kill -0 "$OLD_PID" 2>/dev/null; then
        kill -9 "$OLD_PID" 2>/dev/null
    fi
    rm -f "$WATCHDOG_PID_FILE"
fi

# Belt-and-suspenders: kill any leftover stress-ng workers from a prior
# interrupted run, plus any inner-script bash processes.
pkill -9 -f "stress-ng" 2>/dev/null
pkill -9 -f "stress_inner\." 2>/dev/null

rm -f "$INNER_LOG" "$STATUS_FILE"
echo "Initializing..." > "$STATUS_FILE"
touch "$INNER_LOG"

update_status &
STATUS_PID=$!
echo "$STATUS_PID" > "$STATUS_PID_FILE"

tmux -f "$TMUX_CONFIG" new-session -d -s "$TMUX_SESSION" -x 220 -y 50 \
    "bash $INNER_SCRIPT"

echo "Tmux session '$TMUX_SESSION' started -- streaming output..."
echo "Per-host log file: $LOG"
echo "(to attach manually: tmux attach -t $TMUX_SESSION  --  detach: Ctrl-b d)"
echo "--------------------------------------------------------"

# Tail the inner log -- this is what the frontend sees through SSH
tail -n +1 -F "$INNER_LOG" 2>/dev/null &
TAIL_PID=$!

TIMEOUT=1500
ELAPSED=0
SEEN_DONE=false

while [ $ELAPSED -lt $TIMEOUT ]; do
    if grep -q "__STRESS_TEST_DONE__" "$INNER_LOG" 2>/dev/null; then
        SEEN_DONE=true
        sleep 2  # let tail flush remaining lines
        break
    fi
    sleep 1
    ELAPSED=$(( ELAPSED + 1 ))
done

kill $TAIL_PID 2>/dev/null
wait $TAIL_PID 2>/dev/null
kill $STATUS_PID 2>/dev/null
wait $STATUS_PID 2>/dev/null

rm -f "$INNER_SCRIPT" "$TMUX_CONFIG" "$STATUS_FILE" "$STATUS_PID_FILE" "$WATCHDOG_PID_FILE"

if [ "$SEEN_DONE" = true ]; then
    exit 0
else
    echo "  WARNING: Timed out waiting for stress test to complete."
    echo "  Check tmux session '$TMUX_SESSION' on the remote host manually."
    exit 1
fi