#!/usr/bin/env bash
# Boot the normal (interactive) kernel with no window, type a shell session
# into it through the QEMU monitor's `sendkey` command, then dump the VGA
# screen to a PNG. Doubles as an automated test of the keyboard, the raw
# console mode and the shell's line editor (Backspace, up-arrow history):
# afterwards the serial copy of the console is checked for expected text.
#
# Usage: ./scripts/screenshot.sh [out.png] [session.txt] [expected text...]
# session.txt: one command per line, typed key by key followed by Enter.
# In it, {BS} is Backspace and {UP} is the up arrow.
# Exits 1 if any expected text doesn't show up.
set -euo pipefail
cd "$(dirname "$0")/.."

cargo build
# Force bootimage to repack the disk image (see scripts/runner.sh for why).
touch target/x86_64-rust_kernel/debug/rustbelt
cargo bootimage
IMG=target/x86_64-rust_kernel/debug/bootimage-rustbelt.bin
OUT="${1:-$PWD/boot.png}"
SESSION="${2:-}"
shift $(( $# > 2 ? 2 : $# ))
if [ -z "$SESSION" ]; then
    # The default session (what boot.png shows): two CPU hogs in the
    # background (output to files), sysinfo showing them on two of the four
    # CPUs, and a stack overflow caught by the guard zone. "vx{BS}er" types
    # a typo and fixes it.
    SESSION=$(mktemp /tmp/session.XXXXXX)
    cat > "$SESSION" <<'TXT'
vx{BS}er
spin 150 > A:/tmp/spin1.txt &
spin 150 > A:/tmp/spin2.txt &
sysinfo
overflow
TXT
    set -- "Rustbelt 0.11.0 (stage 11: core polish) x86_64" "4 CPUs online (this program ran on CPU" \
        "CPU load over the last quarter second:" "killed: stack overflow"
fi
SOCK=$(mktemp -u /tmp/qemu-mon.XXXXXX)
SERIAL_LOG=$(mktemp /tmp/serial.XXXXXX)

qemu-system-x86_64 -smp 4 -drive format=raw,file="$IMG" \
    -drive file=target/disk.img,format=raw,if=virtio,snapshot=on \
    -drive file=target/ext2.img,format=raw,if=virtio,snapshot=on \
    -drive file=target/exfat.img,format=raw,if=virtio,snapshot=on \
    -drive file=target/btrfs.img,format=raw,if=virtio,snapshot=on \
    -serial file:"$SERIAL_LOG" -display none -monitor unix:"$SOCK",server,nowait &
QPID=$!
sleep 4   # boot until the shell prompt

# Talk to the QEMU monitor over its UNIX socket (python avoids needing socat).
python3 - "$SOCK" "$OUT" "$SESSION" <<'PY'
import re, socket, sys, time
sock, out, session = sys.argv[1:]
# QEMU key names for the characters we might type (US keyboard layout).
names = {' ': 'spc', ',': 'comma', '.': 'dot', '/': 'slash', '-': 'minus', '=': 'equal',
         ':': 'shift-semicolon', '|': 'shift-backslash', '>': 'shift-dot', '<': 'shift-comma',
         '"': 'shift-apostrophe', "'": 'apostrophe', '_': 'shift-minus', '!': 'shift-1',
         '#': 'shift-3', '&': 'shift-7', '{BS}': 'backspace', '{UP}': 'up', '\n': 'ret'}
def key(c):
    if c in names: return names[c]
    if c.isupper(): return 'shift-' + c.lower()
    return c
s = socket.socket(socket.AF_UNIX); s.connect(sock)
for line in open(session):
    for c in re.findall(r'\{[A-Z]+\}|.|\n', line, re.S):
        s.sendall(f"sendkey {key(c)}\n".encode()); time.sleep(0.08)
    time.sleep(0.8)   # let the command run
time.sleep(0.5)
s.sendall(f"screendump {out} -f png\n".encode()); time.sleep(1)
s.sendall(b"quit\n"); time.sleep(0.5)
PY
wait $QPID || true
echo "---- serial output ----"
cat -v "$SERIAL_LOG"; echo   # cat -v shows Backspace as ^H
echo "-----------------------"
# Replay the serial output like a terminal would (\b = erase one character)
# and check the expected text is there.
if python3 - "$SERIAL_LOG" "$@" <<'PY'
import sys
data = open(sys.argv[1], 'rb').read().decode(errors='replace').replace('\x08 \x08', '\x08')
out = []
for ch in data:
    if ch == '\x08':
        if out and out[-1] != '\n': out.pop()
    else:
        out.append(ch)
text = ''.join(out)
missing = [e for e in sys.argv[2:] if e not in text]
for e in missing: print(f"[screenshot] MISSING: {e}")
sys.exit(1 if missing else 0)
PY
then
    echo "[screenshot] session OK: all $# expected texts found"
else
    rm -f "$SERIAL_LOG"
    exit 1
fi
rm -f "$SERIAL_LOG"
echo "[screenshot] saved $OUT"