#!/usr/bin/env bash
# Validate USB mass-storage autospawn via xHCI in QEMU.

set -euo pipefail

seed_usb_image() {
    local image_path="$1"
    python3 - "$image_path" <<'PY'
import base64
import pathlib
import sys

path = pathlib.Path(sys.argv[1])
payload = (b"REDBEAR-USB-STORAGE-CHECK\0" * 32)[:512]
payload = payload.ljust(512, b'\0')

with path.open("r+b") as fh:
    fh.seek(0)
    fh.write(payload)

print(base64.b64encode(payload).decode("ascii"))
PY
}

find_uefi_firmware() {
    local candidates=(
        "/usr/share/ovmf/x64/OVMF.4m.fd"
        "/usr/share/OVMF/x64/OVMF.4m.fd"
        "/usr/share/ovmf/x64/OVMF_CODE.4m.fd"
        "/usr/share/OVMF/x64/OVMF_CODE.4m.fd"
        "/usr/share/ovmf/OVMF.fd"
        "/usr/share/OVMF/OVMF_CODE.fd"
        "/usr/share/qemu/edk2-x86_64-code.fd"
    )
    local path
    for path in "${candidates[@]}"; do
        if [[ -f "$path" ]]; then
            printf '%s\n' "$path"
            return 0
        fi
    done
    return 1
}

usage() {
    cat <<'USAGE'
Usage: test-usb-storage-qemu.sh [config]

Boot a Red Bear image with a USB storage device attached and verify usbscsid autospawn.
Defaults to redbear-full.
USAGE
}

for arg in "$@"; do
    case "$arg" in
        --help|-h|help)
            usage
            exit 0
            ;;
    esac
done

config="${1:-redbear-full}"
arch="${ARCH:-$(uname -m)}"
# Boot media: prefer the live ISO (canonical build-redbear.sh artifact).
# harddrive.img is only produced by older make-all flows and goes stale —
# booting it tests the wrong tree (this happened in practice with a stale
# bootloader build). extra.img and usb-storage.img stay as data drives.
iso="build/$arch/$config.iso"
image="build/$arch/$config/harddrive.img"
extra="build/$arch/$config/extra.img"
usb_img="build/$arch/$config/usb-storage.img"
log_file="build/$arch/$config/usb-storage-check.log"
firmware="$(find_uefi_firmware)" || {
    echo "ERROR: no usable x86_64 UEFI firmware found" >&2
    exit 1
}

boot_media=""
if [[ -f "$iso" ]]; then
    boot_media="iso"
elif [[ -f "$image" ]]; then
    boot_media="harddrive"
else
    echo "ERROR: no bootable image found ($iso or $image)" >&2
    echo "Build it first with: ./local/scripts/build-redbear.sh $config" >&2
    exit 1
fi

if [[ ! -f "$extra" ]]; then
    truncate -s 1g "$extra"
fi

if [[ ! -f "$usb_img" ]]; then
    truncate -s 64M "$usb_img"
fi

# Boot-disk args: ISO via -cdrom, or legacy harddrive via nvme. Data drives
# (extra, usb-storage) are always attached the same way.
if [[ "$boot_media" == "iso" ]]; then
    boot_args=(-cdrom "$iso")
else
    boot_args=(-drive file="$image",format=raw,if=none,id=drv0,snapshot=on -device nvme,drive=drv0,serial=NVME_SERIAL)
fi

expected_sector_b64="$(seed_usb_image "$usb_img")"

pkill -f "qemu-system-x86_64.*$config" 2>/dev/null || true
sleep 1

rm -f "$log_file"
attempt=1
while true; do
    rc=0
    python3 local/scripts/qemu-login-expect.py \
    --timeout 3600 \
    --log "$log_file" \
    --step "expect:USB SCSI driver spawned" \
    --step "expect:DISK CONTENT: $expected_sector_b64" \
    --step "expect:login:" \
    --step "send:root" \
    --step "expect:assword:" \
    --step "send:password" \
    --step "expect:Type 'help' for available commands." \
    --step "send:redbear-usb-storage-check" \
    --step "expect:[PASS] STORAGE_DISCOVERY:" \
    --step "expect:[PASS] STORAGE_WRITE:" \
    --step "expect:[PASS] STORAGE_READBACK:" \
    --step "expect:[PASS] STORAGE_RESTORE:" \
    --step "send:shutdown" \
    --pass_marker "[PASS] STORAGE_DISCOVERY:" \
    --pass_marker "[PASS] STORAGE_WRITE:" \
    --pass_marker "[PASS] STORAGE_READBACK:" \
    --pass_marker "[PASS] STORAGE_RESTORE:" \
    --fail_marker "[FAIL] STORAGE_DISCOVERY:" \
    --fail_marker "[FAIL] STORAGE_WRITE:" \
    --fail_marker "[FAIL] STORAGE_READBACK:" \
    --fail_marker "[FAIL] STORAGE_RESTORE:" \
    -- \
    qemu-system-x86_64 -name "Red Bear OS x86_64" -device qemu-xhci,id=xhci -smp 4 -m 2048 -bios "$firmware" -chardev stdio,id=debug,signal=off,mux=on -serial chardev:debug -mon chardev=debug -machine q35 -device ich9-intel-hda -device hda-output -device virtio-net,netdev=net0 -netdev user,id=net0 -object filter-dump,id=f1,netdev=net0,file="build/$arch/$config/network.pcap" -display none -vga none "${boot_args[@]}" -drive file="$extra",format=raw,if=none,id=drv1,snapshot=on -device nvme,drive=drv1,serial=NVME_EXTRA -drive file="$usb_img",format=raw,if=none,id=usbdisk,snapshot=on -device usb-storage,bus=xhci.0,drive=usbdisk -enable-kvm -cpu host || rc=$?
    if [[ $rc -eq 0 ]]; then
        break
    fi
    if [[ ( $rc -eq 3 || $rc -eq 137 ) && $attempt -lt 5 ]]; then
        attempt=$((attempt + 1))
        echo "QEMU exited before completing the check; retrying (attempt $attempt/5)"
        sleep 2
        continue
    fi
    exit "$rc"
done

pkill -f "qemu-system-x86_64.*$image" 2>/dev/null || true

if grep -q "panic\|usbscsid: .*IO ERROR\|usbscsid: startup failed\|usbscsid: event queue error\|usbscsid: scheme tick failed\|bulk .* endpoint stalled" "$log_file"; then
    echo "ERROR: USB storage path hit a crash/error; see $log_file" >&2
    exit 1
fi

if grep -q "\[FAIL\]" "$log_file"; then
    echo "ERROR: USB storage read/write check failed; see $log_file" >&2
    exit 1
fi

if ! grep -q "\[PASS\] STORAGE_WRITE:" "$log_file"; then
    echo "ERROR: USB storage write proof not found in log; see $log_file" >&2
    exit 1
fi

if ! grep -q "\[PASS\] STORAGE_READBACK:" "$log_file"; then
    echo "ERROR: USB storage readback proof not found in log; see $log_file" >&2
    exit 1
fi

echo "USB mass-storage write + readback proof verified in $log_file"