#!/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
# screen (the virtio-gpu framebuffer) to a PNG, or a PPM if the file name
# ends in .ppm (scripts/gpu-check.sh reads its pixels). 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|out.ppm] [session.txt] [expected text...]
# session.txt: one command per line, typed key by key followed by Enter.
# In it, {BS} is Backspace, {UP} the up arrow and {^C} Ctrl+C (a line that
# is just {^C} isn't followed by Enter).
# Exits 1 if any expected text doesn't show up.
set -euo pipefail
cd "$(dirname "$0")/.."
T="${CARGO_TARGET_DIR:-target}" # stage 21: the build output (cargo's target dir)

cargo build
# Force bootimage to repack the disk image (see scripts/runner.sh for why).
touch $T/x86_64-rust_kernel/debug/rustbelt
cargo bootimage
IMG=$T/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): the tiling demo moves the
    # text console into the left window of the virtio-gpu framebuffer, a
    # CPU hog stopped with Ctrl+C (the tty turns ^C into SIGINT), then
    # stage 15's libc plumbing (threads summing on all CPUs, the epoll/poll
    # echo server over 127.0.0.1) and stage 15.1's musl prerequisites: the
    # FPU and path-syscall tests and rawabi, a no-libc Linux binary.
    # "vx{BS}er" types a typo and fixes it. A line "{SLEEP n}" types
    # nothing and waits n seconds (fputest runs a while; without it the next
    # command's echo lands in the middle of its output).
    SESSION=$(mktemp /tmp/session.XXXXXX)
    cat > "$SESSION" <<'TXT'
tile
vx{BS}er
spin 1000 > /tmp/hog.txt
{^C}
threads | grep threads
echoserver | grep echoes
fputest | grep passed
{SLEEP 3}
pathtest | grep passed
rawabi | grep -v "ok   "
TXT
    set -- "Rustbelt 0.22.0 (stage 22: arm64 bring-up) x86_64" "tile: 4 windows on 1024x768, console 61x46" \
        "^C" "(interrupted: SIGINT)" "threads: 4 threads summed 1..=40000000 = 800000020000000" \
        "echoserver: epoll: 3 clients, 9 echoes" "[threads] all checks passed" \
        "[fputest] all checks passed" "[pathtest] all checks passed" \
        "rawabi: pi ~ 3.141588" "rawabi: all checks passed"
fi
SOCK=$(mktemp -u /tmp/qemu-mon.XXXXXX)
SERIAL_LOG=$(mktemp /tmp/serial.XXXXXX)

qemu-system-x86_64 -smp 4 -cpu max -drive format=raw,file="$IMG" \
    -drive file=$T/disk.img,format=raw,if=virtio,snapshot=on \
    -drive file=$T/ext2.img,format=raw,if=virtio,snapshot=on \
    -drive file=$T/exfat.img,format=raw,if=virtio,snapshot=on \
    -drive file=$T/btrfs.img,format=raw,if=virtio,snapshot=on \
    -netdev "user,id=net0,guestfwd=tcp:10.0.2.100:80-cmd:sh scripts/http-reply.sh" \
    -device virtio-net-pci,netdev=net0 \
    -vga none -device virtio-vga -device virtio-rng-pci \
    -device virtio-serial-pci -device virtconsole,chardev=hvc -chardev null,id=hvc \
    -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', '%': 'shift-5', '{BS}': 'backspace', '{UP}': 'up',
         '{^C}': 'ctrl-c', '\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):
    m = re.fullmatch(r'\{SLEEP ([0-9.]+)\}', line.strip())
    if m:
        time.sleep(float(m.group(1))); continue   # wait, type nothing
    if line.strip() == '{^C}':
        line = '{^C}'   # Ctrl+C alone, no Enter after it
    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(2.5)   # let the last command finish
fmt = 'ppm' if out.endswith('.ppm') else 'png'
s.sendall(f"screendump {out} -f {fmt}\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"