#!/usr/bin/python3
"""linuxdoors-update-grub -- make the newest LinuxDoors kernel boot by default.

Rewrites /boot/grub/grub.cfg on an installed LinuxDoors system so that:
  * the newest stock kernel (/boot/vmlinuz-<release>-linuxdoors whose
    /lib/modules/<release> is present) is entry 0, the default;
  * every other stock kernel stays as a fallback entry, newest first;
  * kernels built with LinuxDoors Kernel Config (release ending -customN)
    stay as their own "Custom Kernel" entries, never made the default;
  * menu entries that don't boot a LinuxDoors kernel (another OS, etc.) are
    kept unchanged;
  * the kernel command line (root=, rw, console, init, ...) is taken from the
    current default entry, or from --cmdline when there is no grub.cfg yet.

It will NOT move the default away from the running kernel when that kernel
has out-of-tree modules (e.g. the NVIDIA proprietary driver) that the newer
kernel lacks -- booting the newer one would lose that driver. It does nothing
at all when /etc/linuxdoors/no-auto-kernel exists, or on the live ISO.

Run by: the LinuxDoors installer, the package manager after a kernel install,
and linuxdoors-update-grub.service once per boot.

  linuxdoors-update-grub [--root DIR] [--cmdline "root=... rw ..."] [--dry-run]
"""
import argparse
import os
import re
import shutil
import subprocess
import sys

STOCK_RE = re.compile(r"^vmlinuz-(\d+(?:\.\d+)*)-linuxdoors$")
ANY_RE = re.compile(r"^vmlinuz-(.+)-linuxdoors$")
OOT_DIRS = ("extra", "updates", "video")
HEADER = ("insmod part_gpt\n"
          "insmod part_msdos\n"
          "insmod ext2\n"
          "insmod normal\n\n"
          "set timeout=5\n"
          "set default=0\n")
DEFAULT_ARGS = "rw console=tty0 init=/usr/lib/systemd/systemd"


def vkey(release):
    return tuple(int(x) for x in re.findall(r"\d+", release))


def out_of_tree_modules(root, release):
    """Module file names under extra/, updates/ or kernel/drivers/video/nvidia*."""
    names = set()
    base = os.path.join(root, "lib/modules", release)
    for d in OOT_DIRS:
        for top in (os.path.join(base, d), os.path.join(base, "kernel/drivers", d)):
            if not os.path.isdir(top):
                continue
            for dirpath, _, files in os.walk(top):
                for f in files:
                    if ".ko" in f and (d != "video" or f.startswith("nvidia")):
                        names.add(f)
    return names


def split_entries(text):
    """Return (header_text, [entry_text, ...]) for top-level menuentry blocks."""
    entries, i, header_end = [], 0, None
    for m in re.finditer(r"^menuentry\b", text, re.M):
        if m.start() < i:
            continue
        if header_end is None:
            header_end = m.start()
        depth, j = 0, text.index("{", m.start())
        while j < len(text):
            if text[j] == "{":
                depth += 1
            elif text[j] == "}":
                depth -= 1
                if depth == 0:
                    break
            j += 1
        entries.append(text[m.start():j + 1])
        i = j + 1
    return (text if header_end is None else text[:header_end]), entries


def boots_linuxdoors(entry):
    return re.search(r"^\s*linux\s+/boot/vmlinuz-\S+-linuxdoors\b", entry, re.M) is not None


def kernel_args(entry):
    m = re.search(r"^\s*linux\s+/boot/\S+\s+(.*)$", entry, re.M)
    return m.group(1).strip() if m else None


def entry(title, kernel, args):
    return f'menuentry "{title}" {{\n    linux /boot/{kernel} {args}\n}}'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", default="/")
    ap.add_argument("--cmdline", default=None)
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args()
    root = a.root
    cfg_path = os.path.join(root, "boot/grub/grub.cfg")

    if os.path.exists(os.path.join(root, "etc/linuxdoors/no-auto-kernel")):
        print("linuxdoors-update-grub: /etc/linuxdoors/no-auto-kernel present, leaving GRUB alone")
        return 0
    if root == "/":
        fs = subprocess.run(["findmnt", "-n", "-o", "FSTYPE", "/"],
                            capture_output=True, text=True).stdout.strip()
        if fs in ("overlay", "squashfs", "tmpfs"):
            print("linuxdoors-update-grub: live session, nothing to do")
            return 0

    stock, custom = [], []
    for f in os.listdir(os.path.join(root, "boot")):
        m = ANY_RE.match(f)
        if not m or not os.path.isdir(os.path.join(root, "lib/modules", m.group(1))):
            continue
        (stock if STOCK_RE.match(f) else custom).append((m.group(1), f))
    if not stock and not custom:
        print("linuxdoors-update-grub: no kernel with modules found in /boot, not touching GRUB", file=sys.stderr)
        return 1
    stock.sort(key=lambda r: vkey(r[0]), reverse=True)
    custom.sort(key=lambda r: vkey(r[0]), reverse=True)

    old_text = open(cfg_path, encoding="utf-8").read() if os.path.isfile(cfg_path) else ""
    header, old_entries = split_entries(old_text)
    keep = [e for e in old_entries if not boots_linuxdoors(e)]
    args = a.cmdline
    if not args:
        for e in old_entries:
            if boots_linuxdoors(e) and kernel_args(e):
                args = kernel_args(e)
                break
    if not args:
        print("linuxdoors-update-grub: no kernel command line (root=...) known; pass --cmdline", file=sys.stderr)
        return 1

    ordered = stock or custom
    default = ordered[0]
    # Keep the running kernel as default if it carries out-of-tree drivers the
    # newer one lacks (e.g. NVIDIA proprietary on a machine running CUDA).
    # LDUG_RUNNING lets tests stand in for the running kernel on a --root tree.
    running = os.environ.get("LDUG_RUNNING") or (os.uname().release if root == "/" else None)
    if running:
        if running != default[0] and any(r == running for r, _ in ordered):
            missing = out_of_tree_modules(root, running) - out_of_tree_modules(root, default[0])
            if missing:
                print(f"linuxdoors-update-grub: keeping {running} as default -- it has "
                      f"out-of-tree modules {sorted(missing)} that {default[0]} lacks")
                default = next(r for r in ordered if r[0] == running)

    blocks = [entry("LinuxDoors", default[1], args)]
    for rel, k in stock:
        if (rel, k) != default:
            blocks.append(entry(f"LinuxDoors (kernel {rel})", k, args))
    for rel, k in custom:
        if (rel, k) != default:
            blocks.append(entry(f"LinuxDoors (Custom Kernel {rel})", k, args))
    blocks += keep

    new_text = (HEADER if not header.strip() else header.rstrip("\n") + "\n") + "\n" + "\n\n".join(blocks) + "\n"
    print(f"linuxdoors-update-grub: default kernel {default[0]}"
          f" ({len(blocks)} menu entr{'y' if len(blocks) == 1 else 'ies'})")
    if new_text == old_text:
        print("linuxdoors-update-grub: grub.cfg already up to date")
        return 0
    if a.dry_run:
        print(new_text)
        return 0
    os.makedirs(os.path.dirname(cfg_path), exist_ok=True)
    if old_text:
        shutil.copy2(cfg_path, cfg_path + ".ldbak")
    tmp = cfg_path + ".ldnew"
    with open(tmp, "w", encoding="utf-8") as f:
        f.write(new_text)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, cfg_path)
    return 0


if __name__ == "__main__":
    sys.exit(main())
