#!/usr/bin/env bash
# P19-1: Multi-Core Driver Stress Test for Red Bear OS
#
# Validates SMP stability under parallel I/O load:
#   - 4 CPU cores exercising scheduler (P17)
#   - Parallel disk I/O via dd (filesystem scheme)
#   - Parallel scheme reads (IPC paths)
#   - Concurrent process creation/teardown
#   - Panic/hang detection
#
# Usage:
#   ./local/scripts/test-smp-stress-qemu.sh [--check]
#   ./local/scripts/test-smp-stress-qemu.sh --duration 60
#
# Options:
#   --check      Run full QEMU test (default)
#   --duration N Stress duration in seconds (default: 30)
#   --smp N      Number of CPU cores (default: 4)
#   --help       Show this help

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"

# Defaults
DURATION=30
SMP=4
ACTION="check"
IMAGE=""
FIRMWARE="/usr/share/edk2/x64/OVMF_CODE.4m.fd"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --check)    ACTION="check"; shift ;;
        --duration) DURATION="$2"; shift 2 ;;
        --smp)      SMP="$2"; shift 2 ;;
        --image)    IMAGE="$2"; shift 2 ;;
        --help)
            head -20 "$0" | grep '^#' | sed 's/^# \?//'
            exit 0
            ;;
        *) echo "Unknown argument: $1"; exit 1 ;;
    esac
done

# Resolve image
if [[ -z "$IMAGE" ]]; then
    IMAGE="$PROJECT_DIR/build/x86_64/redbear-mini.iso"
fi

if [[ ! -f "$IMAGE" ]]; then
    echo "ERROR: Image not found: $IMAGE"
    echo "Run: make live CONFIG_NAME=redbear-mini"
    exit 1
fi

if [[ ! -f "$FIRMWARE" ]]; then
    echo "ERROR: OVMF firmware not found: $FIRMWARE"
    echo "Install: sudo pacman -S edk2-ovmf (or equivalent)"
    exit 1
fi

echo "=== P19-1: Multi-Core Driver Stress Test ==="
echo "Image:    $IMAGE"
echo "Firmware: $FIRMWARE"
echo "CPUs:     $SMP"
echo "Duration: ${DURATION}s"
echo ""

# Create expect-based test script
export DURATION SMP IMAGE FIRMWARE

expect <<'EXPECT_SCRIPT'
log_user 1
set timeout 300
set duration $env(DURATION)
set smp $env(SMP)
set image $env(IMAGE)
set firmware $env(FIRMWARE)

proc expect_or_fail {pattern description} {
    expect {
        -nocase -re $pattern { return }
        timeout {
            puts stderr "\nERROR: timed out waiting for $description"
            exit 1
        }
        eof {
            puts stderr "\nERROR: QEMU exited while waiting for $description"
            exit 1
        }
    }
}

# Launch QEMU with 4 CPUs, network, serial console
spawn qemu-system-x86_64 \
    -name "Red Bear OS SMP Stress Test" \
    -machine q35 \
    -smp $smp \
    -m 2048 \
    -drive if=pflash,format=raw,readonly=on,file=$firmware \
    -chardev stdio,id=debug,signal=off,mux=on \
    -serial chardev:debug \
    -mon chardev=debug \
    -nographic \
    -vga none \
    -device virtio-net-pci,netdev=net0 \
    -netdev user,id=net0 \
    -drive file=$image,format=raw,if=none,id=drv0,snapshot=on \
    -device virtio-blk-pci,drive=drv0 \
    -enable-kvm \
    -cpu host

# Wait for login prompt
puts "\n>>> Waiting for boot..."
expect_or_fail "login:" "login prompt"

# Login as root
send "root\r"
expect_or_fail "assword:" "password prompt"
send "password\r"
expect_or_fail "#|\\$" "shell prompt"
sleep 1

# --- Phase 1: Verify multi-core boot ---
puts "\n>>> Phase 1: Verifying multi-core boot..."
send "cat /scheme/sys/cpu.log 2>/dev/null || echo 'no cpu.log'\r"
expect_or_fail "#|\\$" "cpu.log output"
sleep 1

send "ls /scheme/sys/ 2>/dev/null | head -20\r"
expect_or_fail "#|\\$" "sys scheme listing"
sleep 1

# --- Phase 2: Parallel disk I/O stress ---
puts "\n>>> Phase 2: Starting parallel disk I/O stress (duration: ${duration}s)..."

# Launch 4 parallel dd writers to tmpfs (exercises filesystem scheme + IPC)
send "dd if=/dev/zero of=/tmp/stress-a bs=1M count=50 2>/tmp/stress-a.err &\r"
expect_or_fail "#|\\$" "dd writer A start"
sleep 0.5

send "dd if=/dev/zero of=/tmp/stress-b bs=1M count=50 2>/tmp/stress-b.err &\r"
expect_or_fail "#|\\$" "dd writer B start"
sleep 0.5

send "dd if=/dev/zero of=/tmp/stress-c bs=512 count=100000 2>/tmp/stress-c.err &\r"
expect_or_fail "#|\\$" "dd writer C start"
sleep 0.5

send "dd if=/dev/zero of=/tmp/stress-d bs=4k count=50000 2>/tmp/stress-d.err &\r"
expect_or_fail "#|\\$" "dd writer D start"
sleep 0.5

# --- Phase 3: Parallel reads (exercises IPC) ---
puts "\n>>> Phase 3: Starting parallel filesystem traversal..."

send "ls -laR / > /tmp/ls-out 2>/tmp/ls-err &\r"
expect_or_fail "#|\\$" "ls traversal start"
sleep 0.5

send "cat /etc/passwd > /dev/null 2>&1 &\r"
expect_or_fail "#|\\$" "passwd read start"
sleep 0.5

send "find /tmp -type f > /dev/null 2>&1 &\r"
expect_or_fail "#|\\$" "find traversal start"
sleep 0.5

# --- Phase 4: Wait for stress duration ---
puts "\n>>> Phase 4: Running stress for ${duration}s..."
sleep $duration

# --- Phase 5: Collect results ---
puts "\n>>> Phase 5: Collecting results..."

# Wait for background jobs and check status
send "wait 2>/dev/null; echo STRESS-WAIT-DONE\r"
expect_or_fail "STRESS-WAIT-DONE" "background jobs completion"
sleep 1

# Check if stress files were created
send "ls -la /tmp/stress-* 2>/dev/null; echo STRESS-FILES-DONE\r"
expect_or_fail "STRESS-FILES-DONE" "stress file listing"
sleep 1

# Verify reads work on stress files (read-back test)
send "dd if=/tmp/stress-a of=/dev/null bs=1M count=50 2>/tmp/readback-a.err; echo READBACK-A-STATUS=$?\r"
expect_or_fail "READBACK-A-STATUS=0" "readback A"
sleep 1

# --- Phase 6: Panic detection ---
puts "\n>>> Phase 6: Checking for panics..."

send "dmesg 2>/dev/null | grep -ic PANIC; echo PANIC-CHECK-DONE\r"
expect {
    -re "(\[0-9\]+)\r\n.*PANIC-CHECK-DONE" {
        set panic_count $expect_out(1,string)
        if {$panic_count > 0} {
            puts "\n>>> FAIL: $panic_count PANIC(S) detected in dmesg!"
            puts "\n>>> Dumping panic context:"
            send "dmesg | grep -i PANIC\r"
            expect "#|\\$"
            exit 1
        } else {
            puts "\n>>> PASS: No panics detected in dmesg"
        }
    }
    timeout {
        puts stderr "\nERROR: timed out during panic check"
        exit 1
    }
}
sleep 1

# --- Phase 7: Summary ---
puts "\n>>> Phase 7: Final system state..."
send "echo \"UPTIME:\"; uptime 2>/dev/null || cat /scheme/sys/time 2>/dev/null; echo UPTIME-DONE\r"
expect_or_fail "UPTIME-DONE" "uptime output"
sleep 1

send "echo \"PROCS:\"; ps 2>/dev/null || echo 'ps not available'; echo PROCS-DONE\r"
expect_or_fail "PROCS-DONE" "process listing"
sleep 1

# Clean shutdown
puts "\n>>> Shutting down..."
send "shutdown\r"
sleep 5
send "\r"
sleep 2

puts "\n"
puts "============================================"
puts "  P19-1 SMP STRESS TEST: PASSED"
puts "  CPUs: $smp | Duration: ${duration}s"
puts "  No panics, no hangs, I/O completed"
puts "============================================"

expect eof
EXPECT_SCRIPT

exit_code=$?

if [[ $exit_code -eq 0 ]]; then
    echo ""
    echo "P19-1: PASSED"
else
    echo ""
    echo "P19-1: FAILED (exit code: $exit_code)"
fi

exit $exit_code