#!/usr/bin/env python3
# build: assemble the installable GRUB theme into dist/udt/.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
# Licensed under the GNU General Public License v2 only.
"""Assemble dist/udt/ from the theme.txt udt-palette rendered.

Usage: build [--selftest]

Writes the PNG slices, the .pf2 fonts and a GRUB-readable copy of the
current wallpaper next to theme.txt. Everything is staged in a temporary
directory and swapped in at the end, so a failed run leaves the previous
dist/udt/ untouched. Never touches /boot: root copies dist/udt/ there.
"""
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile
import zlib
from pathlib import Path

PNG_SIG = b"\x89PNG\r\n\x1a\n"
REPO = Path(__file__).resolve().parent
PLACEHOLDERS = {f"@{r}@" for r in (
    "SCHEME", "BG", "BG_ALT", "BG_DEEP", "SURFACE", "FG", "FG_DIM", "FG_FAINT", "ACCENT")}
# file name: (fc-match pattern, size, the name grub-mkfont records and theme.txt uses)
FONTS = {
    "noto-sans-regular-24.pf2": ("Noto Sans:regular", 24, "Noto Sans Regular 24"),
    "noto-sans-bold-32.pf2": ("Noto Sans:bold", 32, "Noto Sans Bold 32"),
    "inconsolata-18.pf2": ("Inconsolata Nerd Font Mono", 18,
                           "Inconsolata Nerd Font Mono Regular 18"),
}
RENDERED = REPO / "theme.txt"
BUILD = REPO / "dist"
OUT = BUILD / "udt"
FONT_CACHE = BUILD / "fonts"  # fonts do not follow the scheme, so built once
WALLPAPER = Path.home() / ".cache/udt/wpaper"
PANEL_ALPHA = 0xd9  # ~85%, enough to read text over any wallpaper
MAX_SIDE = 2560  # widest mode in use; GRUB crops, a bigger image only costs /boot


def die(msg):
    sys.exit(f"build: {msg}")


def read_colours(text):
    """The colours on the '# udt:' line, as RGB tuples, from a rendered theme."""
    left = sorted(set(re.findall(r"@[A-Z_]+@", text)))
    if left:
        die(f"theme.txt still has placeholders {' '.join(left)}; rerun udt-palette")
    m = re.search(r"^# udt: (.+)$", text, re.M)
    if not m:
        die("theme.txt has no '# udt:' colour line")
    return {k: tuple(bytes.fromhex(v.lstrip("#")))
            for k, v in (kv.split("=") for kv in m.group(1).split())}


def chunk(tag, data):
    return (struct.pack(">I", len(data)) + tag + data
            + struct.pack(">I", zlib.crc32(tag + data)))


def png(w, h, rgba):
    """A w x h PNG of one colour: 8-bit RGBA, not interlaced, which GRUB reads."""
    row = b"\0" + bytes(rgba) * w  # filter byte 0, then the pixels
    return (PNG_SIG
            + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 6, 0, 0, 0))
            + chunk(b"IDAT", zlib.compress(row * h))
            + chunk(b"IEND", b""))


def unpng(data):
    """Decode what png() writes: (width, height, set of RGBA pixels)."""
    assert data[:8] == PNG_SIG
    pos, idat = 8, b""
    while pos < len(data):
        n, tag = struct.unpack(">I4s", data[pos:pos + 8])
        body = data[pos + 8:pos + 8 + n]
        assert struct.unpack(">I", data[pos + 8 + n:pos + 12 + n])[0] == zlib.crc32(tag + body)
        if tag == b"IHDR":
            w, h = struct.unpack(">II", body[:8])
        elif tag == b"IDAT":
            idat += body
        pos += 12 + n
    raw = zlib.decompress(idat)
    stride = 1 + 4 * w
    px = set()
    for y in range(h):
        row = raw[y * stride:(y + 1) * stride]
        assert row[0] == 0
        px |= {tuple(row[1 + 4 * x:5 + 4 * x]) for x in range(w)}
    return w, h, px


def run(*cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode:
        die(f"{cmd[0]} failed:\n{p.stderr.strip()}")
    return p.stdout


def pf2_name(path):
    """The NAME section of a .pf2, which is what theme.txt refers to."""
    data = path.read_bytes()
    pos = 0
    while pos < len(data):
        tag, n = struct.unpack(">4sI", data[pos:pos + 8])
        if tag == b"NAME":
            return data[pos + 8:pos + 8 + n].rstrip(b"\0").decode()
        pos += 8 + n
    die(f"{path.name} has no NAME section")


def font(fname, pattern, size, name):
    """A cached .pf2, built from the TTF fontconfig resolves the pattern to."""
    cached = FONT_CACHE / fname
    if cached.exists():
        return cached
    family = pattern.split(":")[0]
    # fc-match never fails: a missing family comes back as a fallback font.
    if family not in run("fc-match", "-f", "%{family}", pattern).split(","):
        die(f"font {family!r} is not installed")
    ttf = run("fc-match", "-f", "%{file}", pattern)
    FONT_CACHE.mkdir(parents=True, exist_ok=True)
    tmp = cached.with_suffix(".tmp")
    run("grub-mkfont", "-s", str(size), "-o", str(tmp), ttf)
    if pf2_name(tmp) != name:
        die(f"{fname} is named {pf2_name(tmp)!r}, theme.txt expects {name!r}")
    tmp.rename(cached)
    return cached


def os_name():
    for line in Path("/etc/os-release").read_text().splitlines():
        if line.startswith("NAME="):
            return line[5:].strip('"')
    return "GNU/Linux"


def assemble(stage, text, c, wallpaper, fonts):
    (stage / "theme.txt").write_text(text.replace("%TITLE%", os_name()))
    slices = {
        "panel.png": (8, 8, (*c["bg_alt"], PANEL_ALPHA)),
        # The selected item: a 4px accent bar on the left, surface behind.
        "select_w.png": (4, 1, (*c["accent"], 255)),
        "select_c.png": (1, 1, (*c["surface"], 255)),
        "select_e.png": (8, 1, (*c["surface"], 255)),
        # The countdown bar: see theme.txt.in for why the top slice is clear.
        "bar_n.png": (1, 24, (0, 0, 0, 0)),
        "bar_c.png": (1, 1, (*c["surface"], 255)),
        "hl_c.png": (1, 1, (*c["accent"], 255)),
    }
    for s in ("nw", "n", "ne", "w", "c", "e", "sw", "s", "se"):
        slices[f"terminal_box_{s}.png"] = (1, 1, (*c["bg_deep"], 255))
    for fname, (w, h, rgba) in slices.items():
        (stage / fname).write_bytes(png(w, h, rgba))
    for f in fonts:
        shutil.copy(f, stage / f.name)
    # [0]: first frame only. Baseline JPEG, alpha flattened: GRUB reads
    # neither progressive JPEG nor webp, and has no use for transparency here.
    bg = "#%02x%02x%02x" % c["bg_deep"]
    run("magick", f"{wallpaper}[0]", "-auto-orient",
        "-background", bg, "-alpha", "remove",
        "-resize", f"{MAX_SIDE}x{MAX_SIDE}>", "-strip",
        "-interlace", "none", "-quality", "90", str(stage / "background.jpg"))


def selftest():
    w, h, px = unpng(png(3, 2, (0x12, 0x34, 0x56, 0xd9)))
    assert (w, h, px) == (3, 2, {(0x12, 0x34, 0x56, 0xd9)}), (w, h, px)
    good = "# udt: bg_alt=#1e2030 accent=#b7bdf8\ncolor = \"#cad3f5\"\n"
    assert read_colours(good) == {"bg_alt": (0x1e, 0x20, 0x30), "accent": (0xb7, 0xbd, 0xf8)}
    for bad in (good + 'color = "@FG@"\n', 'color = "#cad3f5"\n'):
        try:
            read_colours(bad)
        except SystemExit:
            pass
        else:
            raise AssertionError(f"accepted: {bad!r}")
    tpl = (REPO / "theme.txt.in").read_text()
    assert set(re.findall(r"@[A-Z_]+@", tpl)) <= PLACEHOLDERS, "template token udt does not fill"
    for _, _, name in FONTS.values():
        assert f'"{name}"' in tpl, f"template does not use font {name}"
    print("selftest ok")


def main():
    if sys.argv[1:] == ["--selftest"]:
        return selftest()
    for tool in ("magick", "grub-mkfont", "fc-match"):
        if not shutil.which(tool):
            die(f"{tool} not found")
    if not RENDERED.exists():
        die("theme.txt missing; run udt-palette first")
    text = RENDERED.read_text()
    c = read_colours(text)
    try:
        wallpaper = WALLPAPER.resolve(strict=True)
    except FileNotFoundError:
        die(f"no wallpaper at {WALLPAPER}")
    fonts = [font(f, *spec) for f, spec in FONTS.items()]
    BUILD.mkdir(exist_ok=True)
    stage = Path(tempfile.mkdtemp(dir=BUILD, prefix=".udt-"))
    try:
        assemble(stage, text, c, wallpaper, fonts)
        os.chmod(stage, 0o755)  # mkdtemp makes it 0700; root copies it as is
        shutil.rmtree(OUT, ignore_errors=True)
        stage.rename(OUT)
    finally:
        shutil.rmtree(stage, ignore_errors=True)
    print(f"build: {OUT}")


if __name__ == "__main__":
    main()
