#!/usr/bin/env bash
VERSION="0.10.1"
# ==============================================================================
# Solus Linux Kernel Sanity & Smoke Test Script
# ==============================================================================

# Defaults
RUN_TEST=1
RUN_NVIDIA=0
NV_LOAD_SECS=30
SCRIPT_NAME="${0##*/}"

# Basic cli options
for arg in "$@"; do
    case "$arg" in
        --help|help)
            echo "solus-kst v$VERSION"
            echo "Solus Linux Kernel Sanity & Smoke Test"
            echo ""
            echo "Usage: sudo solus-kst [options]"
            echo ""
            echo "Available Options:"
            echo ""
            echo "--skip-test : Audit only, no tests."
            echo "--nvidia    : Add in-depth NVIDIA driver checks. With --skip-test,"
            echo "              only the passive checks run (no GPU load test)."
            echo "--help      : Help"
            exit 0
            ;;
        --skip-test)
            RUN_TEST=0
            ;;
        --nvidia)
            RUN_NVIDIA=1
            ;;
        *)
            echo "Unknown option: $arg (see '$SCRIPT_NAME --help')" >&2
            exit 2
            ;;
    esac
done

# Re-exec as root if not sudo user
if [[ $EUID -ne 0 ]]; then
    exec sudo bash "$0" "$@"
fi

# ==============================================================================
# Dependency Check
# ==============================================================================
MISSING_DEPS=0
for cmd in sysbench stress-ng; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
        echo "Error: Required dependency '$cmd' is not installed." >&2
        MISSING_DEPS=1
    fi
done

if [ "$MISSING_DEPS" -ne 0 ]; then
    echo "Please install missing dependencies and try again." >&2
    exit 1
fi

clear
echo -e "\033[1;36m────────────────────────────────────────────────────────────"
echo -e " Solus Kernel Smoke Test v$VERSION"
echo -e "────────────────────────────────────────────────────────────\033[0m"

# Safely create a temporary file
LOG_FILE=$(mktemp /tmp/solus_kernel_test_XXXXXX.log)
LOG_MOVED=0

# ==============================================================================
# Cleanup Trap
# ==============================================================================
cleanup() {
    # Stop the NVIDIA load test if the script died while it was running.
    if [[ -n "${NV_STRESS_PID:-}" ]]; then
        kill "$NV_STRESS_PID" 2>/dev/null
        wait "$NV_STRESS_PID" 2>/dev/null
    fi
    [[ -n "${NV_STRESS_LOG:-}" ]] && rm -f "$NV_STRESS_LOG"

    # Unmount the loopback image if the script died mid-test.
    if [[ -n "${MNT_DIR:-}" ]] && mountpoint -q "$MNT_DIR" 2>/dev/null; then
        umount "$MNT_DIR" 2>/dev/null
    fi
    [[ -n "${IMG_FILE:-}"  ]] && rm -f "$IMG_FILE"
    [[ -n "${TEST_FILE:-}" ]] && rm -f "$TEST_FILE"
    [[ -n "${MNT_DIR:-}" && -d "${MNT_DIR:-}" ]] && rmdir "$MNT_DIR" 2>/dev/null

    if [[ "$LOG_MOVED" -eq 0 && -f "${LOG_FILE:-}" ]]; then
        echo -e "\nInterrupted. Partial log retained at: $LOG_FILE" >&2
    fi
    return 0
}
trap cleanup EXIT
trap 'exit 130' INT TERM

# ==============================================================================
# Output & Tally Helpers
# ==============================================================================
PASS_COUNT=0
WARN_COUNT=0
FAIL_COUNT=0
declare -a WARNED_TESTS=()
declare -a FAILED_TESTS=()

C_GREEN="\033[1;32m"
C_YELLOW="\033[1;33m"
C_RED="\033[1;31m"
C_CYAN="\033[1;36m"
C_OFF="\033[0m"

# Echo to console (with color) and to the log file (color stripped).
log_echo() {
    echo -e "$1"
    echo -e "$1" | sed 's/\x1b\[[0-9;]*m//g' >> "$LOG_FILE"
}

result_pass() {
    PASS_COUNT=$((PASS_COUNT + 1))
    log_echo " -> $1: ${C_GREEN}PASSED${C_OFF}${2:+ ($2)}"
}

result_warn() {
    WARN_COUNT=$((WARN_COUNT + 1))
    WARNED_TESTS+=("$1")
    log_echo " -> $1: ${C_YELLOW}WARNING${C_OFF}${2:+ ($2)}"
}

result_fail() {
    FAIL_COUNT=$((FAIL_COUNT + 1))
    FAILED_TESTS+=("$1")
    log_echo " -> $1: ${C_RED}FAILED${C_OFF}${2:+ ($2)}"
}

result_info() {
    log_echo " -> $1: ${C_CYAN}INFO${C_OFF}${2:+ ($2)}"
}

log_echo "────────────────────────────────────────────────────────────"
log_echo " Starting Kernel Smoke Test"
log_echo " Date: $(date)"
log_echo " Temporary Log: $LOG_FILE"
log_echo "────────────────────────────────────────────────────────────\n"


# ==============================================================================
# System & Kernel Information
# ==============================================================================
log_echo "### SYSTEM INFORMATION ###"
log_echo "────────────────────────────────────────────────────────────\n"
log_echo "OS Version: $(grep PRETTY_NAME /etc/os-release | cut -d= -f2 | tr -d '\"')"
log_echo "Kernel Version: $(uname -r)"
log_echo "Architecture: $(uname -m)"
log_echo "Uptime: $(uptime -p)"
log_echo "Effective UID: $EUID (running as root)"
log_echo "Invoking User: ${SUDO_USER:-$USER}\n"


# ==============================================================================
# Kernel Currency (running vs. installed)
# ==============================================================================
log_echo "### KERNEL CURRENCY CHECK ###"
log_echo "────────────────────────────────────────────────────────────\n"

RUNNING_KERNEL=$(uname -r)
KERNEL_VARIANT="${RUNNING_KERNEL##*.}"
LATEST_INSTALLED=""
INSTALLED_LIST=""
REBOOT_REQUIRED=0


if compgen -G "/usr/lib/kernel/com.solus-project.*" >/dev/null 2>&1; then
    INSTALLED_LIST=$(
        for k in /usr/lib/kernel/com.solus-project.*; do
            [ -f "$k" ] || continue
            base="${k##*/}"
            rest="${base#com.solus-project.}"
            variant="${rest%%.*}"
            version="${rest#*.}"
            if [ -n "$variant" ] && [ -n "$version" ] && [ "$variant" != "$version" ]; then
                echo "${version}.${variant}"
            fi
        done | sort -V -u
    )
    # Only compare against the same kernel type currently booted. A machine can
    # have both .current and .lts installed, and those are not upgrades of
    # each other.
    LATEST_INSTALLED=$(echo "$INSTALLED_LIST" | grep "\.${KERNEL_VARIANT}\$" | tail -n 1)
    # Fall back to the newest of any type if the running variant isn't on disk.
    [ -z "$LATEST_INSTALLED" ] && LATEST_INSTALLED=$(echo "$INSTALLED_LIST" | tail -n 1)
elif compgen -G "/usr/lib/kernel/vmlinuz-*" >/dev/null 2>&1; then
    INSTALLED_LIST=$(ls -1 /usr/lib/kernel/vmlinuz-* 2>/dev/null \
        | sed 's#.*/vmlinuz-##' | sort -V -u)
    LATEST_INSTALLED=$(echo "$INSTALLED_LIST" | tail -n 1)
fi

log_echo "Running Kernel:   $RUNNING_KERNEL"
log_echo "Kernel Type:      ${KERNEL_VARIANT}"
log_echo "Newest Installed: ${LATEST_INSTALLED:-unknown}"

if [ -n "$INSTALLED_LIST" ]; then
    log_echo "All Installed Kernels:"
    echo "$INSTALLED_LIST" | while IFS= read -r k; do
        [ -n "$k" ] || continue
        if [ "$k" = "$RUNNING_KERNEL" ]; then
            log_echo "    * $k  (running)"
        else
            log_echo "      $k"
        fi
    done
fi

if [ -z "$LATEST_INSTALLED" ]; then
    result_warn "Kernel Currency" "Could not enumerate installed kernels in /usr/lib/kernel"
elif [ "$RUNNING_KERNEL" = "$LATEST_INSTALLED" ]; then
    result_pass "Kernel Currency" "Running the newest installed ${KERNEL_VARIANT} kernel"
else
    REBOOT_REQUIRED=1
    result_warn "Kernel Currency" "Newer kernel installed but not booted"
    log_echo ""
    log_echo "${C_YELLOW}    ****************************************************${C_OFF}"
    log_echo "${C_YELLOW}    *  REBOOT REQUIRED                                 *${C_OFF}"
    log_echo "${C_YELLOW}    *  You are running an older kernel than the one    *${C_OFF}"
    log_echo "${C_YELLOW}    *  installed on disk. Reboot and re-run this test  *${C_OFF}"
    log_echo "${C_YELLOW}    *  to validate the new kernel.                     *${C_OFF}"
    log_echo "${C_YELLOW}    ****************************************************${C_OFF}"
fi

# Boot manager kernel registration (Solus-specific, informational).
if command -v clr-boot-manager >/dev/null 2>&1; then
    if CBM_OUT=$(clr-boot-manager list-kernels 2>/dev/null) && [ -n "$CBM_OUT" ]; then
        log_echo "\nBoot Manager Registered Kernels:"
        echo "$CBM_OUT" | sed 's/^/    /' | tee -a "$LOG_FILE"
    fi
fi
log_echo "\n"


# ==============================================================================
# Kernel Taint Flags
# ==============================================================================
log_echo "### KERNEL TAINT FLAGS ###"
log_echo "────────────────────────────────────────────────────────────\n"

TAINT_VALUE=$(cat /proc/sys/kernel/tainted 2>/dev/null)

if [ -z "$TAINT_VALUE" ]; then
    result_warn "Kernel Taint" "/proc/sys/kernel/tainted unreadable"
elif [ "$TAINT_VALUE" -eq 0 ]; then
    result_pass "Kernel Taint" "Kernel is untainted (0)"
else
    # Bit -> "severity:description". 'crit' bits indicate the kernel has
    # actually misbehaved; 'note' bits are normal on a desktop (Nvidia, DKMS).
    declare -A TAINT_BITS=(
        [0]="note:Proprietary module loaded (G/P)"
        [1]="note:Module force-loaded (F)"
        [2]="crit:SMP kernel on unsupported CPU (S)"
        [3]="note:Module force-unloaded (R)"
        [4]="crit:Machine Check Exception occurred (M)"
        [5]="crit:Bad page detected (B)"
        [6]="note:Taint requested by userspace (U)"
        [7]="crit:Kernel oops/BUG occurred (D)"
        [8]="note:ACPI table overridden (A)"
        [9]="crit:Kernel WARN_ON() issued (W)"
        [10]="note:Staging driver loaded (C)"
        [11]="note:Firmware bug workaround applied (I)"
        [12]="note:Out-of-tree module loaded (O)"
        [13]="note:Unsigned module loaded (E)"
        [14]="crit:Soft lockup occurred (L)"
        [15]="note:Kernel live-patched (K)"
        [16]="note:Auxiliary taint (X)"
        [17]="note:Struct randomization plugin (T)"
        [18]="crit:In-kernel test module loaded (N)"
    )

    log_echo "Taint value: $TAINT_VALUE"
    log_echo "Decoded flags:"

    TAINT_CRIT=0
    for bit in $(printf '%s\n' "${!TAINT_BITS[@]}" | sort -n); do
        if (( (TAINT_VALUE >> bit) & 1 )); then
            SEVERITY="${TAINT_BITS[$bit]%%:*}"
            DESC="${TAINT_BITS[$bit]#*:}"
            if [ "$SEVERITY" = "crit" ]; then
                TAINT_CRIT=1
                log_echo "    ${C_RED}[bit $bit]${C_OFF} $DESC"
            else
                log_echo "    [bit $bit] $DESC"
            fi
        fi
    done

    if [ "$TAINT_CRIT" -eq 1 ]; then
        result_fail "Kernel Taint" "Kernel tainted by a fault condition"
    else
        result_warn "Kernel Taint" "Tainted, but only by benign flags (modules/firmware)"
    fi
fi
log_echo "\n"


# ==============================================================================
# Kernel Boot Flags & Modules
# ==============================================================================
log_echo "### KERNEL BOOT FLAGS ###"
log_echo "────────────────────────────────────────────────────────────\n"
cat /proc/cmdline | tee -a "$LOG_FILE"
log_echo "\n"

log_echo "### LOADED KERNEL MODULES (Summary) ###"
log_echo "────────────────────────────────────────────────────────────\n"
# lsmod prints a header line, so subtract it from the count.
MODULE_COUNT=$(($(lsmod | wc -l) - 1))
log_echo "Total Modules Loaded: $MODULE_COUNT"
log_echo "Top 10 Loaded Modules (See log file for complete list):"
lsmod | head -n 11 | tee -a "$LOG_FILE"
lsmod | tail -n +12 >> "$LOG_FILE"
log_echo "\n"


if [[ $RUN_TEST -eq 1 ]]; then
    # ==============================================================================
    # Core Function Tests - Start
    # ==============================================================================
    log_echo "### CORE SUBSYSTEM TESTS ###"
    log_echo "────────────────────────────────────────────────────────────\n"

    NPROC=$(nproc)

    # --- CPU TEST ---
    log_echo "[*] Testing CPU Computation (${NPROC} threads, ~10s)..."
    SYSBENCH_OUT=$(sysbench cpu \
        --cpu-max-prime=20000 \
        --threads="$NPROC" \
        --time=10 \
        run 2>&1)

    if [ $? -eq 0 ]; then
        CPU_EPS=$(echo "$SYSBENCH_OUT" | grep -i "events per second" | awk '{print $NF}')
        CPU_LAT=$(echo "$SYSBENCH_OUT" | awk '/95th percentile/ {print $NF}')
        result_pass "CPU Test" "${CPU_EPS:-?} events/sec across $NPROC threads"
        [ -n "$CPU_LAT" ] && log_echo "    95th percentile latency: ${CPU_LAT} ms"
        echo "$SYSBENCH_OUT" >> "$LOG_FILE"
    else
        result_fail "CPU Test" "sysbench returned an error"
        echo "$SYSBENCH_OUT" >> "$LOG_FILE"
    fi

    # --- MEMORY TEST ---
    log_echo "[*] Testing Memory Allocation & Integrity (~15s)..."
    if stress-ng --vm 2 --vm-bytes 512M --vm-method all --verify \
                 --timeout 15s >> "$LOG_FILE" 2>&1; then
        result_pass "Memory Test" "stress-ng verified 2x512M, no corruption"
    else
        result_fail "Memory Test" "stress-ng reported errors or verification failure"
    fi

    # --- DISK I/O TEST ---
    log_echo "[*] Testing Physical Disk I/O & Data Integrity..."
    TEST_FILE=$(mktemp /var/tmp/kernel_io_test_XXXXXX)
    IO_FS=$(findmnt -no FSTYPE -T "$TEST_FILE" 2>/dev/null)
    log_echo "    Target: $TEST_FILE (filesystem: ${IO_FS:-unknown})"

    WRITE_OUT=$(dd if=/dev/urandom of="$TEST_FILE" bs=1M count=50 \
                conv=fdatasync 2>&1)
    if [ $? -eq 0 ]; then
        WRITE_SPEED=$(echo "$WRITE_OUT" | tail -n 1 | sed 's/.*, //')

        # Checksum as the kernel currently sees the file (page cache).
        SUM_CACHED=$(sha256sum "$TEST_FILE" | awk '{print $1}')

        # Checksum read straight off the device, bypassing cache.
        DIRECT_MODE="O_DIRECT"
        READ_OK=1
        if READ_OUT=$(dd if="$TEST_FILE" of=/dev/null bs=1M iflag=direct 2>&1); then
            SUM_DISK=$(dd if="$TEST_FILE" bs=1M iflag=direct 2>/dev/null \
                       | sha256sum | awk '{print $1}') || READ_OK=0
        else
            # Some filesystems (tmpfs, certain overlays) reject O_DIRECT.
            DIRECT_MODE="buffered (O_DIRECT unsupported here)"
            if READ_OUT=$(dd if="$TEST_FILE" of=/dev/null bs=1M 2>&1); then
                SUM_DISK=$(sha256sum "$TEST_FILE" | awk '{print $1}') || READ_OK=0
            else
                READ_OK=0
            fi
        fi

        if [ "$READ_OK" -eq 0 ] || [ -z "$SUM_DISK" ]; then
            result_fail "Disk I/O Test" "Read-back failed"
        elif [ "$SUM_CACHED" != "$SUM_DISK" ]; then
            result_fail "Disk I/O Test" "DATA CORRUPTION: checksum mismatch on read-back"
            log_echo "    Cached: $SUM_CACHED"
            log_echo "    Disk:   $SUM_DISK"
        else
            READ_SPEED=$(echo "$READ_OUT" | tail -n 1 | sed 's/.*, //')
            result_pass "Disk I/O Test" "50MB verified, read mode: $DIRECT_MODE"
            log_echo "    Write (fdatasync): $WRITE_SPEED"
            log_echo "    Read:              $READ_SPEED"
            log_echo "    SHA256 matched:    ${SUM_CACHED:0:16}..."
        fi
    else
        result_fail "Disk I/O Test" "Write failed"
    fi
    rm -f "$TEST_FILE"

    # --- FILESYSTEM & BLOCK DEVICE TEST ---
    log_echo "[*] Testing Loopback Mount & Filesystem Drivers..."
    MNT_DIR=$(mktemp -d /tmp/mnt_test_XXXXXX)
    IMG_FILE=$(mktemp /var/tmp/img_test_XXXXXX.img)

    if dd if=/dev/zero of="$IMG_FILE" bs=1M count=10 > /dev/null 2>&1 && \
       mkfs.ext4 -q -F "$IMG_FILE" > /dev/null 2>&1 && \
       mount -o loop "$IMG_FILE" "$MNT_DIR" > /dev/null 2>&1; then

       # Prove the mounted filesystem is actually writable, not just mounted.
       if echo "kernel-smoke-test" > "$MNT_DIR/probe" 2>/dev/null && \
          [ "$(cat "$MNT_DIR/probe" 2>/dev/null)" = "kernel-smoke-test" ]; then
           umount "$MNT_DIR"
           result_pass "Mount Test" "ext4 loopback mounted and read/write verified"
       else
           umount "$MNT_DIR"
           result_fail "Mount Test" "Mounted but read/write probe failed"
       fi
    else
       result_fail "Mount Test" "Block/VFS layer error"
    fi

    # --- VIRTUAL FILESYSTEM TEST ---
    log_echo "[*] Testing Virtual Filesystems..."
    if [ -d "/proc/1" ] && [ -d "/sys/kernel" ]; then
        result_pass "VFS Test" "/proc and /sys are populated"
    else
        result_fail "VFS Test" "Missing critical VFS structures"
    fi

    # --- RNG / ENTROPY TEST ---
    log_echo "[*] Testing Kernel Random Number Generator..."
    if dd if=/dev/urandom of=/dev/null bs=1K count=10 > /dev/null 2>&1; then
        result_pass "RNG Test" "/dev/urandom accessible"
    else
        result_fail "RNG Test" "Entropy pool issue"
    fi

    # --- CGROUPS & NAMESPACES TEST ---
    log_echo "[*] Testing Cgroups & Namespaces..."
    if [ -d "/sys/fs/cgroup" ] && unshare --user --pid --map-root-user echo >/dev/null 2>&1; then
        result_pass "Cgroups/Namespaces" "Sandboxing supported"
    else
        result_warn "Cgroups/Namespaces" "Cgroups missing or namespace creation failed"
    fi

    # --- AUDIO SUBSYSTEM TEST ---
    log_echo "[*] Testing Audio Subsystem..."
    if [ -d "/dev/snd" ] && grep -qE '^[[:space:]]*[0-9]+[[:space:]]' /proc/asound/cards 2>/dev/null; then
        CARD_COUNT=$(grep -cE '^[[:space:]]*[0-9]+[[:space:]]' /proc/asound/cards)
        result_pass "Audio Test" "$CARD_COUNT sound card(s) detected"
    else
        result_warn "Audio Test" "No sound cards found or ALSA inactive"
    fi

    # --- INPUT SUBSYSTEM TEST ---
    log_echo "[*] Testing Input Subsystem (evdev)..."
    if ls /dev/input/event* >/dev/null 2>&1; then
        EVENT_COUNT=$(ls /dev/input/event* 2>/dev/null | wc -l)
        result_pass "Input Test" "$EVENT_COUNT evdev device(s) detected"
    else
        result_warn "Input Test" "No evdev input devices found"
    fi

    # --- RTC (Real-Time Clock) TEST ---
    log_echo "[*] Testing Real-Time Clock (RTC)..."
    if cat /sys/class/rtc/rtc0/time >/dev/null 2>&1 || cat /proc/driver/rtc >/dev/null 2>&1; then
        result_pass "RTC Test" "Hardware clock accessible"
    else
        result_fail "RTC Test" "RTC node missing"
    fi

    # --- GRAPHICS / DRM TEST ---
    log_echo "[*] Testing Graphics (DRM/KMS) Subsystem..."
    if [ -d "/dev/dri" ] && ls /dev/dri/card* >/dev/null 2>&1; then
        result_pass "Graphics Test" "DRM nodes found in /dev/dri"
    else
        result_warn "Graphics Test" "No DRM nodes found (driver failure or headless?)"
    fi

    # --- GPU DRIVER BINDING TEST ---
    # A DRM node can exist while the real driver failed to bind, so confirm the
    # kernel has actually attached a driver to each display-class PCI device.
    log_echo "[*] Testing GPU Driver Binding..."
    GPU_DEVICES=$(lspci -D | grep -iE "VGA compatible controller|3D controller|Display controller")

    if [ -z "$GPU_DEVICES" ]; then
        result_warn "GPU Driver Binding" "No display-class PCI devices found"
    else
        GPU_UNBOUND=0
        while IFS= read -r gpu_line; do
            GPU_ADDR=$(echo "$gpu_line" | awk '{print $1}')
            GPU_NAME=$(echo "$gpu_line" | cut -d: -f4- | sed 's/^ *//')
            GPU_DRIVER=$(lspci -s "$GPU_ADDR" -k 2>/dev/null \
                | awk -F': ' '/Kernel driver in use/ {print $2}')

            if [ -n "$GPU_DRIVER" ]; then
                log_echo "    $GPU_ADDR  driver: ${C_GREEN}${GPU_DRIVER}${C_OFF}  ($GPU_NAME)"
            else
                GPU_UNBOUND=1
                log_echo "    $GPU_ADDR  driver: ${C_RED}NONE${C_OFF}  ($GPU_NAME)"
            fi
        done <<< "$GPU_DEVICES"

        if [ "$GPU_UNBOUND" -eq 1 ]; then
            result_warn "GPU Driver Binding" "A GPU has no kernel driver bound"
        else
            result_pass "GPU Driver Binding" "All GPUs have a kernel driver bound"
        fi
    fi

    # --- HARDWARE ENUMERATION TEST ---
    log_echo "[*] Testing PCI Bus Enumeration..."
    if lspci | grep -q "Host bridge"; then
        PCI_COUNT=$(lspci | wc -l)
        result_pass "PCI Test" "$PCI_COUNT devices enumerated"
    else
        result_fail "PCI Test" "lspci returned empty or failed"
    fi

    # --- USB ENUMERATION TEST ---
    log_echo "[*] Testing USB Bus Enumeration..."
    if command -v lsusb >/dev/null 2>&1; then
        USB_OUT=$(lsusb 2>/dev/null)
        USB_HUBS=$(echo "$USB_OUT" | grep -ci "root hub")
        USB_TOTAL=$(echo "$USB_OUT" | grep -c .)

        if [ "$USB_HUBS" -gt 0 ]; then
            USB_PERIPHERALS=$((USB_TOTAL - USB_HUBS))
            result_pass "USB Test" "$USB_HUBS root hub(s), $USB_PERIPHERALS peripheral(s)"
            echo "$USB_OUT" >> "$LOG_FILE"
        elif [ "$USB_TOTAL" -gt 0 ]; then
            result_warn "USB Test" "Devices listed but no root hub (xHCI driver issue?)"
        else
            result_warn "USB Test" "No USB devices enumerated"
        fi
    else
        result_warn "USB Test" "lsusb not available"
    fi

    rm -f "$IMG_FILE"
    rmdir "$MNT_DIR" 2>/dev/null

    # --- NETWORK TEST ---
    log_echo "[*] Testing Network Stack..."
    if ip link show lo | grep -q "state UNKNOWN\|state UP"; then
        if ping -c 3 -W 3 1.1.1.1 > /dev/null 2>&1; then
            result_pass "Network Test" "External ping succeeded"
        else
            result_warn "Network Test" "Loopback up, external ping failed (routing/offline?)"
        fi
    else
        result_fail "Network Test" "Loopback interface down"
    fi

    log_echo "\n"
    # ==============================================================================
    # Core Function Tests - End
    # ==============================================================================
fi


# ==============================================================================
# NVIDIA Driver Tests (optional, enabled with --nvidia)
# ==============================================================================
# Kernel updates are the most common way the NVIDIA driver breaks, so these
# checks focus on the kernel side: is there a module built for this kernel,
# did it load and bind, does it match userspace, and does the GPU do work.
NV_STRESS_PID=""
NV_STRESS_LOG=""

nvidia_tests() {
    local kver a drv render node
    kver=$(uname -r)

    log_echo "### NVIDIA DRIVER TESTS ###"
    log_echo "────────────────────────────────────────────────────────────\n"

    # --- NVIDIA hardware present? (PCI vendor 10de, display class 03xx) ---
    local -a nv_addrs
    mapfile -t nv_addrs < <(lspci -Dn -d 10de: 2>/dev/null | awk '$2 ~ /^03/ {print $1}')

    if [ "${#nv_addrs[@]}" -eq 0 ]; then
        result_info "NVIDIA" "No NVIDIA GPU on the PCI bus; skipping NVIDIA tests"
        log_echo "\n"
        return
    fi

    log_echo "NVIDIA GPU(s) found: ${#nv_addrs[@]}"
    for a in "${nv_addrs[@]}"; do
        log_echo "    $a  $(lspci -s "$a" 2>/dev/null | cut -d: -f3- | sed 's/^ *//')"
    done
    log_echo ""

    # --- Module built for the running kernel? ---
    log_echo "[*] Checking for an NVIDIA module built for $kver..."
    local nv_mod_path nv_disk_ver nv_license nv_flavor
    nv_mod_path=$(modinfo -k "$kver" -n nvidia 2>/dev/null)

    if [ -z "$nv_mod_path" ]; then
        # Present for some other kernel but not this one is the classic
        # "kernel updated, driver didn't follow" breakage.
        local others="" d k
        for d in /lib/modules/*/; do
            k=$(basename "$d")
            [ "$k" = "$kver" ] && continue
            modinfo -k "$k" -n nvidia >/dev/null 2>&1 && others+="$k "
        done
        if [ -n "$others" ]; then
            result_fail "NVIDIA Module Present" "No nvidia module for $kver (found for: ${others% })"
            log_echo "    The NVIDIA driver package for this kernel type is missing or out of date."
        else
            result_info "NVIDIA Module Present" "Proprietary driver not installed; skipping remaining NVIDIA tests"
        fi
        log_echo "\n"
        return
    fi

    nv_disk_ver=$(modinfo -k "$kver" -F version nvidia 2>/dev/null)
    nv_license=$(modinfo -k "$kver" -F license nvidia 2>/dev/null)
    case "$nv_license" in
        *GPL*) nv_flavor="open kernel modules" ;;
        *)     nv_flavor="proprietary kernel modules" ;;
    esac
    result_pass "NVIDIA Module Present" "v${nv_disk_ver:-?} for $kver ($nv_flavor)"
    log_echo "    $nv_mod_path"

    # --- If a newer kernel is waiting, make sure the driver is there for it too ---
    if [ "${REBOOT_REQUIRED:-0}" -eq 1 ] && [ -n "$LATEST_INSTALLED" ] \
       && [ -d "/lib/modules/$LATEST_INSTALLED" ]; then
        if modinfo -k "$LATEST_INSTALLED" -n nvidia >/dev/null 2>&1; then
            result_pass "NVIDIA Pending Kernel" "Module present for $LATEST_INSTALLED"
        else
            result_fail "NVIDIA Pending Kernel" "No nvidia module for pending kernel $LATEST_INSTALLED"
            log_echo "    Rebooting into it now would likely leave you without the NVIDIA driver."
        fi
    fi

    # --- Module actually loaded? ---
    log_echo "[*] Checking that the NVIDIA module loaded..."
    if [ ! -d /sys/module/nvidia ]; then
        if [ -d /sys/module/nouveau ]; then
            result_fail "NVIDIA Module Loaded" "nvidia not loaded; nouveau has taken over"
        else
            result_fail "NVIDIA Module Loaded" "nvidia module is installed but not loaded"
        fi
        local why
        why=$(dmesg 2>/dev/null | grep -iE '(nvidia|nvrm).*(fail|error|unknown symbol|mismatch)|Lockdown:.*module' | tail -n 5)
        if [ -n "$why" ]; then
            log_echo "    Possibly related kernel messages:"
            echo "$why" | sed 's/^/    /' | tee -a "$LOG_FILE"
        fi
        log_echo "\n"
        return
    fi

    local nv_loaded_ver
    nv_loaded_ver=$(cat /sys/module/nvidia/version 2>/dev/null)
    if [ -n "$nv_disk_ver" ] && [ -n "$nv_loaded_ver" ] && [ "$nv_loaded_ver" != "$nv_disk_ver" ]; then
        result_warn "NVIDIA Module Loaded" "Running v$nv_loaded_ver, but v$nv_disk_ver is on disk (reboot pending)"
    else
        result_pass "NVIDIA Module Loaded" "v${nv_loaded_ver:-?} loaded"
    fi

    # --- Companion modules & KMS ---
    local missing="" m
    for m in nvidia_modeset nvidia_drm; do
        [ -d "/sys/module/$m" ] || missing+="$m "
    done
    if [ -n "$missing" ]; then
        result_warn "NVIDIA Companion Modules" "Not loaded: ${missing% }"
    else
        local uvm="nvidia_uvm not loaded yet (normal until a CUDA app runs)"
        [ -d /sys/module/nvidia_uvm ] && uvm="nvidia_uvm loaded"
        result_pass "NVIDIA Companion Modules" "nvidia_modeset, nvidia_drm loaded; $uvm"
    fi

    case "$(cat /sys/module/nvidia_drm/parameters/modeset 2>/dev/null)" in
        Y) result_pass "NVIDIA DRM KMS" "nvidia-drm modeset enabled" ;;
        N) result_warn "NVIDIA DRM KMS" "nvidia-drm modeset disabled (Wayland, PRIME sync, GBM unavailable)" ;;
    esac

    # --- Each NVIDIA GPU bound to nvidia and exposed through DRM ---
    log_echo "[*] Checking NVIDIA driver binding per GPU..."
    local bad_bind=0
    NV_LOAD_ADDR=""
    NV_RENDER_NODE=""
    for a in "${nv_addrs[@]}"; do
        drv=""
        [ -e "/sys/bus/pci/devices/$a/driver" ] \
            && drv=$(basename "$(readlink -f "/sys/bus/pci/devices/$a/driver")")
        render=""
        for node in /sys/bus/pci/devices/"$a"/drm/renderD*; do
            [ -e "$node" ] && render="${node##*/}" && break
        done

        if [ "$drv" = "nvidia" ]; then
            log_echo "    $a  driver: ${C_GREEN}nvidia${C_OFF}  render node: ${render:-none}"
            if [ -n "$render" ] && [ -z "$NV_RENDER_NODE" ]; then
                NV_LOAD_ADDR="$a"
                NV_RENDER_NODE="/dev/dri/$render"
            fi
        else
            bad_bind=1
            log_echo "    $a  driver: ${C_RED}${drv:-NONE}${C_OFF}"
        fi
    done
    if [ "$bad_bind" -eq 1 ]; then
        result_fail "NVIDIA Binding" "An NVIDIA GPU is not bound to the nvidia driver"
    else
        result_pass "NVIDIA Binding" "All NVIDIA GPUs bound to nvidia"
    fi

    if [ -d /sys/module/nouveau ]; then
        result_warn "nouveau Conflict" "nouveau is loaded alongside nvidia (check blacklist)"
    fi

    # --- nvidia-smi: userspace <-> kernel path works and versions agree ---
    log_echo "[*] Querying GPUs through nvidia-smi (NVML)..."
    local smi_ok=0
    if ! command -v nvidia-smi >/dev/null 2>&1; then
        result_warn "nvidia-smi" "nvidia-smi not found (userspace driver incomplete?)"
    else
        local smi_out smi_rc smi_count user_ver
        smi_out=$(timeout 30 nvidia-smi \
            --query-gpu=pci.bus_id,name,driver_version,pstate,temperature.gpu,power.draw,memory.used,memory.total \
            --format=csv,noheader 2>&1)
        smi_rc=$?

        if [ "$smi_rc" -eq 124 ]; then
            result_fail "nvidia-smi" "Timed out after 30s (GPU or driver hung?)"
        elif [ "$smi_rc" -ne 0 ]; then
            if grep -qi "version mismatch" <<< "$smi_out"; then
                result_fail "nvidia-smi" "Driver/library version mismatch (userspace updated, reboot needed)"
            else
                result_fail "nvidia-smi" "Failed with exit code $smi_rc"
            fi
            echo "$smi_out" | head -n 5 | sed 's/^/    /' | tee -a "$LOG_FILE"
        else
            smi_count=$(grep -c . <<< "$smi_out")
            user_ver=$(head -n 1 <<< "$smi_out" | awk -F', ' '{print $3}')
            echo "$smi_out" | sed 's/^/    /' | tee -a "$LOG_FILE"

            if [ "$smi_count" -ne "${#nv_addrs[@]}" ]; then
                result_warn "nvidia-smi" "NVML sees $smi_count of ${#nv_addrs[@]} NVIDIA GPU(s)"
            elif [ -n "$nv_loaded_ver" ] && [ "$user_ver" != "$nv_loaded_ver" ]; then
                result_warn "nvidia-smi" "Userspace v$user_ver differs from kernel module v$nv_loaded_ver"
            else
                result_pass "nvidia-smi" "$smi_count GPU(s) responding, userspace v$user_ver"
            fi
            smi_ok=1
        fi
    fi

    # --- Vulkan ICD sees the GPU (works headless, no display needed) ---
    if command -v vulkaninfo >/dev/null 2>&1; then
        if timeout 30 vulkaninfo --summary 2>/dev/null | grep -q "DRIVER_ID_NVIDIA_PROPRIETARY"; then
            result_pass "NVIDIA Vulkan" "Vulkan enumerates the NVIDIA GPU"
        else
            result_warn "NVIDIA Vulkan" "vulkaninfo did not report an NVIDIA device"
        fi
    else
        result_info "NVIDIA Vulkan" "vulkaninfo not installed; skipped"
    fi

    # --- Xid / fatal driver messages already in the ring buffer ---
    nvidia_xid_audit "NVIDIA Xid Audit" "$(dmesg 2>/dev/null)"

    # --- Active GPU load test (skipped with --skip-test) ---
    if [[ $RUN_TEST -eq 1 ]]; then
        if [ "$smi_ok" -eq 1 ]; then
            nvidia_load_test
        else
            result_info "NVIDIA Load Test" "Skipped because nvidia-smi is not working"
        fi
    fi

    log_echo "\n"
}

# Scan kernel log text for NVIDIA Xid events and fatal NVRM messages.
# Most Xids (13, 31, 43...) are faults raised by an application; the ones
# treated as fatal here mean the GPU, its memory, or its firmware failed.
nvidia_xid_audit() {
    local label="$1" text="$2"
    local nvrm xid_lines fatal_xids fatal_msgs
    nvrm=$(grep 'NVRM:' <<< "$text")
    xid_lines=$(grep 'NVRM: Xid' <<< "$nvrm")
    fatal_xids=$(grep -oE 'Xid \([^)]*\): [0-9]+' <<< "$xid_lines" \
        | awk '{print $NF}' | grep -xE '48|62|79|95|119|120' | sort -u | tr '\n' ' ')
    fatal_msgs=$(grep -iE 'fallen off the bus|RmInitAdapter failed|API mismatch' <<< "$nvrm")

    if [ -n "$fatal_xids" ] || [ -n "$fatal_msgs" ]; then
        result_fail "$label" "Fatal GPU/driver errors${fatal_xids:+ (Xid ${fatal_xids% })}"
    elif [ -n "$xid_lines" ]; then
        result_warn "$label" "$(grep -c . <<< "$xid_lines") Xid event(s), likely application faults"
    else
        result_pass "$label" "No Xid events or fatal NVRM messages"
        return
    fi
    { echo "$xid_lines"; echo "$fatal_msgs"; } | grep . | awk '!seen[$0]++' | tail -n 5 | sed 's/^/    /' | tee -a "$LOG_FILE"
}

# Run stress-ng's GPU stressor (EGL/GLES via GBM) against the NVIDIA render
# node and watch nvidia-smi to confirm the GPU actually did the work.
nvidia_load_test() {
    log_echo "[*] NVIDIA GPU load test on $NV_LOAD_ADDR (~${NV_LOAD_SECS}s)..."

    if [ -z "$NV_RENDER_NODE" ]; then
        result_info "NVIDIA Load Test" "No NVIDIA render node found; skipped"
        return
    fi
    if ! stress-ng --stressors 2>/dev/null | tr ' ' '\n' | grep -qx gpu; then
        result_info "NVIDIA Load Test" "This stress-ng has no gpu stressor; skipped"
        return
    fi

    local dmesg_base st_rc sample util temp
    local max_util=0 max_temp=0 bus="${NV_LOAD_ADDR#*:}"
    dmesg_base=$(dmesg 2>/dev/null | wc -l)
    NV_STRESS_LOG=$(mktemp /tmp/kst_nv_stress_XXXXXX.log)

    stress-ng --gpu 1 --gpu-devnode "$NV_RENDER_NODE" \
              --timeout "${NV_LOAD_SECS}s" --metrics-brief \
              > "$NV_STRESS_LOG" 2>&1 &
    NV_STRESS_PID=$!

    # Sample utilization and temperature for this GPU while the load runs.
    while kill -0 "$NV_STRESS_PID" 2>/dev/null; do
        sample=$(timeout 10 nvidia-smi --query-gpu=pci.bus_id,utilization.gpu,temperature.gpu \
                    --format=csv,noheader,nounits 2>/dev/null \
                 | awk -F', ' -v id="$bus" 'tolower($1) ~ (":" tolower(id) "$") {print $2, $3}')
        read -r util temp <<< "$sample"
        [[ "$util" =~ ^[0-9]+$ ]] && (( util > max_util )) && max_util=$util
        [[ "$temp" =~ ^[0-9]+$ ]] && (( temp > max_temp )) && max_temp=$temp
        sleep 2
    done
    wait "$NV_STRESS_PID"
    st_rc=$?
    NV_STRESS_PID=""

    {
        echo "--- STRESS-NG GPU OUTPUT ---"
        cat "$NV_STRESS_LOG"
    } >> "$LOG_FILE"

    # stress-ng reports which GL implementation it actually rendered with.
    local renderer
    renderer=$(grep -oE 'GL_RENDERER: .*' "$NV_STRESS_LOG" | head -n 1 | cut -d' ' -f2-)
    [ -n "$renderer" ] && log_echo "    Renderer: $renderer"

    # Exit 3/4 = no resource / not implemented (e.g. built without EGL/GBM).
    # Match "skipped: N" only when N > 0; a good run prints "skipped: 0".
    if [ "$st_rc" -eq 3 ] || [ "$st_rc" -eq 4 ] \
       || grep -qiE 'not implemented|unimplemented|skipped: [1-9]' "$NV_STRESS_LOG"; then
        result_info "NVIDIA Load Test" "stress-ng gpu stressor unavailable here (needs EGL/GBM); skipped"
    elif [ "$st_rc" -ne 0 ]; then
        result_fail "NVIDIA Load Test" "stress-ng gpu stressor failed (exit $st_rc)"
        tail -n 5 "$NV_STRESS_LOG" | sed 's/^/    /' | tee -a "$LOG_FILE"
    elif [ -n "$renderer" ] && ! grep -qi nvidia <<< "$renderer"; then
        result_warn "NVIDIA Load Test" "Load ran on '$renderer', not the NVIDIA GPU"
    elif [ "$max_util" -eq 0 ]; then
        result_warn "NVIDIA Load Test" "Load ran, but nvidia-smi never saw GPU utilization"
    else
        result_pass "NVIDIA Load Test" "Peak ${max_util}% utilization, ${max_temp}°C"
    fi
    rm -f "$NV_STRESS_LOG"
    NV_STRESS_LOG=""

    # Anything the driver logged while it was under load?
    nvidia_xid_audit "NVIDIA Under Load" "$(dmesg 2>/dev/null | tail -n +"$((dmesg_base + 1))")"
}

if [[ $RUN_NVIDIA -eq 1 ]]; then
    nvidia_tests
fi



# ==============================================================================
# Kernel Log (dmesg) Error Audit
# ==============================================================================
log_echo "### KERNEL LOG AUDIT (dmesg warnings/errors) ###"
log_echo "────────────────────────────────────────────────────────────\n"
log_echo "Checking for emergency/alert/critical/error entries in the kernel ring buffer..."

DMESG_ERRORS=$(dmesg --level=emerg,alert,crit,err 2>/dev/null)

if [ -z "$DMESG_ERRORS" ]; then
    result_pass "dmesg Audit" "No error-level entries found"
else
    DMESG_COUNT=$(echo "$DMESG_ERRORS" | grep -c .)
    result_warn "dmesg Audit" "$DMESG_COUNT error-level entries found"
    log_echo "    Last 5 entries:"
    echo "$DMESG_ERRORS" | tail -n 5 | sed 's/^/    /' | tee -a "$LOG_FILE"
    log_echo "    (Full list written to log file)"
    echo "--- FULL DMESG ERROR LIST ---" >> "$LOG_FILE"
    echo "$DMESG_ERRORS" >> "$LOG_FILE"
fi
log_echo "\n"


# ==============================================================================
# Firmware Load Audit
# ==============================================================================
log_echo "### FIRMWARE LOAD AUDIT ###"
log_echo "────────────────────────────────────────────────────────────\n"
log_echo "Checking for failed firmware loads in the kernel ring buffer..."

FIRMWARE_ERRORS=$(dmesg 2>/dev/null | grep -iE \
    "Direct firmware load for .* failed|firmware: failed to load|failed to load firmware|firmware.*not found|Falling back to (sysfs|user) helper")

if [ -z "$FIRMWARE_ERRORS" ]; then
    result_pass "Firmware Audit" "No firmware load failures detected"
else
    FW_COUNT=$(echo "$FIRMWARE_ERRORS" | grep -c .)
    result_warn "Firmware Audit" "$FW_COUNT firmware load failure(s) detected"
    log_echo "    Last 5 entries:"
    echo "$FIRMWARE_ERRORS" | tail -n 5 | sed 's/^/    /' | tee -a "$LOG_FILE"
    log_echo "    Missing firmware usually means a linux-firmware package is absent,"
    log_echo "    or the new kernel expects a newer firmware blob than is installed."
    echo "--- FULL FIRMWARE ERROR LIST ---" >> "$LOG_FILE"
    echo "$FIRMWARE_ERRORS" >> "$LOG_FILE"
fi
log_echo "\n"


# ==============================================================================
# CPU Vulnerability Mitigations (informational)
# ==============================================================================
log_echo "### CPU VULNERABILITY MITIGATIONS ###"
log_echo "────────────────────────────────────────────────────────────\n"

if [ -d /sys/devices/system/cpu/vulnerabilities ]; then
    VULN_TOTAL=0
    VULN_UNMITIGATED=0

    for vuln_file in /sys/devices/system/cpu/vulnerabilities/*; do
        [ -r "$vuln_file" ] || continue
        VULN_NAME=$(basename "$vuln_file")
        VULN_STATE=$(cat "$vuln_file" 2>/dev/null)
        VULN_TOTAL=$((VULN_TOTAL + 1))

        case "$VULN_STATE" in
            "Not affected"*)
                log_echo "    $(printf '%-24s' "$VULN_NAME") $VULN_STATE"
                ;;
            "Vulnerable"*)
                VULN_UNMITIGATED=$((VULN_UNMITIGATED + 1))
                log_echo "    $(printf '%-24s' "$VULN_NAME") ${C_YELLOW}${VULN_STATE}${C_OFF}"
                ;;
            *)
                log_echo "    $(printf '%-24s' "$VULN_NAME") ${C_CYAN}${VULN_STATE}${C_OFF}"
                ;;
        esac
    done

    result_info "CPU Mitigations" \
        "$VULN_TOTAL vulnerabilities reported, $VULN_UNMITIGATED currently unmitigated"
    log_echo "    Informational only. Mitigation state changes between kernel versions"
    log_echo "    and can be altered by boot flags (e.g. mitigations=off)."
else
    result_info "CPU Mitigations" "/sys/devices/system/cpu/vulnerabilities not present"
fi
log_echo "\n"


# ==============================================================================
# Summary
# ==============================================================================
TOTAL_TESTS=$((PASS_COUNT + WARN_COUNT + FAIL_COUNT))

log_echo "────────────────────────────────────────────────────────────"
log_echo " SUMMARY"
log_echo "────────────────────────────────────────────────────────────\n"
log_echo " Kernel Tested: $RUNNING_KERNEL"
log_echo " Tests/Audits Run:   $TOTAL_TESTS"
log_echo " ${C_GREEN}Passed:${C_OFF}             $PASS_COUNT"
log_echo " ${C_YELLOW}Warnings:${C_OFF}           $WARN_COUNT"
log_echo " ${C_RED}Failed:${C_OFF}             $FAIL_COUNT"

if [ "$WARN_COUNT" -gt 0 ]; then
    log_echo "\n ${C_YELLOW}Warnings in:${C_OFF}"
    for t in "${WARNED_TESTS[@]}"; do
        log_echo "   - $t"
    done
fi

if [ "$FAIL_COUNT" -gt 0 ]; then
    log_echo "\n ${C_RED}Failures in:${C_OFF}"
    for t in "${FAILED_TESTS[@]}"; do
        log_echo "   - $t"
    done
fi

log_echo ""
if [ "$FAIL_COUNT" -gt 0 ]; then
    log_echo " ${C_RED}RESULT: FAIL${C_OFF} - This kernel has problems that need investigation."
    EXIT_CODE=1
elif [ "$WARN_COUNT" -gt 0 ]; then
    log_echo " ${C_YELLOW}RESULT: PASS WITH WARNINGS${C_OFF} - Review the warnings above."
    EXIT_CODE=0
else
    log_echo " ${C_GREEN}RESULT: PASS${C_OFF} - Kernel looks healthy."
    EXIT_CODE=0
fi

if [ "$REBOOT_REQUIRED" -eq 1 ]; then
    log_echo "\n ${C_YELLOW}NOTE: A newer kernel is installed. Reboot and re-run to test it.${C_OFF}"
fi

log_echo "\n────────────────────────────────────────────────────────────"
log_echo " Test Complete."
log_echo "────────────────────────────────────────────────────────────"


# ==============================================================================
# Move Log to User Home Directory
# ==============================================================================
if [ -n "$SUDO_USER" ]; then
    REAL_USER="$SUDO_USER"
    USER_HOME=$(getent passwd "$REAL_USER" | cut -d: -f6)
else
    REAL_USER="$USER"
    USER_HOME="$HOME"
fi

# Fall back to /tmp if the home directory can't be resolved.
if [ -z "$USER_HOME" ] || [ ! -d "$USER_HOME" ]; then
    USER_HOME="/tmp"
fi

FINAL_LOG="$USER_HOME/solus_kernel_test_$(date +%Y%m%d_%H%M%S).log"

if mv "$LOG_FILE" "$FINAL_LOG" 2>/dev/null; then
    LOG_MOVED=1
    chown "$REAL_USER:" "$FINAL_LOG" 2>/dev/null
    chmod 644 "$FINAL_LOG" 2>/dev/null
    echo -e "\nLog file saved to: $FINAL_LOG"
else
    echo -e "\nCould not move log. It remains at: $LOG_FILE" >&2
fi

exit "$EXIT_CODE"
