#!/usr/bin/env python3
"""Maintain Limine entries for Gentoo kernel installs and Portage snapshots."""

from __future__ import annotations

import fcntl
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path

CONFIG = Path("/boot/limine.conf")
BOOT = Path("/boot")
MODULE_ROOTS = (Path("/usr/lib/modules"), Path("/lib/modules"))
SNAPSHOT_ROOT = Path("/.snapshots/portage")
CLONE_ROOT = Path("/.snapshots/bootable")
MARKER_START = "# BEGIN PORTAGE SNAPSHOT "
MARKER_END = "# END PORTAGE SNAPSHOT "
ENTRY_RE = re.compile(r"(?m)^(/+)([^\r\n]*)\r?\n")


class MenuError(RuntimeError):
    pass


@dataclass
class Entry:
    depth: int
    title: str
    start: int
    end: int
    block: str


def run(argv: list[str], *, capture: bool = True) -> str:
    result = subprocess.run(
        argv,
        check=False,
        text=True,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.PIPE if capture else None,
    )
    if result.returncode:
        details = (result.stderr or result.stdout or "").strip()
        raise MenuError(f"command failed ({result.returncode}): {' '.join(argv)}\n{details}")
    return result.stdout or ""


def split_entries(text: str) -> tuple[str, list[Entry]]:
    matches = list(ENTRY_RE.finditer(text))
    if not matches:
        raise MenuError("Limine configuration has no menu entries")
    entries: list[Entry] = []
    for index, match in enumerate(matches):
        end = matches[index + 1].start() if index + 1 < len(matches) else len(text)
        entries.append(
            Entry(
                depth=len(match.group(1)),
                title=match.group(2),
                start=match.start(),
                end=end,
                block=text[match.start():end],
            )
        )
    return text[:matches[0].start()], entries


def value(entry: Entry, key: str) -> str | None:
    match = re.search(rf"(?mi)^\s*{re.escape(key)}\s*:\s*(.*?)\s*$", entry.block)
    return match.group(1) if match else None


def all_entries(text: str) -> list[Entry]:
    return split_entries(text)[1]


def children_of(entries: list[Entry], parent: Entry) -> list[Entry]:
    try:
        index = entries.index(parent)
    except ValueError:
        return []
    children: list[Entry] = []
    for entry in entries[index + 1:]:
        if entry.depth <= parent.depth:
            break
        children.append(entry)
    return children


def kernel_release(kernel_path: str | None) -> str | None:
    if not kernel_path:
        return None
    path = kernel_path.rsplit(":", 1)[-1]
    name = Path(path).name
    if not name.startswith("vmlinuz-") or name.endswith(".old"):
        return None
    return name.removeprefix("vmlinuz-")


def root_subvol(cmdline: str | None) -> str | None:
    if not cmdline:
        return None
    match = re.search(r"(?:^|\s)rootflags=([^\s]+)", cmdline)
    if not match:
        return None
    for option in match.group(1).split(","):
        if option.startswith("subvol="):
            return option.removeprefix("subvol=")
    return None


def read_config() -> str:
    try:
        return CONFIG.read_text(encoding="utf-8")
    except OSError as error:
        raise MenuError(f"cannot read {CONFIG}: {error}") from error


def backup_config() -> None:
    backup = Path(str(CONFIG) + ".pre-portage-recovery-menu")
    if not backup.exists():
        shutil.copy2(CONFIG, backup)
        print(f"Saved original Limine config: {backup}")


def write_config(text: str) -> None:
    if not text.endswith("\n"):
        text += "\n"
    validate_config(text)
    backup_config()
    fd, temporary = tempfile.mkstemp(prefix="limine.conf.portage.", dir=CONFIG.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as stream:
            stream.write(text)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, CONFIG)
    except BaseException:
        try:
            os.unlink(temporary)
        except OSError:
            pass
        raise


def validate_config(text: str) -> None:
    prefix, entries = split_entries(text)
    if not re.search(r"(?mi)^\s*timeout\s*:\s*", prefix):
        raise MenuError("Limine config validation failed: global timeout is missing")
    roots = [entry for entry in entries if entry.depth == 1]
    if not roots:
        raise MenuError("Limine config validation failed: no top-level entries")
    current_path = None
    default = re.search(r"(?mi)^\s*default_entry\s*:\s*(.*?)\s*$", prefix)
    if default:
        current_path = default.group(1).split("#", 1)[0]
    if any(children_of(entries, root) for root in roots) and not current_path:
        raise MenuError("Limine config validation failed: nested kernel entries require an explicit default_entry path")
    if current_path:
        if current_path.isdecimal():
            raise MenuError("Limine config validation failed: use a leaf entry path for default_entry when kernel directories exist")
        parts = current_path.split("/")
        parent_title = parts[0]
        child_title = parts[1] if len(parts) > 1 else None
        parent = next((entry for entry in roots if entry.title == parent_title), None)
        if parent is None:
            raise MenuError(f"Limine config validation failed: default parent is missing: {parent_title}")
        if child_title and not any(
            entry.depth == 2 and entry.title == child_title
            for entry in children_of(entries, parent)
        ):
            raise MenuError(f"Limine config validation failed: default subentry is missing: {current_path}")
        if child_title is None and children_of(entries, parent):
            raise MenuError(f"Limine config validation failed: default_entry points to a directory: {current_path}")
    for entry in entries:
        if entry.depth == 2 and entry.title.startswith("Portage "):
            for key in ("protocol", "kernel_path", "module_path", "cmdline"):
                if value(entry, key) is None:
                    raise MenuError(f"Limine config validation failed: generated entry {entry.title!r} lacks {key}")


def migrate() -> None:
    text = read_config()
    prefix, entries = split_entries(text)
    if any(entry.depth > 1 for entry in entries):
        validate_config(text)
        if kernel_groups(text):
            print("Limine kernel groups are already configured.")
            return
        raise MenuError("Limine config has nested entries but no managed kernel groups; refusing migration")

    kernels: list[Entry] = []
    recovery: list[Entry] = []
    other: list[Entry] = []
    for entry in entries:
        if value(entry, "protocol") != "linux" or not value(entry, "kernel_path"):
            other.append(entry)
            continue
        root = root_subvol(value(entry, "cmdline"))
        if root == "@":
            kernels.append(entry)
        elif root and root.startswith("@snapshots/bootable/"):
            recovery.append(entry)
        else:
            other.append(entry)

    dist = next((entry for entry in kernels if "Dist Kernel" in entry.title), None)
    if dist is None:
        raise MenuError("expected Gentoo Dist Kernel entry was not found; refusing to change the boot menu")
    if not any("CachyOS Kernel" in entry.title for entry in kernels):
        raise MenuError("expected CachyOS Kernel entry was not found; refusing to change the boot menu")

    groups: dict[str, list[str]] = {}
    releases: dict[str, str] = {}
    for entry in kernels:
        title = entry.title
        release = kernel_release(value(entry, "kernel_path"))
        if not release:
            raise MenuError(f"cannot derive kernel release for {title}")
        if title in groups:
            raise MenuError(f"duplicate kernel group title: {title}")
        body = entry.block[entry.block.find("\n") + 1:]
        groups[title] = [f"/{title}\n", f"//Current boot\n{body}"]
        releases[release] = title

    for entry in recovery:
        release = kernel_release(value(entry, "kernel_path"))
        title = releases.get(release or "")
        if title is None:
            other.append(entry)
            continue
        body = entry.block[entry.block.find("\n") + 1:]
        child_title = entry.title
        match = re.match(r"Gentoo Linux \((.*)\)", child_title)
        if match:
            child_title = match.group(1)
        groups[title].append(f"//{child_title}\n{body}")

    result_prefix = prefix
    if re.search(r"(?mi)^\s*default_entry\s*:", result_prefix):
        result_prefix = re.sub(
            r"(?mi)^\s*default_entry\s*:\s*[^\r\n]*",
            "default_entry: Gentoo Linux (Dist Kernel)/Current boot",
            result_prefix,
            count=1,
        )
    else:
        result_prefix = result_prefix.rstrip() + "\ndefault_entry: Gentoo Linux (Dist Kernel)/Current boot\n\n"

    pieces = [result_prefix.rstrip() + "\n\n"]
    written: set[str] = set()
    for entry in entries:
        if entry.title in groups and entry.title not in written:
            pieces.extend(groups[entry.title])
            pieces.append("\n")
            written.add(entry.title)
        elif entry in other:
            pieces.append(entry.block)
    for title, group in groups.items():
        if title not in written:
            pieces.extend(group)
            pieces.append("\n")
    write_config("".join(pieces))
    print("Migrated existing kernel entries into Limine kernel groups; Dist Kernel remains the default.")


def kernel_groups(text: str) -> list[tuple[Entry, Entry, str]]:
    _, entries = split_entries(text)
    groups: list[tuple[Entry, Entry, str]] = []
    for parent in entries:
        if parent.depth != 1:
            continue
        for child in children_of(entries, parent):
            if child.depth != 2 or child.title != "Current boot":
                continue
            release = kernel_release(value(child, "kernel_path"))
            if release:
                groups.append((parent, child, release))
    return groups


def option_block(title: str, kernel: str, module_paths: list[str], cmdline: str, depth: int = 2) -> str:
    slashes = "/" * depth
    content = (
        f"{slashes}{title}\n"
        "    protocol: linux\n"
        f"    kernel_path: boot():/{kernel}\n"
    )
    content += "".join(f"    module_path: boot():/{module}\n" for module in module_paths)
    return content + f"    cmdline: {cmdline}\n"


def module_releases(root: Path) -> set[str]:
    releases: set[str] = set()
    for base in (root / "usr/lib/modules", root / "lib/modules"):
        try:
            releases.update(path.name for path in base.iterdir() if path.is_dir())
        except OSError:
            pass
    return releases


def compatible_groups(root: Path, text: str) -> list[tuple[Entry, Entry, str]]:
    releases = module_releases(root)
    compatible: list[tuple[Entry, Entry, str]] = []
    for group, current, release in kernel_groups(text):
        kernel_path = value(current, "kernel_path")
        module_paths = re.findall(r"(?mi)^\s*module_path\s*:\s*(.*?)\s*$", current.block)
        if release not in releases or not kernel_path or not module_paths:
            continue
        if not boot_asset_exists(kernel_path) or not all(boot_asset_exists(path) for path in module_paths):
            continue
        compatible.append((group, current, release))
    return compatible


def boot_asset_exists(limine_path: str) -> bool:
    if not limine_path.startswith("boot():/"):
        return False
    return (BOOT / limine_path.removeprefix("boot():/")).is_file()


def legacy_clone_for(source: Path, group: Entry, text: str) -> tuple[Path, str] | None:
    """Reuse a pre-manager recovery clone when it clearly matches this package."""
    lowered = source.name.lower()
    token = "fastfetch" if "install-app-misc_fastfetch" in lowered else None
    if token is None:
        return None
    _, entries = split_entries(text)
    fresh_group = next((entry for entry in entries if entry.depth == 1 and entry.title == group.title), None)
    if fresh_group is None:
        return None
    for child in children_of(entries, fresh_group):
        if child.depth != 2 or child.title == "Current boot" or child.title.startswith("Portage "):
            continue
        if token not in child.title.lower():
            continue
        subvol = root_subvol(value(child, "cmdline"))
        if not subvol or not subvol.startswith("@snapshots/bootable/"):
            continue
        clone = CLONE_ROOT / subvol.removeprefix("@snapshots/bootable/")
        if not clone.is_dir():
            continue
        try:
            if run(["btrfs", "property", "get", "-ts", str(clone), "ro"]).strip() == "ro=false":
                return clone, subvol
        except MenuError:
            continue
    return None


def marker_block(snapshot_name: str, release: str, clone_subvol: str, group_title: str, child_title: str) -> str:
    _, entries = split_entries(read_config())
    group = next((entry for entry in entries if entry.depth == 1 and entry.title == group_title), None)
    if group is None:
        raise MenuError(f"kernel group disappeared during menu update: {group_title}")
    current = next((entry for entry in children_of(entries, group) if entry.depth == 2 and entry.title == "Current boot"), None)
    if current is None:
        raise MenuError(f"kernel group has no Current boot subentry: {group.title}")
    cmdline = value(current, "cmdline")
    if not cmdline or not re.search(r"(?:^|\s)rootflags=[^\s]+", cmdline):
        raise MenuError(f"kernel group has no rootflags in its command line: {group.title}")
    cmdline = re.sub(r"((?:^|\s)rootflags=)[^\s]+", rf"\g<1>subvol={clone_subvol}", cmdline, count=1)
    modules = re.findall(r"(?mi)^\s*module_path\s*:\s*(.*?)\s*$", current.block)
    module_names = [Path(module.split(":", 1)[-1]).name for module in modules]
    body = option_block(child_title, Path(value(current, "kernel_path").split(":", 1)[-1]).name,
                        module_names, cmdline)
    marker_id = f"{snapshot_name} {release}"
    return f"{MARKER_START}{marker_id}\n{body}{MARKER_END}{marker_id}\n"


def clone_name(snapshot_name: str) -> str:
    digest = hashlib.sha256(snapshot_name.encode("utf-8")).hexdigest()[:20]
    return f"portage-{digest}"


def subvol_path(path: Path) -> str:
    listing = run(["btrfs", "subvolume", "list", "-a", "/"])
    wanted = path.name
    found: list[str] = []
    for line in listing.splitlines():
        if " path " not in line:
            continue
        subvol = line.split(" path ", 1)[1].strip()
        if subvol.startswith("<FS_TREE>/"):
            subvol = subvol[len("<FS_TREE>/"):]
        if subvol == wanted or subvol.endswith("/" + wanted):
            found.append(subvol)
    if len(found) != 1:
        raise MenuError(f"cannot uniquely resolve Btrfs top-level path for {path}")
    return found[0]


def writable_clone(source: Path) -> tuple[Path, str]:
    if not source.is_dir():
        raise MenuError(f"snapshot source is missing: {source}")
    ro = run(["btrfs", "property", "get", "-ts", str(source), "ro"]).strip()
    if ro != "ro=true":
        raise MenuError(f"refusing to clone a non-read-only source snapshot: {source} ({ro})")
    CLONE_ROOT.mkdir(mode=0o700, parents=True, exist_ok=True)
    target = CLONE_ROOT / clone_name(source.name)
    if not target.exists():
        available = shutil.disk_usage("/.snapshots").free
        if available < 2 * 1024**3:
            raise MenuError("less than 2 GiB is free; refusing to create a writable snapshot clone")
        run(["btrfs", "subvolume", "snapshot", str(source), str(target)], capture=False)
        print(f"Created writable recovery clone: {target}")
    clone_ro = run(["btrfs", "property", "get", "-ts", str(target), "ro"]).strip()
    if clone_ro != "ro=false":
        raise MenuError(f"recovery clone is not writable: {target} ({clone_ro})")
    return target, subvol_path(target)


def insert_child(text: str, parent_title: str, child: str) -> str:
    prefix, entries = split_entries(text)
    parent_index = next((i for i, entry in enumerate(entries) if entry.depth == 1 and entry.title == parent_title), None)
    if parent_index is None:
        raise MenuError(f"kernel group disappeared during menu update: {parent_title}")
    insert_at = len(text)
    for entry in entries[parent_index + 1:]:
        if entry.depth == 1:
            insert_at = entry.start
            break
    before = text[:insert_at]
    if before and not before.endswith("\n\n"):
        before = before.rstrip("\n") + "\n\n"
    after = text[insert_at:]
    if after and not child.endswith("\n\n"):
        child += "\n"
    return before + child + after


def add_snapshot(source: Path) -> bool:
    name = source.name
    text = read_config()
    compatible = compatible_groups(source, text)
    if not compatible:
        print(f"Snapshot has no matching kernel/initramfs menu group; keeping snapshot only: {source}", file=sys.stderr)
        return False
    adopted: tuple[Path, str] | None = None
    for group, _current, _release in compatible:
        adopted = legacy_clone_for(source, group, text)
        if adopted:
            break
    clone, subvol = adopted if adopted else writable_clone(source)
    changed = False
    for group, _current, _release in compatible:
        marker = f"{MARKER_START}{name} {_release}\n"
        if marker in text:
            continue
        if legacy_clone_for(source, group, text):
            print(f"Existing recovery entry under {group.title} already covers {source.name}")
            continue
        title = f"Portage {name.removeprefix('portage-')}"
        child = marker_block(name, _release, subvol, group.title, title)
        text = insert_child(text, group.title, child)
        changed = True
    if changed:
        write_config(text)
        print(f"Added Limine recovery option(s) for {source.name} using {clone.name}")
    return changed


def remove_snapshot(source: Path) -> None:
    name = source.name
    text = read_config()
    pattern = re.compile(
        rf"(?ms)^# BEGIN PORTAGE SNAPSHOT {re.escape(name)} [^\n]+\n.*?^# END PORTAGE SNAPSHOT {re.escape(name)} [^\n]+\n?"
    )
    new_text, count = pattern.subn("", text)
    if count:
        write_config(new_text)
        print(f"Removed {count} Limine recovery option(s) for expired snapshot {name}")
    clone = CLONE_ROOT / clone_name(name)
    if clone.exists():
        show = run(["btrfs", "subvolume", "show", str(clone)])
        if "Subvolume ID:" not in show:
            raise MenuError(f"refusing to delete unverified clone path: {clone}")
        run(["btrfs", "subvolume", "delete", str(clone)], capture=False)
        print(f"Deleted expired writable recovery clone: {clone}")


def ensure_kernel_groups() -> int:
    text = read_config()
    groups = kernel_groups(text)
    releases = module_releases(Path("/"))
    created = 0
    template = next((entry for _group, entry, _release in groups if value(entry, "cmdline")), None)
    if template is None:
        raise MenuError("cannot add kernel groups without an existing Current boot command line")
    template_cmdline = value(template, "cmdline") or ""
    for release in sorted(releases):
        kernel = BOOT / f"vmlinuz-{release}"
        initramfs = BOOT / f"initramfs-{release}.img"
        if not kernel.is_file():
            continue
        if not initramfs.is_file():
            print(f"Kernel {release} has no matching initramfs at {initramfs}; no Limine entry was created.", file=sys.stderr)
            continue
        text = read_config()
        groups = kernel_groups(text)
        if any(existing == release for _group, _entry, existing in groups):
            continue
        title = f"Gentoo Linux (Kernel {release})"
        if any(entry.depth == 1 and entry.title == title for entry in all_entries(text)):
            continue
        cmdline = template_cmdline
        if not re.search(r"(?:^|\s)rootflags=", cmdline):
            raise MenuError("existing kernel template has no rootflags; refusing to invent a root command line")
        body = option_block("Current boot", kernel.name, [initramfs.name], cmdline)
        text = text.rstrip() + f"\n\n/{title}\n{body}\n"
        write_config(text)
        created += 1
        print(f"Added Limine kernel group for {release}")

    text = read_config()
    try:
        snapshots = sorted(
            path for path in SNAPSHOT_ROOT.iterdir()
            if path.is_dir() and path.name.startswith("portage-")
        )
    except OSError:
        snapshots = []
    for snapshot in snapshots:
        add_snapshot(snapshot)
    return created


def reconcile() -> None:
    ensure_kernel_groups()


def main(argv: list[str]) -> int:
    if os.geteuid() != 0:
        raise MenuError("must run as root from the Portage hook or installer")
    if len(argv) < 2 or argv[1] not in {"migrate", "reconcile", "sync-kernels", "add-snapshot", "remove-snapshot"}:
        raise MenuError("usage: portage-limine-menu {migrate|reconcile|sync-kernels|add-snapshot|remove-snapshot} [SNAPSHOT]")
    action = argv[1]
    os.makedirs("/run/lock", exist_ok=True)
    with open("/run/lock/portage-limine-menu.lock", "w", encoding="ascii") as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        if action == "migrate":
            migrate()
        elif action == "reconcile":
            migrate_if_needed()
            reconcile()
        elif action == "sync-kernels":
            ensure_kernel_groups()
        elif action in {"add-snapshot", "remove-snapshot"}:
            migrate_if_needed()
            source = Path(argv[2]) if len(argv) > 2 else None
            if source is None:
                raise MenuError(f"{action} requires a snapshot path")
            if action == "add-snapshot":
                add_snapshot(source)
            else:
                remove_snapshot(source)
    return 0


def migrate_if_needed() -> None:
    text = read_config()
    _, entries = split_entries(text)
    if not any(entry.depth > 1 for entry in entries):
        migrate()
        return
    validate_config(text)
    if not kernel_groups(text):
        raise MenuError("Limine menu is nested but has no managed kernel groups")


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv))
    except (MenuError, OSError, subprocess.SubprocessError) as error:
        print(f"portage-limine-menu: {error}", file=sys.stderr)
        sys.exit(1)
