#!/usr/bin/env python3
"""qemu_session.py -- one typed QEMU session against a kernel boot, and its verdict.

    scripts/qemu_session.py MODE ARCH [--label L] [--make "VAR=val ..."] [--no-zuse]

Builds and boots `make ARCH=<arch> clean qemu` (one QEMU, in the foreground of
this script), with copies of the minted Zuse pair (disk/zuse-scratch/) made
under build/qemu-session/<label>/ so the committed images are never booted.
It reads the serial socket continuously, drives the boot over the serial and
QMP sockets the `qemu` goal opens, prints PASS:/FAIL: lines and
RESULT: PASS|FAIL, and quits QEMU over QMP. The serial log is the usual
logs/<TS>/<arch>/qemu-<arch>-<TS>.log.

Modes (the identity gate's, plan 2026-10-26-engine-step4-identity-gate.md, and
the stick's, plan 2026-10-26-stick-1-screen-and-keyboard.md):
  headless   no Zuse drive: POST: PASSED, then no `ok> ` for 60 s
  zuse       the minted pair: identity confirmed, [zuse@Hera] ok>, `3 4 + .` -> 7  ok
  hotplug    no drive; Zuse added over QMP -> prompt; pulled -> session ended,
             headless; added again -> prompt
  midpull    pulled while a line is in hand: what is typed after does not run
  waitpull   pulled while a line WAITS: the wait ends at logout; re-attach -> prompt
  typeahead  serial bytes typed while headless do not run once Zuse attaches
  emergency  an EMERGENCY_CONSOLE_ENABLED=1 build: ok> with no drive
  v3_keys    system=v3: `3 4 + .` typed on QEMU's keyboard (send-key) -> 7  ok
  screen_dark_until_login     no drive: 30 s after POST the screen is all black
  screen_console_after_login  Zuse added: the screen shows the console; keys typed
                              on QEMU's keyboard answer 7  ok and change the screen
  screen_dark_after_logout    ... and when her drive is pulled, the screen is black again
  kbd_typeahead               keys pressed while headless do not run once Zuse attaches
  usb_kbd                     the only keyboard a USB one (--make "QEMU_EXTRA='-machine
                              i8042=off -device usb-kbd,bus=xhci0.0,port=3'"): it types
  no_framebuffer              no display (-vga none): serial alone, 7  ok
  alttab_after_logout         after logout, Alt+TAB twice: the screen stays black
  alttab_before_pull          Alt+TAB, then the drive pulled: black; re-attach draws the console
  emergency_pull              EMERGENCY_CONSOLE_ENABLED=1: Zuse added and pulled, the screen
                              still shows the running console, and it types
  gpt_zuse     a GPT drive, Zuse in its StarshipOS identity partition: recognised; the partition
               is attached (run with --make KERNEL_ARGS=--log-level=info)
  gpt_hotplug  ... pulled: session ended; added again: prompt
  gpt_foreign  a GPT drive with only a Linux partition: no identity, not attached (info level)
  gpt_bad_crc  the gpt_zuse image with its header CRC broken: not used, said so, not attached
  virtio_kbd                  aarch64/riscv64: QEMU's virtio keyboard types 7  ok

Needs only Python 3. One QEMU at a time (.claude/CLAUDE.md).
"""
import argparse, glob, json, os, re, shlex, shutil, socket, subprocess, sys, threading, time

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
PROMPT = r"\[zuse@Hera\][^\n]*ok> "
ANYPROMPT = r"ok> "

ap = argparse.ArgumentParser()
ap.add_argument("mode")
ap.add_argument("arch")
ap.add_argument("--label", default=None)
ap.add_argument("--make", default="")
ap.add_argument("--no-zuse", action="store_true")
a = ap.parse_args()

NO_ZUSE_MODES = {"gpt_zuse", "gpt_hotplug", "gpt_foreign", "gpt_bad_crc", "headless", "hotplug", "midpull", "waitpull", "typeahead", "emergency",
                 "alttab_after_logout", "alttab_before_pull", "emergency_pull",
                 "screen_dark_until_login", "screen_console_after_login", "screen_dark_after_logout",
                 "kbd_typeahead"}
label = a.label or f"{a.mode}-{a.arch}"
work = os.path.join(ROOT, "build", "qemu-session", label)
shutil.rmtree(work, ignore_errors=True)
os.makedirs(work)
for f in ("artemis.img", "zuse-thumb-ident.img"):
    subprocess.run(["cp", "--sparse=always", os.path.join(ROOT, "disk/zuse-scratch", f), work], check=True)
art = os.path.join(work, "artemis.img")
zimg = os.path.join(work, "zuse-thumb-ident.img")
attach_zuse = not (a.no_zuse or a.mode in NO_ZUSE_MODES)

make_args = ["make", "-f", "kernel/Makefile", f"ARCH={a.arch}", "clean", "qemu", f"ARTDISK={art}",
             f"ZUSEDISK={zimg if attach_zuse else ''}"] + shlex.split(a.make)
if a.mode == "v3_keys":
    make_args.append("KERNEL_ARGS=system=v3")
start = time.time()
make_log = open(os.path.join(work, "make.txt"), "w")
mk = subprocess.Popen(make_args, cwd=ROOT, stdout=make_log, stderr=subprocess.STDOUT)
deadline = time.time() + 1800
ok = True


def newest(pattern):
    c = [p for p in glob.glob(pattern) if os.path.getmtime(p) >= start]
    return max(c, key=os.path.getmtime) if c else None


log = ser = qmp = None
while time.time() < deadline and not (log and ser and qmp):
    if mk.poll() is not None:
        print(f"MAKE FAILED (exit {mk.returncode}); see {make_log.name}", flush=True)
        sys.exit(1)
    log = log or newest(f"{ROOT}/logs/*/{a.arch}/qemu-{a.arch}-*.log")
    ser = ser or newest(f"{ROOT}/build/{a.arch}/kernel/qemu-serial-{a.arch}-*.sock")
    qmp = qmp or newest(f"{ROOT}/build/{a.arch}/kernel/qemu-qmp-{a.arch}-*.sock")
    time.sleep(1)
print(f"LOG: {os.path.relpath(log, ROOT)}", flush=True)

s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
for _ in range(60):
    try:
        s.connect(ser)
        break
    except OSError:
        time.sleep(1)


def drain():
    while True:
        try:
            if not s.recv(4096):
                return
        except OSError:
            return


threading.Thread(target=drain, daemon=True).start()

q = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
q.connect(qmp)
qf = q.makefile("r")
qf.readline()


def qcmd(cmd, **args):
    msg = {"execute": cmd}
    if args:
        msg["arguments"] = args
    q.send(json.dumps(msg).encode())
    while True:
        r = json.loads(qf.readline())
        if "return" in r or "error" in r:
            if "error" in r:
                print(f"QMP {cmd}: {json.dumps(r)}", flush=True)
            return r


qcmd("qmp_capabilities")


def text():
    with open(log, errors="replace") as f:
        return f.read()


def count(pat):
    return len(re.findall(pat, text()))


def wait_for(pat, secs, after=0):
    t = time.time()
    while time.time() - t < secs:
        if count(pat) > after:
            return True
        if "POST: FAILED" in text() or "PARITY:FAIL" in text():
            return False
        time.sleep(0.5)
    return False


def check(name, cond):
    global ok
    print(f"{'PASS' if cond else 'FAIL'}: {name}", flush=True)
    ok = ok and cond


def type_line(line):
    """Over the serial line, a character at a time."""
    for ch in line:
        s.send(ch.encode())
        time.sleep(0.05)
    s.send(b"\r")


# QEMU qcodes for what the tests type on a keyboard
QCODE = {" ": ["spc"], ".": ["dot"], "+": ["shift", "equal"], "-": ["minus"], "\n": ["ret"]}


def key_line(line):
    """On QEMU's keyboard (send-key), a key at a time."""
    for ch in line + "\n":
        keys = QCODE.get(ch, [ch.lower()])
        qcmd("send-key", keys=[{"type": "qcode", "data": k} for k in keys])
        time.sleep(0.1)


# The two StarshipOS partition type GUIDs (docs/v4.0.0/plans/2026-10-26-stick-2-gpt-reader.md)
GPT_IDENTITY = "23BEBFE1-4817-4CF3-A537-9069BF159788"
GPT_LINUX = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"


def make_gpt_image(path, ptype, payload, mib=72, break_crc=False):
    """A GPT disk image: one partition at LBA 2048, 131072 sectors (64 MiB), of type
    `ptype`, holding `payload` from its start; optionally the primary header's CRC broken."""
    subprocess.run(["truncate", "-s", f"{mib}M", path], check=True)
    sf = shutil.which("sfdisk") or "/usr/sbin/sfdisk"
    subprocess.run([sf, "--quiet", path], input=f"label: gpt\nstart=2048, size=131072, type={ptype}\n",
                   text=True, check=True)
    subprocess.run(["dd", f"if={payload}", f"of={path}", "bs=512", "seek=2048", "conv=notrunc,sparse",
                    "status=none"], check=True)
    if break_crc:
        with open(path, "r+b") as f:
            f.seek(512 + 16)                         # the primary header's HeaderCRC32
            b = f.read(1)
            f.seek(512 + 16)
            f.write(bytes([b[0] ^ 0xFF]))


def blk_lines():
    return count(r"blk: ")


def alt_tab():
    qcmd("send-key", keys=[{"type": "qcode", "data": "alt"}, {"type": "qcode", "data": "tab"}])
    time.sleep(1)


SHOTS = 0


def screen():
    """A screendump of QEMU's display: (width, height, number of non-black pixels, bytes)."""
    global SHOTS
    SHOTS += 1
    path = os.path.join(work, f"screen-{SHOTS}.ppm")
    qcmd("screendump", filename=path)
    for _ in range(20):
        if os.path.exists(path) and os.path.getsize(path) > 0:
            break
        time.sleep(0.5)
    with open(path, "rb") as f:
        data = f.read()
    parts = data.split(b"\n", 3)                      # P6, "w h", maxval, pixels
    w, h = (int(x) for x in parts[1].split())
    px = parts[3]
    lit = sum(1 for i in range(0, len(px) - 2, 3) if px[i] or px[i + 1] or px[i + 2])
    return w, h, lit, px


def add_zuse(n, img=None):
    qcmd("blockdev-add", driver="raw", **{"node-name": f"zuse{n}"},
         file={"driver": "file", "filename": img or zimg})
    qcmd("device_add", driver="usb-storage", bus="xhci0.0", port="1",
         drive=f"zuse{n}", id=f"zusedev{n}")


def del_zuse(n):
    qcmd("device_del", id=f"zusedev{n}")
    time.sleep(3)
    qcmd("blockdev-del", **{"node-name": f"zuse{n}"})


check("POST: PASSED", wait_for(r"POST: PASSED", 900))
m = a.mode

if m == "headless":
    check("no ok> for 60 s after POST", not wait_for(ANYPROMPT, 60))

elif m == "zuse":
    check("Zuse: identity confirmed", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120))
    check("[zuse@Hera] ok>", wait_for(PROMPT, 60))
    n = count(PROMPT)
    type_line("3 4 + .")
    check("7  ok, and a prompt", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m == "hotplug":
    check("no ok> for 20 s with no drive", not wait_for(ANYPROMPT, 20))
    add_zuse(1)
    check("identity confirmed on attach", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120))
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 60))
    n = count(PROMPT)
    type_line("3 4 + .")
    check("7  ok", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)
    del_zuse(1)
    check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
    n = count(ANYPROMPT)
    check("no new ok> for 30 s after detach", not wait_for(ANYPROMPT, 30, n))
    c = count(r"Zuse: identity confirmed from attached thumbdrive")
    p = count(PROMPT)
    add_zuse(2)
    check("identity confirmed on re-attach", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120, c))
    check("[zuse@Hera] ok> again", wait_for(PROMPT, 60, p))

elif m == "midpull":
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    qcmd("device_del", id="zusedev1")
    type_line("3 4 + .")
    type_line("5 6 + .")
    check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
    time.sleep(10)
    t = text()
    check("the line typed after the pull did not run (no 11)", re.search(r"\b11\s+ok", t) is None)
    check("no ok> after the session ended", "ok> " not in t[t.rfind("Zuse: session ended"):])

elif m == "waitpull":
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    type_line("5 AWAIT")
    time.sleep(5)
    del_zuse(1)
    check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
    check("the waiting line ended at logout (Interrupted)", wait_for(r"Interrupted", 30))
    p = count(PROMPT)
    add_zuse(2)
    check("identity confirmed on re-attach", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120, 1))
    check("[zuse@Hera] ok> again", wait_for(PROMPT, 60, p))
    n = count(PROMPT)
    type_line("3 4 + .")
    check("7  ok after re-attach", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m == "typeahead":
    check("no ok> for 10 s with no drive", not wait_for(ANYPROMPT, 10))
    type_line("3 4 + .")
    time.sleep(5)
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    time.sleep(5)
    check("what was typed while headless did not run (no 7)", re.search(r"\b7\s+ok", text()) is None)
    n = count(PROMPT)
    type_line("1 2 + .")
    check("3  ok typed at the prompt", wait_for(PROMPT, 60, n) and re.search(r"\b3\s+ok", text()) is not None)

elif m == "emergency":
    check("ok> with no drive", wait_for(ANYPROMPT, 60))
    n = count(ANYPROMPT)
    type_line("3 4 + .")
    check("7  ok", wait_for(ANYPROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m == "v3_keys":
    check("[zuse@Hera] ok>", wait_for(PROMPT, 300))
    time.sleep(2)
    n = count(PROMPT)
    key_line("3 4 + .")
    check("7  ok typed on the keyboard", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m == "screen_dark_until_login":
    time.sleep(30)
    w, h, lit, _ = screen()
    check(f"the screen ({w}x{h}) is all black before login ({lit} lit)", lit == 0)

elif m in ("screen_console_after_login", "screen_dark_after_logout"):
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    time.sleep(3)
    w, h, lit1, px1 = screen()
    check(f"the screen shows the console ({lit1} lit pixels)", lit1 >= 500)
    n = count(PROMPT)
    key_line("3 4 + .")
    check("7  ok typed on the keyboard", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)
    time.sleep(3)
    _, _, lit2, px2 = screen()
    check("the screen changed with what was typed", px2 != px1)
    if m == "screen_dark_after_logout":
        del_zuse(1)
        check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
        time.sleep(5)
        _, _, lit3, _ = screen()
        check(f"the screen is black after logout ({lit3} lit)", lit3 == 0)

elif m == "kbd_typeahead":
    check("no ok> for 10 s with no drive", not wait_for(ANYPROMPT, 10))
    key_line("3 4 + .")
    time.sleep(5)
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    time.sleep(5)
    check("keys pressed while headless did not run (no 7)", re.search(r"\b7\s+ok", text()) is None)
    n = count(PROMPT)
    key_line("1 2 + .")
    check("3  ok typed on the keyboard", wait_for(PROMPT, 60, n) and re.search(r"\b3\s+ok", text()) is not None)

elif m == "usb_kbd":
    check("[zuse@Hera] ok>", wait_for(PROMPT, 300))
    time.sleep(2)
    n = count(PROMPT)
    key_line("3 4 + .")          # the USB keyboard is the only one (-machine i8042=off)
    check("7  ok typed on the USB keyboard", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m in ("no_framebuffer", "virtio_kbd"):
    check("[zuse@Hera] ok>", wait_for(PROMPT, 300))
    time.sleep(2)
    n = count(PROMPT)
    if m == "no_framebuffer":
        type_line("3 4 + .")
    else:
        key_line("3 4 + .")
    check("7  ok", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)

elif m in ("alttab_after_logout", "alttab_before_pull"):
    add_zuse(1)
    check("[zuse@Hera] ok> after attach", wait_for(PROMPT, 120))
    n = count(PROMPT)
    key_line("3 4 + .")
    check("7  ok typed", wait_for(PROMPT, 60, n))
    if m == "alttab_before_pull":
        alt_tab()                                   # to the graphics view, her session on screen
    del_zuse(1)
    check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
    time.sleep(5)
    if m == "alttab_after_logout":
        alt_tab()
        alt_tab()
        time.sleep(3)
    _, _, lit, _ = screen()
    check(f"the screen is black after logout ({lit} lit)", lit == 0)
    if m == "alttab_before_pull":
        p = count(PROMPT)
        add_zuse(2)
        check("[zuse@Hera] ok> again", wait_for(PROMPT, 120, p))
        time.sleep(3)
        _, _, lit1, px1 = screen()
        check(f"the console is drawn again ({lit1} lit)", lit1 >= 500)
        n = count(PROMPT)
        key_line("1 2 + .")
        check("3  ok typed", wait_for(PROMPT, 60, n))
        time.sleep(3)
        _, _, _, px2 = screen()
        check("the screen changed with what was typed", px2 != px1)

elif m == "emergency_pull":
    check("ok> with no drive", wait_for(ANYPROMPT, 60))
    add_zuse(1)
    check("Zuse: identity confirmed", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120))
    time.sleep(3)
    del_zuse(1)
    check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
    time.sleep(5)
    _, _, lit1, px1 = screen()
    check(f"the screen still shows the console ({lit1} lit)", lit1 >= 500)
    n = count(ANYPROMPT)
    key_line("3 4 + .")
    check("7  ok typed after the pull", wait_for(ANYPROMPT, 60, n) and re.search(r"7\s+ok", text()) is not None)
    time.sleep(3)
    _, _, _, px2 = screen()
    check("the screen changed with what was typed", px2 != px1)

elif m in ("gpt_zuse", "gpt_hotplug", "gpt_foreign", "gpt_bad_crc"):
    gimg = os.path.join(work, "gpt.img")
    make_gpt_image(gimg, GPT_LINUX if m == "gpt_foreign" else GPT_IDENTITY, zimg,
                   break_crc=(m == "gpt_bad_crc"))
    check("no ok> for 10 s with no drive", not wait_for(ANYPROMPT, 10))
    b0 = blk_lines()
    add_zuse(1, gimg)
    if m in ("gpt_zuse", "gpt_hotplug"):
        check("Zuse: identity confirmed", wait_for(r"Zuse: identity confirmed from attached thumbdrive", 120))
        check("[zuse@Hera] ok>", wait_for(PROMPT, 60))
        # the identity partition (64 MiB), not the whole 72 MiB disk: a 64 MiB identity drive
        # shows 48724 user blocks (the raw minted drive, logs/20261026-091010); the disk, 54868
        users = re.findall(r"blk: disk .*?\((\d+) user", text())
        check(f"the partition, not the disk, was attached (blk: lines {b0} -> {blk_lines()}, "
              f"user blocks {users[-1] if users else None})",
              blk_lines() > b0 and users and users[-1] == "48724")
        n = count(PROMPT)
        type_line("3 4 + .")
        check("7  ok", wait_for(PROMPT, 60, n) and re.search(r"7\s+ok", text()[-400:]) is not None)
        if m == "gpt_hotplug":
            del_zuse(1)
            check("session ended on detach", wait_for(r"Zuse: session ended -- reattach to re-authenticate", 60))
            p = count(PROMPT)
            add_zuse(2, gimg)
            check("[zuse@Hera] ok> again", wait_for(PROMPT, 120, p))
    else:
        if m == "gpt_bad_crc":
            check("the damaged table is said so", wait_for(r"GPT: a USB drive has a damaged partition table -- not used", 60))
        check("no ok> for 30 s", not wait_for(ANYPROMPT, 30))
        check("no identity", count(r"Zuse: identity confirmed") == 0)
        check(f"not attached (blk: lines {b0} -> {blk_lines()})", blk_lines() == b0)

else:
    check(f"known mode {m}", False)

print(f"RESULT: {'PASS' if ok else 'FAIL'}", flush=True)
q.send(json.dumps({"execute": "quit"}).encode())
mk.wait()
for f in ("v3/include/version.h", "capsules/BLOCK_MAP.md"):
    subprocess.run(["git", "checkout", "-q", "--", f], cwd=ROOT)
sys.exit(0 if ok else 2)