#!/usr/bin/env python3
# preview: boot dist/udt in real GRUB under QEMU, before root installs it.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License v2 only.
"""Boot the built theme in GRUB x86_64-efi under QEMU and OVMF.

Usage: preview [--live] [WIDTHxHEIGHT]

Default: screenshot to dist/preview.png, headless. --live opens a QEMU
window instead, to move through the menu and try the console ('c').
The mode defaults to 2560x1080. The menu entries are stand-ins.
"""
import shutil
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent
THEME = REPO / "dist" / "udt"
SHOT = REPO / "dist" / "preview.png"
OVMF = Path("/usr/share/OVMF")
MODULES = ("part_gpt part_msdos fat normal configfile all_video efi_gop gfxterm "
           "gfxmenu png jpeg font true").split()
BOOT_WAIT = 12  # ponytail: fixed wait for OVMF + GRUB; raise it if the shot is blank
CFG = """\
insmod all_video
insmod gfxterm
set gfxmode={mode},auto
terminal_output gfxterm
{fonts}
set theme=($root)/boot/grub/themes/udt/theme.txt
set timeout=30
menuentry 'Slackware (preview)' --class slackware --class os {{ true }}
submenu 'Advanced options for Slackware' {{ menuentry 'Recovery' {{ true }} }}
menuentry 'UEFI Firmware Settings' --class efi {{ true }}
"""


def main():
    args = sys.argv[1:]
    live = "--live" in args
    args = [a for a in args if a != "--live"]
    mode = args[0] if args else "2560x1080"
    xres, yres = mode.split("x")
    if not (THEME / "theme.txt").exists():
        sys.exit("preview: dist/udt missing; run ./build first")
    with tempfile.TemporaryDirectory() as tmp:
        tmp = Path(tmp)
        root = tmp / "esp"
        grub = root / "boot/grub"
        (root / "EFI/BOOT").mkdir(parents=True)
        shutil.copytree(THEME, grub / "themes/udt")
        # grub-mkconfig writes one loadfont per .pf2 in the theme dir; do the same.
        fonts = "\n".join(f"loadfont ($root)/boot/grub/themes/udt/{f.name}"
                          for f in sorted(THEME.glob("*.pf2")))
        (grub / "grub.cfg").write_text(CFG.format(mode=mode, fonts=fonts))
        subprocess.run(["grub-mkimage", "-O", "x86_64-efi", "-p", "/boot/grub",
                        "-o", str(root / "EFI/BOOT/BOOTX64.EFI"), *MODULES], check=True)
        vars_fd = tmp / "vars.fd"
        vars_fd.write_bytes((OVMF / "OVMF_VARS.fd").read_bytes())
        mon = tmp / "mon.sock"
        cmd = ["qemu-system-x86_64", "-machine", "q35", "-m", "512",
               "-drive", f"if=pflash,format=raw,readonly=on,file={OVMF / 'OVMF_CODE.fd'}",
               "-drive", f"if=pflash,format=raw,file={vars_fd}",
               "-drive", f"file=fat:rw:{root},format=raw",
               "-vga", "none", "-device", f"VGA,edid=on,xres={xres},yres={yres}"]
        if live:
            return subprocess.run(cmd).returncode
        q = subprocess.Popen(cmd + ["-display", "none", "-monitor", f"unix:{mon},server,nowait"])
        try:
            time.sleep(BOOT_WAIT)
            s = socket.socket(socket.AF_UNIX)
            s.connect(str(mon))
            ppm = tmp / "shot.ppm"
            s.sendall(f"screendump {ppm}\n".encode())
            for _ in range(50):  # the dump is written asynchronously
                if ppm.exists() and ppm.stat().st_size:
                    break
                time.sleep(0.2)
            time.sleep(0.5)
            s.sendall(b"quit\n")
            q.wait(timeout=10)
        finally:
            if q.poll() is None:
                q.kill()
        subprocess.run(["magick", str(ppm), str(SHOT)], check=True)
    print(f"preview: {SHOT}")


if __name__ == "__main__":
    main()