#!/usr/bin/env python3
# montauk-imager.py - write and update MontaukOS SD cards for the Raspberry Pi 4/400
#
# Usage: ./scripts/montauk-imager.py [--image <file>]          the graphical imager
#        ./scripts/montauk-imager.py list                       list SD cards / USB drives
#        sudo ./scripts/montauk-imager.py flash  <image> <device> [--no-verify] [--expand] [--yes]
#        sudo ./scripts/montauk-imager.py update <image> <device> [--keep-config] [--firmware]
#                                                     [--prune] [--expand] [--yes]
#        sudo ./scripts/montauk-imager.py expand <device> [--yes]
#
# The image is what `make ARCH=aarch64 rpi4-sd` builds (montauk-rpi4.img, see
# scripts/mk-rpi4-sd.sh and docs/raspberry-pi.md), raw or compressed (.gz, .xz,
# .zip). The GUI runs as the user and does the work by starting this script
# again through pkexec with --machine, reading its progress from stdout.
#
# flash   Writes the image over the whole card (cancellable), wipes the end of
#         the card (a stale backup GPT from an earlier OS), reads it back to
#         verify, and rereads the partition table.
#
# update  Refreshes the installation on a card already written from an image,
#         in place: the card's partition table is never touched, so a system
#         partition that has grown to fill the card (first boot, or --expand)
#         stays as it is. Both partitions are found by their labels.
#           ESP (FAT32, MONTAUK)  every file of the image's ESP is written;
#                                 config.txt is backed up to config.txt.bak
#                                 if it differs, and the UEFI firmware
#                                 RPI_EFI.fd (which holds the firmware's own
#                                 settings) is only replaced with --firmware.
#           root (ext2, MONTAUKROOT)
#                                 os/, apps/, sdk/, boot/ and any other system
#                                 directory: new and changed files written.
#                                 config/ and www/: only files the card lacks
#                                 are added, except that the image's init.toml
#                                 and capabilities.toml are merged in (below).
#                                 users/ and state/ are never touched.
#         init.toml becomes the image's, with each service's `enabled` taken
#         from the card's copy and services only the card has appended;
#         capabilities.toml becomes the image's with grants only the card has
#         appended. The card's old copy is kept as <name>.old. --keep-config
#         leaves both alone. os/firstboot.elf is only written if the card still
#         has it: it deletes itself once the volume has grown. --prune deletes
#         files of the system directories the image no longer ships (except
#         os/btkeys.bin, the Bluetooth bonds the kernel keeps there).
#         Unchanged files are not rewritten; each file is written to a temporary
#         name, flushed and renamed over the old one.
#
# expand  Grows the system partition to the end of the card (on a 1 MiB
#         boundary, at most 1 TiB, the size mk-rpi4-sd.sh's resize_inode
#         reserves for) and the ext2 volume with resize2fs, which is what
#         0:/os/firstboot.elf otherwise does on the Pi. firstboot.elf then
#         finds nothing to grow and deletes itself. flash and update take
#         --expand to do this as a last (flash) or first (update) step.
#
# The ext2 volume is checked with e2fsck before it is changed and afterwards,
# and its feature flags are checked against what Fs/Ext2.cpp mounts (no
# incompat feature but filetype).

import argparse
import contextlib
import errno
import fcntl
import filecmp
import gzip
import hashlib
import json
import lzma
import os
import queue
import re
import shutil
import signal
import stat
import struct
import subprocess
import sys
import tempfile
import threading
import time
import zipfile

ROOT_LABEL = "MONTAUKROOT"     # RootLabel in kernel/src/Fs/Boot.cpp
ESP_LABEL = "MONTAUK"          # mformat -v in scripts/mk-rpi4-sd.sh
SECTOR = 512
MIB = 1 << 20
ALIGN = MIB // SECTOR          # partitions start and end on 1 MiB boundaries
MAX_ROOT_BYTES = (1 << 28) * 4096   # mke2fs -E resize=2^28 (4 KiB blocks); the kernel's grow cap too
CHUNK = 4 * MIB

FAT_TYPES = {0x01, 0x04, 0x06, 0x0B, 0x0C, 0x0E, 0xEF}
TYPE_ESP = 0xEF
TYPE_LINUX = 0x83

# ext2/3/4 incompat features (dumpe2fs names). Fs/Ext2.cpp refuses all but filetype.
EXT_INCOMPAT = {
    "compression", "filetype", "needs_recovery", "journal_dev", "meta_bg",
    "extent", "extents", "64bit", "mmp", "flex_bg", "ea_inode", "dirdata",
    "metadata_csum_seed", "large_dir", "inline_data", "encrypt", "casefold",
}
SUPPORTED_INCOMPAT = {"filetype"}

# Mount points that mark a disk as the host's own: never offered, never written.
SYSTEM_MOUNTS = {"/", "/boot", "/boot/efi", "/efi", "/usr", "/var", "/home", "[SWAP]"}

# Update policy for the root volume (see the header).
MERGE_DIRS = {"config", "www"}
STATE_DIRS = {"users", "state", "lost+found"}
PRUNE_DIRS = {"os", "apps", "sdk", "boot"}
PRESERVE = {"os/btkeys.bin", "os/firstboot.elf"}
FIRSTBOOT = "os/firstboot.elf"
INIT_TOML = "config/init.toml"
CAPS_TOML = "config/capabilities.toml"
FIRMWARE_FILE = "RPI_EFI.fd"
ESP_BACKUP = {"config.txt"}

BLKRRPART = 0x125F
BLKFLSBUF = 0x1261

SCRIPT = os.path.realpath(__file__)
REPO = os.path.dirname(os.path.dirname(SCRIPT))
if not os.path.isfile(os.path.join(REPO, "scripts", "mk-rpi4-sd.sh")):
    REPO = None
DEFAULT_IMAGE = os.path.join(REPO, "montauk-rpi4.img") if REPO else ""


class Fail(Exception):
    pass


def human(n):
    """Decimal units, as SD cards are labelled."""
    n = float(n)
    for unit in ("B", "kB", "MB", "GB", "TB"):
        if n < 1000 or unit == "TB":
            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1000


def tool_env():
    env = dict(os.environ)
    env["PATH"] = env.get("PATH", "/usr/bin:/bin") + ":/usr/sbin:/sbin"
    env["LC_ALL"] = "C"
    return env


# ==== Output ====

class Reporter:
    """Progress for the GUI (--machine: one '@tag text' line each) or a terminal."""

    def __init__(self, machine):
        self.machine = machine
        self.cancel = threading.Event()
        self.cancellable = False
        self._last = 0.0
        self._inline = False
        self._tty = sys.stdout.isatty()

    def _write(self, line, end="\n"):
        try:
            sys.stdout.write(line + end)
            sys.stdout.flush()
        except (BrokenPipeError, OSError):
            # The GUI went away. Finish the job anyway: an update stopped
            # halfway is worse than one nobody watches.
            devnull = os.open(os.devnull, os.O_WRONLY)
            os.dup2(devnull, sys.stdout.fileno())

    def _emit(self, tag, text):
        text = " ".join(str(text).splitlines()) if tag != "log" else str(text).rstrip("\n")
        if self.machine:
            for part in text.splitlines() or [""]:
                self._write(f"@{tag} {part}")
            return
        if self._inline:
            self._write("")
            self._inline = False
        prefix = {"stage": "==> ", "log": "    ", "warn": "warning: ",
                  "error": "error: ", "done": "==> "}[tag]
        self._write(prefix + text)

    def stage(self, text): self._emit("stage", text)
    def log(self, text): self._emit("log", text)
    def warn(self, text): self._emit("warn", text)
    def error(self, text): self._emit("error", text)
    def done(self, text): self._emit("done", text)

    def progress(self, done, total, text, force=False):
        now = time.monotonic()
        if not force and done < total and now - self._last < 0.2:
            return
        self._last = now
        if self.machine:
            self._write(f"@progress {int(done)} {int(max(total, 1))} {text}")
        elif self._tty:
            pct = 100 * done / max(total, 1)
            self._write(f"\r[{pct:5.1f}%] {text}".ljust(78), end="")
            self._inline = True

    def check_cancel(self):
        if self.cancellable and self.cancel.is_set():
            raise Fail("cancelled")


def watch_stdin(rep):
    """The GUI cancels by writing 'cancel' (or by closing the pipe)."""
    def run():
        try:
            for line in sys.stdin:
                if line.strip() == "cancel":
                    rep.cancel.set()
        except (OSError, ValueError):
            pass
        rep.cancel.set()
    threading.Thread(target=run, daemon=True).start()


def run(rep, cmd, ok=(0,), quiet=False):
    try:
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                           stdin=subprocess.DEVNULL, text=True, errors="replace", env=tool_env())
    except FileNotFoundError:
        raise Fail(f"{cmd[0]} is not installed")
    if not quiet or p.returncode not in ok:
        for line in p.stdout.splitlines():
            if line.strip():
                rep.log(line)
    if p.returncode not in ok:
        raise Fail(f"{os.path.basename(cmd[0])} failed (exit {p.returncode})")
    return p


# ==== Partition tables ====

class Part:
    def __init__(self, number, ptype, start, sectors):
        self.number, self.type, self.start, self.sectors = number, ptype, start, sectors
        self.node = None
        self.label = ""
        self.fstype = ""

    @property
    def offset(self): return self.start * SECTOR

    @property
    def size(self): return self.sectors * SECTOR

    @property
    def end(self): return self.start + self.sectors


def parse_mbr(sector0):
    """The four primary entries of an MBR, or None. Raises on a GPT disk."""
    if len(sector0) < 512 or sector0[510:512] != b"\x55\xaa":
        return None
    parts = []
    for i in range(4):
        e = sector0[446 + 16 * i: 462 + 16 * i]
        ptype = e[4]
        start, count = struct.unpack_from("<II", e, 8)
        if ptype == 0xEE:
            raise Fail("the card uses a GPT partition table, not the MBR a MontaukOS image has")
        if ptype and count:
            parts.append(Part(i + 1, ptype, start, count))
    return parts


# ==== Images ====

class Image:
    """A raw, gzip, xz or zip image, streamed without unpacking it to disk."""

    def __init__(self, path):
        self.path = os.path.abspath(path)
        if not os.path.isfile(self.path):
            raise Fail(f"image {path} not found")
        self.file_size = os.path.getsize(self.path)
        with open(self.path, "rb") as f:
            magic = f.read(6)
        if magic[:2] == b"\x1f\x8b":
            self.kind = "gzip"
        elif magic == b"\xfd7zXZ\x00":
            self.kind = "xz"
        elif magic[:4] == b"PK\x03\x04":
            self.kind = "zip"
        else:
            self.kind = "raw"

    def _zip_member(self, zf):
        infos = [i for i in zf.infolist() if not i.is_dir()]
        imgs = [i for i in infos if i.filename.lower().endswith(".img")]
        pick = imgs or infos
        if not pick:
            raise Fail("the zip file is empty")
        return max(pick, key=lambda i: i.file_size)

    def size(self):
        """The uncompressed size, or None where it is not cheap to find out."""
        if self.kind == "raw":
            return self.file_size
        if self.kind == "zip":
            with zipfile.ZipFile(self.path) as zf:
                return self._zip_member(zf).file_size
        if self.kind == "gzip" and self.file_size < (1 << 32):
            with open(self.path, "rb") as f:
                f.seek(-4, os.SEEK_END)
                return struct.unpack("<I", f.read(4))[0] or None
        return None

    @contextlib.contextmanager
    def open(self):
        """Yields (stream, raw) where raw.tell() / file_size is the progress."""
        raw = open(self.path, "rb")
        try:
            if self.kind == "raw":
                yield raw, raw
            elif self.kind == "gzip":
                with gzip.GzipFile(fileobj=raw) as s:
                    yield s, raw
            elif self.kind == "xz":
                with lzma.LZMAFile(raw) as s:
                    yield s, raw
            else:
                with zipfile.ZipFile(raw) as zf, zf.open(self._zip_member(zf)) as s:
                    yield s, raw
        finally:
            raw.close()

    def materialize(self, rep, workdir):
        """A raw image file: this one, or a decompressed copy in workdir."""
        if self.kind == "raw":
            return self.path
        out = os.path.join(workdir, "image.img")
        rep.stage(f"Decompressing {os.path.basename(self.path)}")
        with self.open() as (s, raw), open(out, "wb") as o:
            while True:
                buf = s.read(CHUNK)
                if not buf:
                    break
                o.write(buf)
                rep.progress(raw.tell(), self.file_size, "Decompressing")
        return out

    def partitions(self, raw_path):
        with open(raw_path, "rb") as f:
            parts = parse_mbr(f.read(512))
        if not parts:
            raise Fail("the image has no MBR partition table; is it a MontaukOS Raspberry Pi image?")
        return parts


# ==== Block devices ====

def disk_of(device):
    """The /sys/block name of a whole disk, or Fail."""
    real = os.path.realpath(device)
    name = os.path.basename(real)
    if not real.startswith("/dev/") or not os.path.isdir(f"/sys/block/{name}"):
        raise Fail(f"{device} is not a whole disk (give /dev/sdX or /dev/mmcblkN, not a partition)")
    if not stat.S_ISBLK(os.stat(real).st_mode):
        raise Fail(f"{device} is not a block device")
    return name


def sysfs_int(path, default=0):
    try:
        with open(path) as f:
            return int(f.read().strip())
    except (OSError, ValueError):
        return default


def disk_bytes(name):
    return sysfs_int(f"/sys/block/{name}/size") * SECTOR


def sys_partitions(name):
    """{number: Part} as the kernel currently sees the disk's partitions."""
    parts = {}
    base = f"/sys/block/{name}"
    for entry in os.listdir(base):
        num = sysfs_int(f"{base}/{entry}/partition", 0)
        if num:
            p = Part(num, 0, sysfs_int(f"{base}/{entry}/start"), sysfs_int(f"{base}/{entry}/size"))
            p.node = f"/dev/{entry}"
            parts[num] = p
    return parts


def holders(sysname):
    """Everything stacked on a block device (dm-crypt, LVM, md), recursively."""
    out = []
    for base in (f"/sys/block/{sysname}", f"/sys/class/block/{sysname}"):
        hdir = os.path.join(base, "holders")
        if os.path.isdir(hdir):
            for h in os.listdir(hdir):
                out.append(h)
                out.extend(holders(h))
            break
    return out


def disk_mounts(name):
    """[(device node, mount point)] of the disk, its partitions and their holders."""
    nodes = {f"/dev/{name}"} | {p.node for p in sys_partitions(name).values()}
    for n in list(nodes):
        nodes |= {f"/dev/{h}" for h in holders(os.path.basename(n))}
    real = {os.path.realpath(n): n for n in nodes}
    found = []
    with open("/proc/self/mounts") as f:
        for line in f:
            dev, mnt = line.split()[:2]
            mnt = mnt.replace("\\040", " ")
            if dev.startswith("/dev/") and os.path.realpath(dev) in real:
                found.append((dev, mnt))
    with open("/proc/swaps") as f:
        for line in f.readlines()[1:]:
            dev = line.split()[0]
            if os.path.realpath(dev) in real:
                found.append((dev, "[SWAP]"))
    return found


def check_target(rep, device):
    """The disk name of a device it is safe to write, or Fail."""
    name = disk_of(device)
    loop_ok = name.startswith("loop") and os.environ.get("MONTAUK_IMAGER_TEST_LOOP") == "1"
    if name.startswith(("loop", "zram", "ram", "sr", "dm-", "md")) and not loop_ok:
        raise Fail(f"{device} is not an SD card or USB drive")
    for _, mnt in disk_mounts(name):
        if mnt in SYSTEM_MOUNTS:
            raise Fail(f"{device} holds this computer's own system ({mnt}); refusing to touch it")
    if sysfs_int(f"/sys/block/{name}/ro"):
        raise Fail(f"{device} is write-protected (check the lock switch on the side of the card)")
    if disk_bytes(name) == 0:
        raise Fail(f"{device} has no medium (is the card inserted?)")
    lbs = sysfs_int(f"/sys/block/{name}/queue/logical_block_size", SECTOR)
    if lbs != SECTOR:
        raise Fail(f"{device} has {lbs}-byte sectors; the image's partition table assumes 512")
    return name


def unmount_disk(rep, name):
    for part in sys_partitions(name).values():
        if holders(os.path.basename(part.node)):
            raise Fail(f"{part.node} is in use by device mapper/LVM/RAID; close it first")
    for attempt in range(3):
        mounts = disk_mounts(name)
        if not mounts:
            return
        for dev, mnt in sorted(mounts, key=lambda m: -len(m[1])):
            if mnt == "[SWAP]":
                raise Fail(f"{dev} is in use as swap")
            rep.log(f"unmounting {dev} from {mnt}")
            run(rep, ["umount", mnt], ok=(0, 32), quiet=True)
        time.sleep(0.5)
    if disk_mounts(name):
        raise Fail(f"could not unmount /dev/{name}; close any program using it")


def settle():
    with contextlib.suppress(FileNotFoundError, subprocess.SubprocessError):
        subprocess.run(["udevadm", "settle", "--timeout=10"], env=tool_env(),
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=15)


def reread_partitions(rep, name, want):
    """Make the kernel load the new table; wait for the partitions we expect."""
    dev = f"/dev/{name}"
    for attempt in range(10):
        settle()
        unmount_disk(rep, name)
        fd = os.open(dev, os.O_RDONLY | os.O_CLOEXEC)
        try:
            fcntl.ioctl(fd, BLKRRPART)
            err = None
        except OSError as e:
            err = e
        finally:
            os.close(fd)
        if err is None:
            break
        if err.errno != errno.EBUSY:
            raise Fail(f"could not reread the partition table of {dev}: {err.strerror}")
        time.sleep(1)
    else:
        # The kernel keeps the old table while a partition is open (a desktop
        # probing it); partx updates the entries one by one instead.
        run(rep, ["partx", "-u", dev], ok=(0, 1), quiet=True)
    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        settle()
        parts = sys_partitions(name)
        if all(n in parts and os.path.exists(parts[n].node) for n in want):
            return parts
        time.sleep(0.5)
    raise Fail(f"the partitions of {dev} did not appear after writing the table")


@contextlib.contextmanager
def no_automount(rep, name):
    """Keep the desktop from mounting the card's partitions while we work.

    Closing a disk opened for writing makes udev re-probe it, and a desktop
    then mounts the partitions it finds (and Linux writes to an ext2 volume as
    soon as it is mounted read-write). A runtime rule turns udisks' automount
    off for this disk until we are done."""
    rules = "/run/udev/rules.d"
    path = os.path.join(rules, f"90-montauk-imager-{name}.rules")
    sep = "p" if name[-1].isdigit() else ""
    try:
        os.makedirs(rules, exist_ok=True)
        with open(path, "w") as f:
            f.write("# Written by MontaukOS's montauk-imager.py while it writes this disk.\n"
                    f'KERNEL=="{name}", ENV{{UDISKS_AUTO}}="0"\n'
                    f'KERNEL=="{name}{sep}[0-9]*", ENV{{UDISKS_AUTO}}="0", ENV{{UDISKS_IGNORE}}="1"\n')
        subprocess.run(["udevadm", "control", "--reload"], env=tool_env(),
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except (OSError, subprocess.SubprocessError):
        path = None
    try:
        yield
    finally:
        if path:
            with contextlib.suppress(OSError, subprocess.SubprocessError):
                os.unlink(path)
                subprocess.run(["udevadm", "control", "--reload"], env=tool_env(),
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                # Let the desktop see the finished card.
                subprocess.run(["udevadm", "trigger", "--action=change", f"/dev/{name}"],
                               env=tool_env(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def probe(node):
    """blkid's LABEL and TYPE of a partition (low-level probe, no cache)."""
    p = subprocess.run(["blkid", "-p", "-o", "export", node], stdout=subprocess.PIPE,
                       stderr=subprocess.DEVNULL, text=True, env=tool_env())
    info = {}
    for line in p.stdout.splitlines():
        k, _, v = line.partition("=")
        info[k] = v
    return info.get("LABEL", ""), info.get("TYPE", "")


def card_layout(rep, name):
    """(esp, root) Parts of a MontaukOS card, root None for the ramdisk layout."""
    with open(f"/dev/{name}", "rb") as f:
        parts = parse_mbr(f.read(512))
    if not parts:
        raise Fail(f"/dev/{name} has no partition table; write a fresh image instead")
    nodes = sys_partitions(name)
    esp = root = None
    for p in parts:
        if p.number not in nodes:
            continue
        p.node = nodes[p.number].node
        p.label, p.fstype = probe(p.node)
        if p.label == ROOT_LABEL and p.fstype.startswith("ext"):
            root = root or p
        elif p.fstype == "vfat" and (p.label == ESP_LABEL or (p.type == TYPE_ESP and not esp)):
            esp = p
    if root and root.fstype != "ext2":
        raise Fail(f"{root.node} is {root.fstype}, but MontaukOS needs ext2; write a fresh image")
    if not esp:
        raise Fail(f"/dev/{name} has no MontaukOS boot partition (FAT32, label {ESP_LABEL})")
    return esp, root, parts


# ==== ext2 ====

def fs_info(rep, node):
    p = run(rep, ["dumpe2fs", "-h", node], quiet=True)
    info = {}
    for line in p.stdout.splitlines():
        k, _, v = line.partition(":")
        info[k.strip()] = v.strip()
    return (int(info.get("Block count", 0)), int(info.get("Block size", 4096)),
            set(info.get("Filesystem features", "").split()))


def check_features(rep, node):
    _, _, features = fs_info(rep, node)
    bad = (features & EXT_INCOMPAT) - SUPPORTED_INCOMPAT
    if bad:
        raise Fail(f"{node} now has ext2 features MontaukOS cannot mount: {', '.join(sorted(bad))}")


def fsck(rep, node, check_only=False):
    if check_only:
        p = run(rep, ["e2fsck", "-f", "-n", node], ok=range(0, 256), quiet=True)
        if p.returncode != 0:
            for line in p.stdout.splitlines()[-20:]:
                rep.log(line)
            raise Fail(f"e2fsck reports problems on {node} after the update (exit {p.returncode})")
        return
    rep.log(f"checking {node}")
    p = run(rep, ["e2fsck", "-f", "-p", node], ok=range(0, 256), quiet=True)
    if p.returncode & ~3:
        for line in p.stdout.splitlines()[-20:]:
            rep.log(line)
        raise Fail(f"{node} has errors e2fsck cannot fix by itself; run `sudo e2fsck -f {node}`")
    if p.returncode:
        rep.warn(f"e2fsck repaired {node} (it was probably not shut down cleanly)")


def expand_root(rep, name):
    """Grow the root partition to the end of the card, then the volume."""
    unmount_disk(rep, name)
    esp, root, parts = card_layout(rep, name)
    if not root:
        raise Fail("this card has no ext2 system partition to expand (ramdisk layout)")
    if any(p.start > root.start for p in parts):
        raise Fail(f"partition {root.number} is not the last one; cannot expand it")
    rep.stage("Expanding the system partition")
    total = disk_bytes(name) // SECTOR
    end = total // ALIGN * ALIGN
    count = min(end - root.start, MAX_ROOT_BYTES // SECTOR, 0xFFFFFFFF // ALIGN * ALIGN)
    if count > root.sectors:
        dev = f"/dev/{name}"
        fd = os.open(dev, os.O_RDWR | os.O_EXCL | os.O_CLOEXEC)
        try:
            mbr = bytearray(os.pread(fd, SECTOR, 0))
            struct.pack_into("<I", mbr, 446 + 16 * (root.number - 1) + 12, count)
            os.pwrite(fd, bytes(mbr), 0)
            os.fsync(fd)
        finally:
            os.close(fd)
        rep.log(f"partition {root.number}: {human(root.size)} -> {human(count * SECTOR)}")
        reread_partitions(rep, name, [p.number for p in parts])
        root.sectors = count
    else:
        rep.log(f"partition {root.number} already fills the card ({human(root.size)})")
    fsck(rep, root.node)
    blocks, bs, _ = fs_info(rep, root.node)
    target = min(root.size, MAX_ROOT_BYTES) // bs
    if target > blocks:
        rep.log(f"growing the volume: {human(blocks * bs)} -> {human(target * bs)}")
        rep.progress(0, 1, "Growing the volume (resize2fs)", force=True)
        run(rep, ["resize2fs", root.node, str(target)])
        check_features(rep, root.node)
        fsck(rep, root.node, check_only=True)
    else:
        rep.log(f"the volume already fills the partition ({human(blocks * bs)})")


# ==== Flash ====

def cmd_flash(rep, args):
    img = Image(args.image)
    name = check_target(rep, args.device)
    dev = f"/dev/{name}"
    card = disk_bytes(name)
    known = img.size()
    if known and known > card:
        raise Fail(f"the image ({human(known)}) does not fit on the card ({human(card)})")
    confirm(args, f"Everything on {dev} ({human(card)}) will be erased.")

    rep.stage(f"Preparing {dev}")
    with no_automount(rep, name):
        unmount_disk(rep, name)
        try:
            fd = os.open(dev, os.O_RDWR | os.O_EXCL | os.O_CLOEXEC)
        except OSError as e:
            raise Fail(f"cannot open {dev}: {e.strerror} (is something still using it?)")
        # One descriptor from the first write to the reread: udev re-probes
        # the disk when a writer closes it, and a volume mounted before the
        # read-back would be changed under the verification.
        try:
            written, digest = write_image(rep, img, fd, card, dev)
            if not args.no_verify:
                verify_image(rep, fd, written, digest)
            parts = parse_mbr(os.pread(fd, SECTOR, 0)) or []
            with contextlib.suppress(OSError):
                fcntl.ioctl(fd, BLKRRPART)
        finally:
            os.close(fd)
        reread_partitions(rep, name, [p.number for p in parts])
        if args.expand:
            expand_root(rep, name)
        settle()
    rep.done(f"MontaukOS is on {dev}. You can remove the card.")


def write_image(rep, img, fd, card, dev):
    digest = hashlib.sha256()
    written = 0
    rep.stage(f"Writing {os.path.basename(img.path)} to {dev}")
    rep.cancellable = True
    try:
        start = time.monotonic()
        synced = 0
        with img.open() as (stream, raw):
            while True:
                rep.check_cancel()
                buf = stream.read(CHUNK)
                if not buf:
                    break
                if written == 0:
                    try:
                        ok = parse_mbr(buf[:SECTOR]) is not None
                    except Fail:
                        ok = False
                    if not ok:
                        raise Fail(f"{os.path.basename(img.path)} is not a MontaukOS Raspberry Pi image "
                                   "(no MBR partition table); nothing was written")
                if written + len(buf) > card:
                    raise Fail(f"the image does not fit on the card ({human(card)})")
                view = memoryview(buf)
                while view:
                    n = os.write(fd, view)
                    view = view[n:]
                digest.update(buf)
                written += len(buf)
                if written - synced >= 32 * MIB:
                    os.fsync(fd)   # keep the progress honest: page cache is not the card
                    synced = written
                rate = written / max(time.monotonic() - start, 0.001)
                rep.progress(raw.tell(), img.file_size,
                             f"Writing: {written // MIB} MiB ({rate / 1e6:.1f} MB/s)")
        rep.progress(1, 1, "Flushing to the card", force=True)
        os.fsync(fd)
    except Fail as e:
        if str(e) == "cancelled":
            os.fsync(fd)
            raise Fail("cancelled; the card is incomplete now, write it again before using it")
        raise
    finally:
        rep.cancellable = False
    # Clear the last MiB: a backup GPT left by an earlier OS would otherwise
    # outlive the MBR this image brings.
    if card - written > 2 * MIB:
        os.pwrite(fd, bytes(MIB), card - MIB)
        os.fsync(fd)
    rep.log(f"wrote {human(written)}")
    return written, digest.digest()


def verify_image(rep, fd, written, digest):
    rep.stage("Verifying")
    fcntl.ioctl(fd, BLKFLSBUF)    # read the card, not the page cache
    os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
    check = hashlib.sha256()
    done = 0
    start = time.monotonic()
    while done < written:
        if rep.cancel.is_set():
            rep.warn("verification cancelled; the card was written completely but not checked")
            return
        buf = os.pread(fd, min(CHUNK, written - done), done)
        if not buf:
            break
        check.update(buf)
        done += len(buf)
        rate = done / max(time.monotonic() - start, 0.001)
        rep.progress(done, written, f"Verifying: {done // MIB} MiB ({rate / 1e6:.1f} MB/s)")
    if done != written or check.digest() != digest:
        raise Fail("verification failed: the card does not read back what was written "
                   "(a failing or counterfeit card?)")
    rep.log("the card reads back correctly")


# ==== Update ====

def toml_sections(text):
    """[(table name or None, [lines])]: each table with the lines up to the next."""
    header = re.compile(r"^\[([A-Za-z0-9_.\-\"]+)\]\s*(#.*)?$")
    secs = [(None, [])]
    for line in text.splitlines():
        m = header.match(line)
        if m:
            secs.append((m.group(1), [line]))
        else:
            secs[-1][1].append(line)
    return secs


def merge_init(new, old):
    """The image's init.toml, keeping the card's on/off choices and own services."""
    enabled = re.compile(r"^\s*enabled\s*=\s*(true|false)\b")
    old_secs = {n: l for n, l in toml_sections(old) if n}
    new_secs = toml_sections(new)
    names = {n for n, _ in new_secs}
    out = []
    for n, lines in new_secs:
        lines = list(lines)
        if n and n.startswith("services.") and n in old_secs:
            choice = next((m.group(1) for m in map(enabled.match, old_secs[n]) if m), None)
            if choice is not None:
                idx = [i for i, l in enumerate(lines) if enabled.match(l)]
                if idx:
                    lines[idx[0]] = f"enabled = {choice}"
                else:
                    lines.insert(1, f"enabled = {choice}")
        out.extend(lines)
    extra = [(n, l) for n, l in toml_sections(old)
             if n and n.startswith("services.") and n not in names]
    if extra:
        out += ["", "# Services from the previous installation."]
        for _, lines in extra:
            out.extend(lines)
    return "\n".join(out).rstrip("\n") + "\n"


def merge_caps(new, old):
    """The image's capabilities.toml plus the grants only the card has."""
    names = {n for n, _ in toml_sections(new)}
    extra = [(n, l) for n, l in toml_sections(old)
             if n and n.startswith("grant.") and n not in names]
    out = new.rstrip("\n").splitlines()
    if extra:
        out += ["", "# ==== Grants from the previous installation ====", ""]
        for _, lines in extra:
            out.extend(lines)
    return "\n".join(out).rstrip("\n") + "\n"


class Op:
    def __init__(self, kind, rel, dst, src=None, data=None, size=0, replaced=0):
        self.kind, self.rel, self.dst, self.src, self.data = kind, rel, dst, src, data
        self.size, self.replaced = size, replaced


def alloc(size, bs, ext):
    blocks = -(-size // bs)
    if ext and blocks > 12:
        blocks += -(-blocks // (bs // 4))   # indirect blocks
    return blocks * bs


class Plan:
    def __init__(self, what, dst_root, ext):
        self.what, self.dst_root, self.ext = what, dst_root, ext
        self.ops = []
        self.unchanged = 0
        st = os.statvfs(dst_root)
        self.bs = st.f_frsize
        self.free = st.f_bavail * st.f_frsize

    def write(self, rel, src=None, data=None):
        dst = os.path.join(self.dst_root, rel)
        size = len(data) if data is not None else os.path.getsize(src)
        replaced = os.path.getsize(dst) if os.path.isfile(dst) else 0
        self.ops.append(Op("write", rel, dst, src, data, size, replaced))

    def needed(self):
        a = lambda n: alloc(n, self.bs, self.ext)
        net = 0
        biggest = 0
        for op in self.ops:
            if op.kind == "write":
                net += a(op.size) - (a(op.replaced) if op.replaced else 0)
                biggest = max(biggest, a(op.size))
            elif op.kind == "backup":
                net += a(op.size)
            elif op.kind == "delete":
                net -= a(op.size)
        return net + biggest + MIB

    def count(self, kind):
        return sum(1 for op in self.ops if op.kind == kind)


def same_file(a, b):
    try:
        return os.path.getsize(a) == os.path.getsize(b) and filecmp.cmp(a, b, shallow=False)
    except OSError:
        return False


def plan_tree(rep, plan, src_root, decide, prune_dirs=(), preserve=()):
    shipped = set()
    files = []
    for dirpath, dirnames, filenames in os.walk(src_root):
        rel_dir = os.path.relpath(dirpath, src_root)
        rel_dir = "" if rel_dir == "." else rel_dir
        dirnames[:] = sorted(d for d in dirnames
                             if decide(os.path.join(rel_dir, d), True) != "skip")
        for d in dirnames:
            shipped.add(os.path.join(rel_dir, d))
        for fn in sorted(filenames):
            files.append(os.path.join(rel_dir, fn))
            shipped.add(files[-1])

    for i, rel in enumerate(files):
        rep.progress(i, len(files), f"Comparing {plan.what} ({i}/{len(files)})")
        src = os.path.join(src_root, rel)
        dst = os.path.join(plan.dst_root, rel)
        policy = decide(rel, False)
        exists = os.path.lexists(dst)
        if exists and os.path.isdir(dst):
            raise Fail(f"{rel} is a directory on the card but a file in the image")
        if policy == "skip" or (policy == "add" and exists):
            continue
        if policy == "firstboot" and not exists:
            rep.log("not adding os/firstboot.elf: the card's system volume has already grown")
            continue
        if policy in ("merge-init", "merge-caps") and exists:
            with open(src, encoding="utf-8", errors="replace") as f:
                new = f.read()
            with open(dst, encoding="utf-8", errors="replace") as f:
                old = f.read()
            merged = (merge_init if policy == "merge-init" else merge_caps)(new, old)
            if merged == old:
                plan.unchanged += 1
                continue
            plan.ops.append(Op("backup", rel, dst, size=len(old.encode())))
            plan.write(rel, data=merged.encode())
            continue
        if exists and same_file(src, dst):
            plan.unchanged += 1
            continue
        if policy == "backup" and exists:
            plan.ops.append(Op("backup", rel, dst, size=os.path.getsize(dst)))
        plan.write(rel, src=src)

    for top in prune_dirs:
        base = os.path.join(plan.dst_root, top)
        if not os.path.isdir(base) or top not in shipped:
            continue
        for dirpath, dirnames, filenames in os.walk(base, topdown=False):
            rel_dir = os.path.relpath(dirpath, plan.dst_root)
            for fn in filenames:
                rel = os.path.join(rel_dir, fn)
                if rel not in shipped and rel not in preserve:
                    path = os.path.join(dirpath, fn)
                    plan.ops.append(Op("delete", rel, path, size=os.path.getsize(path)))
            if rel_dir not in shipped:
                plan.ops.append(Op("rmdir", rel_dir, dirpath))


def execute(rep, plan, done, total):
    for op in plan.ops:
        if op.kind == "delete":
            rep.log(f"removing {op.rel}")
            os.unlink(op.dst)
    for op in plan.ops:
        if op.kind == "rmdir":
            with contextlib.suppress(OSError):
                os.rmdir(op.dst)   # only if it ended up empty
    for op in plan.ops:
        if op.kind == "backup":
            suffix = ".bak" if op.rel in ESP_BACKUP else ".old"
            rep.log(f"keeping the card's {op.rel} as {os.path.basename(op.rel)}{suffix}")
            shutil.copyfile(op.dst, op.dst + suffix)
    for op in plan.ops:
        if op.kind != "write":
            continue
        os.makedirs(os.path.dirname(op.dst), exist_ok=True)
        tmp = os.path.join(os.path.dirname(op.dst), ".imager-" + os.path.basename(op.dst))
        with open(tmp, "wb") as out:
            if op.data is not None:
                out.write(op.data)
                done += len(op.data)
            else:
                with open(op.src, "rb") as src:
                    while True:
                        buf = src.read(CHUNK)
                        if not buf:
                            break
                        out.write(buf)
                        done += len(buf)
                        rep.progress(done, total, f"Writing {op.rel}")
            out.flush()
            os.fsync(out.fileno())
        if op.src:
            with contextlib.suppress(OSError):   # FAT has no modes
                os.chmod(tmp, os.stat(op.src).st_mode & 0o777)
        os.replace(tmp, op.dst)
        rep.progress(done, total, f"Writing {op.rel}")
    return done


@contextlib.contextmanager
def mounted(rep, workdir, label, source, fstype, opts):
    mnt = os.path.join(workdir, label)
    os.makedirs(mnt, exist_ok=True)
    run(rep, ["mount", "-t", fstype, "-o", opts, source, mnt], quiet=True)
    try:
        yield mnt
    finally:
        os.sync()
        p = subprocess.run(["umount", mnt], env=tool_env(),
                           stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        if p.returncode:
            time.sleep(1)
            subprocess.run(["umount", "-l", mnt], env=tool_env(),
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def cmd_update(rep, args):
    img = Image(args.image)
    name = check_target(rep, args.device)
    with no_automount(rep, name):
        update(rep, args, img, name)


def update(rep, args, img, name):
    dev = f"/dev/{name}"
    workdir = tempfile.mkdtemp(prefix="montauk-imager-")
    try:
        raw = img.materialize(rep, workdir)
        iparts = img.partitions(raw)
        iesp = next((p for p in iparts if p.type in FAT_TYPES), None)
        iroot = next((p for p in iparts if p.type == TYPE_LINUX), None)
        if not iesp:
            raise Fail("the image has no boot partition; is it a MontaukOS Raspberry Pi image?")

        rep.stage(f"Inspecting {dev}")
        unmount_disk(rep, name)
        esp, root, _ = card_layout(rep, name)
        if bool(iroot) != bool(root):
            have = "an ext2 system partition" if root else "the single-partition (ramdisk) layout"
            want = "an ext2 system partition" if iroot else "the single-partition (ramdisk) layout"
            raise Fail(f"the card has {have} but the image uses {want}; write a fresh image instead")
        rep.log(f"boot partition {esp.node} ({human(esp.size)})")
        if root:
            rep.log(f"system partition {root.node} ({human(root.size)}, "
                    f"{'fills the card' if root.end + ALIGN >= disk_bytes(name) // SECTOR else 'not expanded'})")
        confirm(args, f"MontaukOS on {dev} will be updated; users, files and settings are kept.")

        signal.signal(signal.SIGHUP, signal.SIG_IGN)   # an update must not stop halfway
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        if root:
            rep.stage("Checking the system volume")
            fsck(rep, root.node)
            if args.expand:
                expand_root(rep, name)
                esp, root, _ = card_layout(rep, name)

        with contextlib.ExitStack() as stack:
            rep.stage("Mounting")
            loop = lambda p: f"ro,loop,offset={p.offset},sizelimit={p.size}"
            src_esp = stack.enter_context(mounted(rep, workdir, "image-esp", raw, "vfat", loop(iesp)))
            dst_esp = stack.enter_context(mounted(rep, workdir, "card-esp", esp.node, "vfat", "rw,noatime"))
            plans = []

            def esp_rule(rel, is_dir):
                if is_dir:
                    return "system"
                if rel == FIRMWARE_FILE and not args.firmware:
                    return "add"
                return "backup" if rel in ESP_BACKUP else "system"

            rep.stage("Comparing the card with the image")
            plans.append((Plan("the boot partition", dst_esp, False), src_esp, esp_rule, (), ()))

            if root:
                src_root = stack.enter_context(mounted(rep, workdir, "image-root", raw, "ext2", loop(iroot)))
                dst_root = stack.enter_context(mounted(rep, workdir, "card-root", root.node, "ext2", "rw,noatime"))

                def root_rule(rel, is_dir):
                    top = rel.split("/")[0]
                    if top in STATE_DIRS:
                        return "skip"
                    if is_dir:
                        return "system"
                    if rel == FIRSTBOOT:
                        return "firstboot"
                    if rel in (INIT_TOML, CAPS_TOML):
                        if args.keep_config:
                            return "add"
                        return "merge-init" if rel == INIT_TOML else "merge-caps"
                    return "add" if top in MERGE_DIRS else "system"

                plans.append((Plan("the system partition", dst_root, True), src_root, root_rule,
                              PRUNE_DIRS if args.prune else (), PRESERVE))

            for plan, src, rule, prune, keep in plans:
                plan_tree(rep, plan, src, rule, prune, keep)
                need = plan.needed()
                rep.log(f"{plan.what}: {plan.count('write')} to write, {plan.unchanged} unchanged"
                        + (f", {plan.count('delete')} to remove" if plan.count("delete") else ""))
                if need > plan.free:
                    hint = ("enable Expand (--expand) or write a fresh image"
                            if plan.ext else "write a fresh image")
                    raise Fail(f"not enough space on {plan.what}: needs {human(need)}, "
                               f"{human(plan.free)} free; {hint}")

            total = sum(op.size for plan, *_ in plans for op in plan.ops if op.kind == "write")
            if total == 0 and not any(plan.ops for plan, *_ in plans):
                rep.log("the card already matches the image")
            rep.stage("Updating")
            done = 0
            for plan, *_ in plans:
                done = execute(rep, plan, done, total)
            rep.progress(1, 1, "Flushing to the card", force=True)
            os.sync()
            rep.stage("Unmounting")

        if root:
            rep.stage("Checking the system volume")
            fsck(rep, root.node, check_only=True)
            check_features(rep, root.node)
        settle()
        rep.done(f"MontaukOS on {dev} is up to date. You can remove the card.")
    finally:
        remove_workdir(rep, workdir)


def remove_workdir(rep, workdir):
    """rmtree crosses mount points: never let it into a card left mounted."""
    for entry in os.listdir(workdir):
        path = os.path.join(workdir, entry)
        if os.path.ismount(path):
            rep.warn(f"{path} is still mounted; unmount it by hand (sudo umount {path})")
        elif os.path.isdir(path):
            with contextlib.suppress(OSError):
                os.rmdir(path)
        else:
            with contextlib.suppress(OSError):
                os.unlink(path)
    with contextlib.suppress(OSError):
        os.rmdir(workdir)


def cmd_expand(rep, args):
    name = check_target(rep, args.device)
    confirm(args, f"The system partition on /dev/{name} will be grown to fill the card.")
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    with no_automount(rep, name):
        expand_root(rep, name)
        settle()
    rep.done(f"The system partition on /dev/{name} fills the card. You can remove the card.")


def confirm(args, what):
    if args.yes:
        return
    if not sys.stdin.isatty():
        raise Fail(f"{what} Pass --yes to go ahead.")
    answer = input(f"{what} Type 'yes' to continue: ")
    if answer.strip().lower() != "yes":
        raise Fail("aborted")


# ==== Card discovery (unprivileged) ====

def truthy(v):
    return v in (True, 1, "1", "true")


def list_cards(show_all=False):
    """Removable disks (or all non-system disks) as dicts, for the GUI and `list`."""
    cols = "NAME,PATH,SIZE,TYPE,RM,HOTPLUG,TRAN,MODEL,VENDOR,FSTYPE,LABEL,MOUNTPOINT,RO"
    try:
        p = subprocess.run(["lsblk", "-J", "-b", "-o", cols], stdout=subprocess.PIPE,
                           stderr=subprocess.DEVNULL, text=True, env=tool_env(), timeout=10)
        devices = json.loads(p.stdout or "{}").get("blockdevices", [])
    except (OSError, ValueError, subprocess.SubprocessError):
        return []

    def mounts(d):
        out = [d.get("mountpoint")] if d.get("mountpoint") else []
        for c in d.get("children") or []:
            out += mounts(c)
        return out

    cards = []
    for d in devices:
        name = d.get("name", "")
        size = int(d.get("size") or 0)
        if d.get("type") != "disk" or size == 0 or name.startswith(("zram", "loop", "ram")):
            continue
        if any(m in SYSTEM_MOUNTS for m in mounts(d)):
            continue
        removable = (truthy(d.get("rm")) or truthy(d.get("hotplug"))
                     or d.get("tran") in ("usb", "mmc", "sdio") or name.startswith("mmcblk"))
        if not removable and not show_all:
            continue
        parts = []
        for c in d.get("children") or []:
            if c.get("type") != "part":
                continue
            cname = c.get("name", "")
            parts.append({
                "path": c.get("path") or f"/dev/{cname}",
                "size": int(c.get("size") or 0),
                "start": sysfs_int(f"/sys/block/{name}/{cname}/start") * SECTOR,
                "label": c.get("label") or "",
                "fstype": c.get("fstype") or "",
                "mount": c.get("mountpoint") or "",
            })
        model = " ".join(x.strip() for x in (d.get("vendor") or "", d.get("model") or "") if x and x.strip())
        card = {"path": d.get("path") or f"/dev/{name}", "name": name, "size": size,
                "model": model or ("SD card" if name.startswith("mmcblk") else "Drive"),
                "ro": truthy(d.get("ro")), "removable": removable, "parts": parts}
        card.update(describe(card))
        cards.append(card)
    return cards


def describe(card):
    """What is on a card: {'kind': 'ext2'|'ramdisk'|'other'|'empty', 'contents', 'fills'}."""
    parts = card["parts"]
    root = next((p for p in parts if p["label"] == ROOT_LABEL), None)
    esp = next((p for p in parts if p["label"] == ESP_LABEL and p["fstype"] == "vfat"), None)
    if root:
        fills = (root["start"] + root["size"] + 2 * MIB >= card["size"]
                 or root["size"] >= MAX_ROOT_BYTES - MIB)
        text = f"MontaukOS, system partition {human(root['size'])}"
        text += ", fills the card" if fills else ", not expanded yet"
        return {"kind": "ext2", "contents": text, "fills": fills}
    if esp:
        return {"kind": "ramdisk", "contents": "MontaukOS (single partition, ramdisk)", "fills": True}
    if parts:
        names = [p["label"] or p["fstype"] or "unknown" for p in parts]
        return {"kind": "other", "contents": ", ".join(names), "fills": False}
    return {"kind": "empty", "contents": "empty", "fills": False}


def cmd_list(rep, args):
    cards = list_cards(args.all)
    if not cards:
        print("No SD cards or USB drives found." + ("" if args.all else " (--all shows fixed disks too)"))
    for c in cards:
        print(f"{c['path']:<16} {human(c['size']):>9}  {c['model']:<28} {c['contents']}"
              + ("  [write-protected]" if c["ro"] else ""))


def image_summary(path):
    """One line about an image file for the GUI, and its layout if it can tell."""
    try:
        img = Image(path)
    except Fail as e:
        return str(e), None
    when = time.strftime("%Y-%m-%d %H:%M", time.localtime(os.path.getmtime(img.path)))
    if img.kind != "raw":
        return f"{img.kind} compressed, {human(img.file_size)}, {when}", None
    try:
        parts = img.partitions(img.path)
    except Fail as e:
        return str(e), None
    layout = "ext2" if any(p.type == TYPE_LINUX for p in parts) else "ramdisk"
    desc = "boot + ext2 system partition" if layout == "ext2" else "single partition (ramdisk)"
    return f"{human(img.file_size)}, {desc}, built {when}", layout


# ==== GUI ====

def gui_main(initial):
    try:
        import tkinter as tk
        from tkinter import ttk, filedialog, messagebox
    except ImportError:
        sys.exit("montauk-imager: the GUI needs Tkinter (sudo apt install python3-tk)")
    try:
        root = tk.Tk()
    except tk.TclError as e:
        sys.exit(f"montauk-imager: cannot open a window ({e}); use the list/flash/update commands")

    root.title("MontaukOS Imager")
    root.minsize(680, 600)
    style = ttk.Style(root)
    if "clam" in style.theme_names():
        style.theme_use("clam")
    style.configure("Title.TLabel", font=("TkDefaultFont", 15, "bold"))
    style.configure("Dim.TLabel", foreground="#666666")
    style.configure("Accent.TButton", font=("TkDefaultFont", 10, "bold"))

    S = {"busy": None, "proc": None, "cards": [], "sig": None, "layout": None}
    events = queue.Queue()

    image_var = tk.StringVar(value=initial or (DEFAULT_IMAGE if os.path.isfile(DEFAULT_IMAGE) else ""))
    image_info = tk.StringVar()
    show_all = tk.BooleanVar(value=False)
    action = tk.StringVar(value="flash")
    verify = tk.BooleanVar(value=True)
    flash_expand = tk.BooleanVar(value=False)
    merge_cfg = tk.BooleanVar(value=True)
    firmware = tk.BooleanVar(value=False)
    prune = tk.BooleanVar(value=False)
    upd_expand = tk.BooleanVar(value=False)
    status = tk.StringVar(value="Choose an image and an SD card.")
    card_info = tk.StringVar()

    main = ttk.Frame(root, padding=16)
    main.pack(fill="both", expand=True)
    main.columnconfigure(0, weight=1)

    head = ttk.Frame(main)
    head.grid(row=0, column=0, sticky="ew")
    ttk.Label(head, text="MontaukOS Imager", style="Title.TLabel").pack(side="left")
    ttk.Label(head, text="Raspberry Pi 4 / 400", style="Dim.TLabel").pack(side="left", padx=10, pady=(5, 0))

    # Image
    imgf = ttk.LabelFrame(main, text="Image", padding=10)
    imgf.grid(row=1, column=0, sticky="ew", pady=(12, 0))
    imgf.columnconfigure(0, weight=1)
    img_entry = ttk.Entry(imgf, textvariable=image_var)
    img_entry.grid(row=0, column=0, sticky="ew")
    browse_btn = ttk.Button(imgf, text="Browse...", command=lambda: browse())
    browse_btn.grid(row=0, column=1, padx=(6, 0))
    build_btn = ttk.Button(imgf, text="Build", command=lambda: build())
    if REPO:
        build_btn.grid(row=0, column=2, padx=(6, 0))
    ttk.Label(imgf, textvariable=image_info, style="Dim.TLabel").grid(row=1, column=0, columnspan=3, sticky="w", pady=(4, 0))

    # Card
    cardf = ttk.LabelFrame(main, text="SD card", padding=10)
    cardf.grid(row=2, column=0, sticky="nsew", pady=(10, 0))
    main.rowconfigure(2, weight=1)
    cardf.columnconfigure(0, weight=1)
    cardf.rowconfigure(0, weight=1)
    tree = ttk.Treeview(cardf, columns=("device", "size", "model", "contents"), show="headings",
                        height=4, selectmode="browse")
    for col, text, width, anchor in (("device", "Device", 110, "w"), ("size", "Size", 80, "e"),
                                     ("model", "Name", 170, "w"), ("contents", "Contents", 300, "w")):
        tree.heading(col, text=text, anchor=anchor)
        tree.column(col, width=width, anchor=anchor, stretch=(col == "contents"))
    tree.grid(row=0, column=0, columnspan=3, sticky="nsew")
    ttk.Label(cardf, textvariable=card_info, style="Dim.TLabel").grid(row=1, column=0, sticky="w", pady=(4, 0))
    ttk.Checkbutton(cardf, text="Show all drives", variable=show_all,
                    command=lambda: refresh(force=True)).grid(row=1, column=1, padx=6, pady=(4, 0))
    ttk.Button(cardf, text="Refresh", command=lambda: refresh(force=True)).grid(row=1, column=2, pady=(4, 0))

    # Action
    actf = ttk.LabelFrame(main, text="Action", padding=10)
    actf.grid(row=3, column=0, sticky="ew", pady=(10, 0))
    actf.columnconfigure(1, weight=1)
    rb_flash = ttk.Radiobutton(actf, text="Write a fresh image (erases the card)", value="flash",
                               variable=action, command=lambda: update_ui())
    rb_update = ttk.Radiobutton(actf, text="Update the installation on the card (keeps users, files and settings)",
                                value="update", variable=action, command=lambda: update_ui())
    rb_flash.grid(row=0, column=0, columnspan=2, sticky="w")
    rb_update.grid(row=1, column=0, columnspan=2, sticky="w")

    optf = ttk.Frame(actf, padding=(22, 4, 0, 0))
    optf.grid(row=2, column=0, columnspan=2, sticky="ew")
    flash_opts = [
        ttk.Checkbutton(optf, text="Verify the card after writing", variable=verify),
        ttk.Checkbutton(optf, text="Expand the system partition to fill the card now "
                                   "(otherwise MontaukOS does it on first boot)", variable=flash_expand),
    ]
    update_opts = [
        ttk.Checkbutton(optf, text="Merge new services and capability grants into the card's configuration",
                        variable=merge_cfg),
        ttk.Checkbutton(optf, text="Update the UEFI firmware (RPI_EFI.fd; resets the firmware's settings)",
                        variable=firmware),
        ttk.Checkbutton(optf, text="Remove system files the new image no longer ships", variable=prune),
        ttk.Checkbutton(optf, text="Expand the system partition to fill the card", variable=upd_expand),
    ]

    # Progress
    progf = ttk.Frame(main)
    progf.grid(row=4, column=0, sticky="ew", pady=(12, 0))
    progf.columnconfigure(0, weight=1)
    bar = ttk.Progressbar(progf, mode="determinate", maximum=1000)
    bar.grid(row=0, column=0, columnspan=3, sticky="ew")
    ttk.Label(progf, textvariable=status).grid(row=1, column=0, sticky="w", pady=(4, 0))
    log_btn = ttk.Button(progf, text="Show log", command=lambda: toggle_log())
    log_btn.grid(row=1, column=1, padx=6, pady=(4, 0))
    logf = ttk.Frame(main)
    logf.columnconfigure(0, weight=1)
    logf.rowconfigure(0, weight=1)
    log = tk.Text(logf, height=10, wrap="word", state="disabled", font=("TkFixedFont", 9),
                  relief="flat", borderwidth=1)
    scroll = ttk.Scrollbar(logf, orient="vertical", command=log.yview)
    log.configure(yscrollcommand=scroll.set)
    log.grid(row=0, column=0, sticky="nsew")
    scroll.grid(row=0, column=1, sticky="ns")
    log.tag_configure("stage", font=("TkFixedFont", 9, "bold"))
    log.tag_configure("warn", foreground="#a05a00")
    log.tag_configure("error", foreground="#b00020", font=("TkFixedFont", 9, "bold"))
    log.tag_configure("done", foreground="#0a7a2f", font=("TkFixedFont", 9, "bold"))

    btns = ttk.Frame(main)
    btns.grid(row=6, column=0, sticky="e", pady=(12, 0))
    cancel_btn = ttk.Button(btns, text="Cancel", command=lambda: cancel())
    start_btn = ttk.Button(btns, text="Write", style="Accent.TButton", command=lambda: start())
    cancel_btn.pack(side="left", padx=(0, 6))
    start_btn.pack(side="left")

    def toggle_log(show=None):
        visible = bool(logf.winfo_ismapped())
        show = (not visible) if show is None else show
        if show and not visible:
            logf.grid(row=5, column=0, sticky="nsew", pady=(8, 0))
            main.rowconfigure(5, weight=2)
            log_btn.configure(text="Hide log")
        elif not show and visible:
            logf.grid_remove()
            main.rowconfigure(5, weight=0)
            log_btn.configure(text="Show log")

    def append(text, tag=None):
        log.configure(state="normal")
        log.insert("end", text + "\n", (tag,) if tag else ())
        log.see("end")
        log.configure(state="disabled")

    def selected():
        sel = tree.selection()
        return next((c for c in S["cards"] if c["path"] == (sel[0] if sel else None)), None)

    def refresh(force=False):
        if S["busy"] in ("flash", "update", "expand") and not force:
            return
        cards = list_cards(show_all.get())
        sig = json.dumps(cards, sort_keys=True)
        if sig == S["sig"] and not force:
            return
        S["sig"] = sig
        keep = selected()
        S["cards"] = cards
        tree.delete(*tree.get_children())
        for c in cards:
            contents = c["contents"] + ("  [write-protected]" if c["ro"] else "")
            tree.insert("", "end", iid=c["path"],
                        values=(c["path"], human(c["size"]), c["model"], contents))
        if keep and tree.exists(keep["path"]):
            tree.selection_set(keep["path"])
        elif len(cards) == 1:
            tree.selection_set(cards[0]["path"])
        update_ui()

    def poll_cards():
        if not S["busy"]:
            refresh()
        root.after(2000, poll_cards)

    def refresh_image(*_):
        path = image_var.get().strip()
        if not path:
            image_info.set("No image chosen." + ("  Build one, or browse for montauk-rpi4.img." if REPO else ""))
            S["layout"] = None
        else:
            info, S["layout"] = image_summary(path)
            image_info.set(info)
        update_ui()

    def update_ui(*_):
        card = selected()
        installed = bool(card and card["kind"] in ("ext2", "ramdisk"))
        rb_update.configure(state="normal" if installed and not S["busy"] else "disabled")
        if action.get() == "update" and not installed:
            action.set("flash")
        for w in flash_opts + update_opts:
            w.pack_forget()
        opts = flash_opts if action.get() == "flash" else update_opts
        for w in opts:
            w.pack(anchor="w")
        if card:
            if card["kind"] == "ext2" and card["fills"]:
                upd_expand.set(False)
            card_info.set(f"{card['path']}: {card['model']}, {human(card['size'])}, {card['contents']}")
        else:
            card_info.set("Insert an SD card (directly or in a USB reader)." if not S["cards"]
                          else "Select the card to write.")
        state = "disabled" if S["busy"] else "normal"
        for w in opts + [img_entry, browse_btn, build_btn, rb_flash]:
            w.configure(state=state)
        if card and card["kind"] == "ext2" and card["fills"] and not S["busy"]:
            update_opts[3].configure(state="disabled")
        if card and card["kind"] == "ramdisk":
            update_opts[3].configure(state="disabled")
        start_btn.configure(text="Write" if action.get() == "flash" else "Update")
        ready = card and image_var.get().strip() and not S["busy"] and not card["ro"]
        start_btn.configure(state="normal" if ready else "disabled")
        cancel_btn.configure(state="normal" if S["busy"] in ("flash", "build") else "disabled")

    def browse():
        start_dir = os.path.dirname(image_var.get()) or REPO or os.path.expanduser("~")
        path = filedialog.askopenfilename(
            parent=root, title="Choose a MontaukOS image", initialdir=start_dir,
            filetypes=[("Disk images", "*.img *.img.gz *.img.xz *.gz *.xz *.zip"), ("All files", "*")])
        if path:
            image_var.set(path)

    def launch(kind, cmd, cwd=None):
        S["busy"] = kind
        bar.configure(mode="determinate", value=0)
        append("$ " + " ".join(cmd), "stage")
        try:
            proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                    stderr=subprocess.STDOUT, text=True, bufsize=1,
                                    errors="replace", cwd=cwd)
        except OSError as e:
            S["busy"] = None
            fail(f"Could not start {cmd[0]}: {e}")
            return
        S["proc"] = proc

        def reader():
            for line in proc.stdout:
                events.put(("line", line.rstrip("\n")))
            events.put(("exit", proc.wait()))
        threading.Thread(target=reader, daemon=True).start()
        update_ui()

    def fail(msg):
        status.set(msg)
        append(msg, "error")
        toggle_log(True)
        messagebox.showerror("MontaukOS Imager", msg, parent=root)

    def pump():
        try:
            while True:
                kind, val = events.get_nowait()
                if kind == "line":
                    handle_line(val)
                else:
                    finished(val)
        except queue.Empty:
            pass
        root.after(50, pump)

    def handle_line(line):
        if S["busy"] == "build" or not line.startswith("@"):
            append(line)
            if S["busy"] == "build" and line.strip():
                status.set(line.strip()[:90])
            return
        tag, _, text = line[1:].partition(" ")
        if tag == "progress":
            done, total, text = (text.split(" ", 2) + ["", ""])[:3]
            bar.configure(mode="determinate", value=1000 * int(done) / max(int(total), 1))
            status.set(text)
        elif tag in ("stage", "log", "warn", "error", "done"):
            append(("==> " + text) if tag == "stage" else text, tag if tag != "log" else None)
            if tag in ("stage", "done"):
                status.set(text)
            if tag == "error":
                S["error"] = text
            if tag == "warn":
                toggle_log(True)

    def finished(rc):
        kind = S["busy"]
        S["busy"], S["proc"] = None, None
        if kind == "build":
            bar.stop()
            bar.configure(mode="determinate", value=0)
            if rc == 0 and os.path.isfile(DEFAULT_IMAGE):
                status.set("Image built.")
                image_var.set(DEFAULT_IMAGE)
                refresh_image()
            else:
                fail(f"The build failed (make exit {rc}); see the log.")
        elif rc == 0:
            bar.configure(value=1000)
            messagebox.showinfo("MontaukOS Imager", status.get(), parent=root)
        else:
            err = S.pop("error", None)
            if rc == 126 and not err:
                status.set("Authorization was cancelled.")
            elif rc == 127 and not err:
                fail("Could not get administrator rights through pkexec. "
                     "Run the imager with sudo instead.")
            else:
                fail(err or f"Failed (exit {rc}); see the log.")
        S["sig"] = None
        refresh(force=True)

    def privileged(args):
        base = [sys.executable, "-I", SCRIPT, "--machine"] + args + ["--yes"]
        if os.geteuid() == 0:
            return base
        if not shutil.which("pkexec"):
            return None
        return ["pkexec"] + base

    def start():
        card = selected()
        path = os.path.abspath(image_var.get().strip())
        if not card:
            return
        if not os.path.isfile(path):
            fail(f"The image {path} does not exist.")
            return
        what = f"{card['model']} ({card['path']}, {human(card['size'])})"
        if action.get() == "flash":
            if card["kind"] != "empty":
                detail = f"It currently holds: {card['contents']}.\n\n"
            else:
                detail = ""
            if not messagebox.askyesno(
                    "Erase the card?",
                    f"Everything on {what} will be erased and replaced with MontaukOS.\n\n{detail}"
                    "Continue?", icon="warning", default="no", parent=root):
                return
            args = ["flash", path, card["path"]]
            if not verify.get():
                args.append("--no-verify")
            if flash_expand.get():
                args.append("--expand")
        else:
            if S["layout"] and card["kind"] != S["layout"]:
                fail("The card and the image use different layouts "
                     "(ext2 system partition vs single ramdisk partition). Write a fresh image instead.")
                return
            if not messagebox.askyesno(
                    "Update MontaukOS?",
                    f"MontaukOS on {what} will be updated from\n{os.path.basename(path)}.\n\n"
                    "Users, their files and settings are kept. Do not remove the card "
                    "until the update has finished.\n\nContinue?", parent=root):
                return
            args = ["update", path, card["path"]]
            if not merge_cfg.get():
                args.append("--keep-config")
            if firmware.get():
                args.append("--firmware")
            if prune.get():
                args.append("--prune")
            if upd_expand.get():
                args.append("--expand")
        cmd = privileged(args)
        if not cmd:
            fail("Writing a card needs administrator rights, but pkexec is not installed. "
                 "Run the imager with sudo.")
            return
        S.pop("error", None)
        status.set("Waiting for authorization..." if cmd[0] == "pkexec" else "Starting...")
        launch(args[0], cmd)

    def build():
        if not REPO:
            return
        status.set("Building the image (make ARCH=aarch64 rpi4-sd)...")
        toggle_log(True)
        launch("build", ["make", "ARCH=aarch64", "rpi4-sd"], cwd=REPO)
        bar.configure(mode="indeterminate")
        bar.start(15)

    def cancel():
        proc = S["proc"]
        if not proc:
            return
        if S["busy"] == "build":
            proc.terminate()
            bar.stop()
            return
        if S["busy"] == "flash" and messagebox.askyesno(
                "Cancel writing?", "Stop writing the card? It will not boot until it is written again.",
                parent=root):
            with contextlib.suppress(OSError):
                proc.stdin.write("cancel\n")
                proc.stdin.flush()
            status.set("Cancelling...")

    def on_close():
        if S["busy"] == "update" or S["busy"] == "expand":
            messagebox.showwarning("MontaukOS Imager",
                                   "The card is being updated. Wait until it has finished.", parent=root)
            return
        if S["busy"] == "flash":
            if not messagebox.askyesno("Quit?", "The card is still being written. Stop and quit?",
                                       parent=root):
                return
            with contextlib.suppress(OSError):
                S["proc"].stdin.write("cancel\n")
                S["proc"].stdin.flush()
        if S["busy"] == "build":
            S["proc"].terminate()
        root.destroy()

    tree.bind("<<TreeviewSelect>>", update_ui)
    image_var.trace_add("write", refresh_image)
    root.protocol("WM_DELETE_WINDOW", on_close)
    refresh_image()
    refresh(force=True)
    root.after(50, pump)
    root.after(2000, poll_cards)
    root.mainloop()


# ==== Entry ====

def main():
    ap = argparse.ArgumentParser(description="Write and update MontaukOS SD cards for the Raspberry Pi 4/400.")
    ap.add_argument("--machine", action="store_true", help=argparse.SUPPRESS)
    ap.add_argument("--image", help="image to preselect in the GUI")
    sub = ap.add_subparsers(dest="cmd")
    p = sub.add_parser("list", help="list SD cards and USB drives")
    p.add_argument("--all", action="store_true", help="include fixed disks")
    p = sub.add_parser("flash", help="write an image over the whole card")
    p.add_argument("image")
    p.add_argument("device")
    p.add_argument("--no-verify", action="store_true", help="do not read the card back")
    p.add_argument("--expand", action="store_true", help="grow the system partition to fill the card")
    p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
    p = sub.add_parser("update", help="update the installation on a card, keeping user data")
    p.add_argument("image")
    p.add_argument("device")
    p.add_argument("--keep-config", action="store_true", help="leave init.toml and capabilities.toml alone")
    p.add_argument("--firmware", action="store_true", help="also replace RPI_EFI.fd (resets UEFI settings)")
    p.add_argument("--prune", action="store_true", help="delete system files the image no longer ships")
    p.add_argument("--expand", action="store_true", help="grow the system partition to fill the card first")
    p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
    p = sub.add_parser("expand", help="grow the system partition to fill the card")
    p.add_argument("device")
    p.add_argument("--yes", action="store_true", help="do not ask for confirmation")
    args = ap.parse_args()

    if args.cmd is None:
        gui_main(args.image)
        return
    rep = Reporter(args.machine)
    if args.cmd == "list":
        cmd_list(rep, args)
        return
    if os.geteuid() != 0:
        rep.error("writing a card needs root: run with sudo (the GUI asks through pkexec)")
        sys.exit(1)
    if args.machine:
        watch_stdin(rep)
    try:
        {"flash": cmd_flash, "update": cmd_update, "expand": cmd_expand}[args.cmd](rep, args)
    except Fail as e:
        rep.error(str(e))
        sys.exit(1)
    except KeyboardInterrupt:
        rep.error("interrupted")
        sys.exit(130)
    except OSError as e:
        rep.error(f"{e.filename + ': ' if e.filename else ''}{e.strerror or e}")
        sys.exit(1)


if __name__ == "__main__":
    main()