#!/usr/bin/env python3
from __future__ import annotations

import argparse
import configparser
from dataclasses import dataclass, replace
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile

CONFIG_PATH = Path("/etc/catos/firmware-boot.conf")
CMDLINE_PATH = Path("/etc/kernel/cmdline")
INSTALL_MARKER = Path("/run/calamares/bootloadu-installing")
LOCK_PATH = Path("/run/lock/boot-partition.lock")
MODULES_PATH = Path("/usr/lib/modules")
SECURE_BOOT_KEY_PATH = Path("/var/lib/catos-secureboot/keys/machine.key")


class FirmwareBootError(RuntimeError):
    pass


@dataclass(frozen=True)
class Configuration:
    method: str
    esp_path: Path
    efi_disk: str
    efi_partition: int
    label_prefix: str
    default_kernel: str
    install_fallback: bool = False
    firmware_directory: str = "/"


@dataclass(frozen=True)
class Kernel:
    version: str
    package: str
    image: Path
    entry_id: str = ""

    @property
    def name(self) -> str:
        return self.entry_id or self.package


def run(arguments: list[str], *, capture: bool = False) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        arguments,
        check=True,
        text=True,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.PIPE if capture else None,
    )


def load_configuration(path: Path = CONFIG_PATH) -> Configuration:
    parser = configparser.ConfigParser(interpolation=None)
    if not parser.read(path):
        raise FirmwareBootError(f"configuration is missing: {path}")
    section = parser["boot"]
    method = section.get("method", "disabled").strip().lower()
    if method not in {"disabled", "uki", "efistub"}:
        raise FirmwareBootError(f"unsupported firmware boot method: {method}")
    try:
        partition = section.getint("efi_partition", fallback=0)
    except ValueError as error:
        raise FirmwareBootError("efi_partition must be an integer") from error
    return Configuration(
        method=method,
        esp_path=Path(section.get("esp_path", "/boot/efi")),
        efi_disk=section.get("efi_disk", "").strip(),
        efi_partition=partition,
        label_prefix=section.get("label_prefix", "CatOS").strip() or "CatOS",
        default_kernel=section.get("default_kernel", "linux").strip() or "linux",
        install_fallback=section.getboolean("install_fallback", fallback=False),
        firmware_directory=section.get("firmware_directory", "/"),
    )


def discover_kernels() -> list[Kernel]:
    kernels: list[Kernel] = []
    if not MODULES_PATH.is_dir():
        raise FirmwareBootError(f"kernel modules directory is missing: {MODULES_PATH}")
    packages: dict[str, int] = {}
    for directory in sorted(MODULES_PATH.iterdir(), key=lambda item: item.name):
        image = directory / "vmlinuz"
        if not image.is_file():
            continue
        pkgbase = directory / "pkgbase"
        package = pkgbase.read_text(encoding="utf-8").strip() if pkgbase.is_file() else directory.name
        if not re.fullmatch(r"[A-Za-z0-9._+\-]+", package):
            raise FirmwareBootError(f"invalid kernel package name: {pkgbase}")
        packages[package] = packages.get(package, 0) + 1
        kernels.append(Kernel(directory.name, package, image))
    if not kernels:
        raise FirmwareBootError("no installed kernels were found")
    return [replace(kernel, entry_id=f"{kernel.package}-{kernel.version}") if packages[kernel.package] > 1 else kernel for kernel in kernels]


def loader_path(path: str | Path) -> str:
    return "\\" + str(path).strip("/").replace("/", "\\")


@dataclass(frozen=True)
class FirmwareState:
    entries: dict[str, str]
    order: list[str]
    current: str
    next_boot: str


def firmware_state() -> FirmwareState:
    output = run(["efibootmgr", "--verbose"], capture=True).stdout
    entries: dict[str, str] = {}
    order: list[str] = []
    current = next_boot = ""
    for line in output.splitlines():
        match = re.match(r"^Boot([0-9A-Fa-f]{4})\*?\s+(.+)$", line)
        if match:
            entries[match[1].upper()] = match[2]
        elif line.startswith("BootOrder:"):
            order = [number.strip().upper() for number in line.partition(":")[2].split(",") if number.strip()]
        elif line.startswith("BootCurrent:"):
            current = line.partition(":")[2].strip().upper()
        elif line.startswith("BootNext:"):
            next_boot = line.partition(":")[2].strip().upper()
    return FirmwareState(entries, order, current, next_boot)


def digest(path: Path) -> str:
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def state_path(configuration: Configuration) -> Path:
    return configuration.esp_path / ".catos-firmware-boot" / f"{configuration.method}.json"


def owned_path(configuration: Configuration, relative: str) -> Path:
    path = configuration.esp_path / relative
    if Path(relative).is_absolute() or ".." in Path(relative).parts or not path.resolve().is_relative_to(configuration.esp_path.resolve()):
        raise FirmwareBootError(f"invalid managed output path: {relative}")
    return path


def load_state(configuration: Configuration) -> dict:
    path = state_path(configuration)
    if not path.exists():
        return {"version": 1, "files": {}, "entries": []}
    state = json.loads(path.read_text(encoding="utf-8"))
    if state.get("version") != 1 or not isinstance(state.get("files"), dict) or not isinstance(state.get("entries"), list):
        raise FirmwareBootError(f"invalid deployment record: {path}")
    for relative, checksum in state["files"].items():
        owned_path(configuration, relative)
        if not isinstance(checksum, str) or not re.fullmatch(r"[a-f0-9]{64}", checksum):
            raise FirmwareBootError(f"invalid file record: {path}")
    for entry in state["entries"]:
        if not re.fullmatch(r"[0-9A-F]{4}", entry["number"]) or not isinstance(entry["identity"], str):
            raise FirmwareBootError(f"invalid firmware record: {path}")
        for relative in entry["files"]:
            owned_path(configuration, relative)
    return state


def save_state(configuration: Configuration, state: dict) -> None:
    path = state_path(configuration)
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent, delete=False) as stream:
        temporary = Path(stream.name)
        try:
            json.dump(state, stream, indent=2)
            stream.write("\n")
            stream.flush()
            os.fsync(stream.fileno())
            os.replace(temporary, path)
        finally:
            temporary.unlink(missing_ok=True)


def register_entry(configuration: Configuration, state: dict, request: dict, files: list[str]) -> str:
    before = firmware_state()
    for record in state["entries"]:
        if record["request"] == request and before.entries.get(record["number"]) == record["identity"]:
            record["files"] = files
            save_state(configuration, state)
            if record["number"] not in before.order:
                run(["efibootmgr", "--bootorder", ",".join(before.order + [record["number"]])])
            return record["number"]
    command = ["efibootmgr", "--create-only", "--disk", configuration.efi_disk,
        "--part", str(configuration.efi_partition), "--label", request["label"], "--loader", request["loader"]]
    if request["options"]:
        command.extend(["--unicode", request["options"]])
    run(command)
    after = firmware_state()
    created = set(after.entries) - set(before.entries)
    if len(created) != 1:
        raise FirmwareBootError("cannot identify the newly created firmware entry")
    number = created.pop()
    identity = after.entries[number]
    if (not (identity.startswith(request["label"] + "\t") or identity.startswith(request["label"] + " "))
        or ("File(" + request["loader"] + ")").casefold() not in identity.casefold()):
        raise FirmwareBootError("the newly created firmware entry does not match the requested loader")
    state["entries"].append({"number": number, "identity": after.entries[number], "request": request, "files": files})
    save_state(configuration, state)
    # Keep the previous default and order; replace an owned entry in its existing slot.
    replaced = {record["number"] for record in state["entries"] if record["request"]["label"] == request["label"] and record["request"]["disk"] == request["disk"]
        and record["request"]["partition"] == request["partition"] and record["number"] != number
        and before.entries.get(record["number"]) == record["identity"]}
    order: list[str] = []
    for previous in before.order:
        selected = number if previous in replaced else previous
        if selected not in order:
            order.append(selected)
    if number not in order:
        order.append(number)
    run(["efibootmgr", "--bootorder", ",".join(order)])
    return number


def prune_owned(configuration: Configuration, state: dict, desired: set[str], active_entries: set[str] | None) -> None:
    protected: set[str] = set()
    retained = []
    for record in state["entries"]:
        if active_entries is None or record["number"] in active_entries:
            retained.append(record)
            protected.update(record["files"])
            continue
        observed = firmware_state()
        number = record["number"]
        if number in {observed.current, observed.next_boot} or (number in observed.entries and observed.entries[number] != record["identity"]):
            retained.append(record)
            protected.update(record["files"])
        elif observed.entries.get(number) == record["identity"]:
            run(["efibootmgr", "--bootnum", number, "--delete-bootnum"])
    state["entries"] = retained
    save_state(configuration, state)
    for relative, checksum in list(state["files"].items()):
        if relative in desired or relative in protected:
            continue
        path = owned_path(configuration, relative)
        if path.is_file() and not path.is_symlink() and digest(path) == checksum:
            path.unlink()
            parent = path.parent
            while parent != configuration.esp_path and parent.is_dir() and not any(parent.iterdir()):
                parent.rmdir()
                parent = parent.parent
        del state["files"][relative]
    save_state(configuration, state)


def atomic_copy(source: Path, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(dir=destination.parent, prefix=destination.name + ".", delete=False) as temporary:
        temporary_path = Path(temporary.name)
    try:
        shutil.copy2(source, temporary_path)
        os.replace(temporary_path, destination)
    finally:
        temporary_path.unlink(missing_ok=True)


def uki_output_path(configuration: Configuration, kernel: Kernel) -> Path:
    return configuration.esp_path / "EFI/Linux" / f"catos-{kernel.name}.efi"


def efistub_directory(configuration: Configuration, kernel: Kernel) -> Path:
    return configuration.esp_path / "EFI/CatOS" / kernel.name


def firmware_path(configuration: Configuration, path: Path) -> str:
    return loader_path(Path(configuration.firmware_directory.lstrip("/")) / path.relative_to(configuration.esp_path))


def machine_secure_boot_managed(
    configuration: Configuration,
    key_path: Path = SECURE_BOOT_KEY_PATH,
) -> bool:
    return configuration.method == "uki" and key_path.is_file()


def build_uki(configuration: Configuration, kernel: Kernel, initramfs: Path, cmdline: str) -> Path:
    output = uki_output_path(configuration, kernel)
    output.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix=".uki-", dir=output.parent) as directory:
        temporary = Path(directory) / "image.efi"
        run(["ukify", "build", "--linux", str(kernel.image), "--initrd", str(initramfs),
            "--cmdline", cmdline, "--os-release", "@/etc/os-release", "--uname", kernel.version, "--output", str(temporary)])
        if not temporary.is_file() or temporary.stat().st_size == 0:
            raise FirmwareBootError(f"UKI generation produced no image: {kernel.package}")
        os.replace(temporary, output)
    return output


def build_efistub(configuration: Configuration, kernel: Kernel, initramfs: Path) -> tuple[Path, Path]:
    parent = efistub_directory(configuration, kernel)
    parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix=".build-", dir=parent) as directory:
        staging = Path(directory)
        shutil.copyfile(kernel.image, staging / "vmlinuz.efi")
        shutil.copyfile(initramfs, staging / "initramfs.img")
        # The pair has an immutable destination. A failed build leaves the old pair intact.
        checksum = hashlib.sha256((digest(staging / "vmlinuz.efi") + digest(staging / "initramfs.img")).encode()).hexdigest()
        destination = parent / checksum
        if destination.exists():
            if digest(destination / "vmlinuz.efi") != digest(staging / "vmlinuz.efi") or digest(destination / "initramfs.img") != digest(staging / "initramfs.img"):
                raise FirmwareBootError(f"existing kernel deployment was modified: {destination}")
        else:
            os.replace(staging, destination)
    return destination / "vmlinuz.efi", destination / "initramfs.img"


def selected_default(kernels: list[Kernel], package: str) -> Kernel:
    matches = [kernel for kernel in kernels if kernel.name == package]
    if len(matches) != 1:
        raise FirmwareBootError(f"select an installed default kernel: {package}")
    return matches[0]


def verify(configuration: Configuration, kernels: list[Kernel], *, require_entries: bool = True) -> None:
    state = load_state(configuration)
    for relative, checksum in state["files"].items():
        path = owned_path(configuration, relative)
        if (not path.is_file() or path.stat().st_size == 0
            or (digest(path) != checksum and not machine_secure_boot_managed(configuration))):
            raise FirmwareBootError(f"deployed boot file is missing or changed: {path}")
    for kernel in kernels:
        if kernel.name not in state.get("kernels", {}):
            raise FirmwareBootError(f"kernel has not been deployed: {kernel.package}")
    if require_entries:
        observed = firmware_state()
        for kernel in kernels:
            files = state["kernels"][kernel.name]
            if not any(record["files"] == files and observed.entries.get(record["number"]) == record["identity"] for record in state["entries"]):
                raise FirmwareBootError(f"firmware entry is missing: {kernel.package}")


def update(configuration: Configuration, *, no_variables: bool = False) -> None:
    if not CMDLINE_PATH.is_file():
        raise FirmwareBootError(f"kernel command line is missing: {CMDLINE_PATH}")
    kernels = discover_kernels()
    cmdline = CMDLINE_PATH.read_text(encoding="utf-8").strip()
    if configuration.method == "efistub" and configuration.efi_disk and not no_variables and " " in configuration.firmware_directory:
        raise FirmwareBootError("EFI-stub initrd firmware paths cannot contain spaces")
    configuration.esp_path.mkdir(parents=True, exist_ok=True)
    state = load_state(configuration)
    desired: set[str] = set()
    pending: list[tuple[dict, list[str]]] = []
    deployed: dict[str, list[str]] = {}
    register = bool(configuration.efi_disk and configuration.efi_partition > 0) and not no_variables and not machine_secure_boot_managed(configuration)
    outputs: dict[str, Path] = {}
    state_path(configuration).parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix=".build-", dir=state_path(configuration).parent) as directory:
        for kernel in kernels:
            work = Path(directory) / kernel.version
            work.mkdir()
            initramfs = work / "initramfs.img"
            run(["mkinitcpio", "--kernel", kernel.version, "--generate", str(initramfs), "--nopost"])
            if not initramfs.is_file() or initramfs.stat().st_size == 0:
                raise FirmwareBootError(f"initramfs generation produced no image: {kernel.version}")
            if configuration.method == "uki":
                output = build_uki(configuration, kernel, initramfs, cmdline)
                paths = [output]
                options = ""
            else:
                output, initrd = build_efistub(configuration, kernel, initramfs)
                paths = [output, initrd]
                options = f"initrd={firmware_path(configuration, initrd)} {cmdline}"
            files = [str(path.relative_to(configuration.esp_path)) for path in paths]
            outputs[kernel.name] = output
            deployed[kernel.name] = files
            desired.update(files)
            state["files"].update({relative: digest(path) for relative, path in zip(files, paths)})
            save_state(configuration, state)
            pending.append(({"label": f"{configuration.label_prefix} {kernel.name}",
                "loader": firmware_path(configuration, output), "options": options,
                "disk": configuration.efi_disk, "partition": configuration.efi_partition}, files))
    if configuration.method == "uki" and configuration.install_fallback and not machine_secure_boot_managed(configuration):
        default = selected_default(kernels, configuration.default_kernel)
        filename = "BOOTAA64.EFI" if os.uname().machine in {"aarch64", "arm64"} else "BOOTX64.EFI"
        destination = configuration.esp_path / "EFI/BOOT" / filename
        atomic_copy(outputs[default.name], destination)
        relative = str(destination.relative_to(configuration.esp_path))
        desired.add(relative)
        state["files"][relative] = digest(destination)
        save_state(configuration, state)
    state["kernels"] = deployed
    save_state(configuration, state)
    active = {register_entry(configuration, state, request, files) for request, files in pending} if register else None
    prune_owned(configuration, state, desired, active)
    verify(configuration, kernels, require_entries=register)


def save_configuration(path: Path, configuration: Configuration) -> None:
    parser = configparser.ConfigParser(interpolation=None)
    parser.read(path)
    if not parser.has_section("boot"):
        parser.add_section("boot")
    parser["boot"].update({
        "method": configuration.method,
        "esp_path": str(configuration.esp_path),
        "efi_disk": configuration.efi_disk,
        "efi_partition": str(configuration.efi_partition),
        "label_prefix": configuration.label_prefix,
        "default_kernel": configuration.default_kernel,
        "install_fallback": "yes" if configuration.install_fallback else "no",
        "firmware_directory": configuration.firmware_directory,
    })
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent, delete=False) as stream:
        temporary = Path(stream.name)
        try:
            parser.write(stream)
            stream.flush()
            os.fsync(stream.fileno())
            os.replace(temporary, path)
        finally:
            temporary.unlink(missing_ok=True)


def describe(configuration: Configuration) -> dict:
    return {
        "loader": configuration.method, "esp": str(configuration.esp_path),
        "efi_disk": configuration.efi_disk, "efi_partition": configuration.efi_partition,
        "label_prefix": configuration.label_prefix, "default_kernel": configuration.default_kernel,
        "fallback": configuration.install_fallback, "firmware_directory": configuration.firmware_directory,
        "cmdline": CMDLINE_PATH.read_text(encoding="utf-8").strip() if CMDLINE_PATH.is_file() else "",
    }


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Maintain CatOS direct UEFI boot artifacts")
    parser.add_argument("--config", type=Path, default=CONFIG_PATH, help="configuration file")
    parser.add_argument("--force", action="store_true", help="run even during Calamares installation")
    parser.add_argument("--hook", action="store_true", help="consume pacman hook targets from stdin")
    parser.add_argument("--no-variables", action="store_true", help="deploy files without reading or writing firmware variables")
    action = parser.add_mutually_exclusive_group()
    action.add_argument("--verify", action="store_true", help="verify without rebuilding")
    action.add_argument("--describe", action="store_true", help="print configuration as JSON without making changes")
    action.add_argument("--configure-only", action="store_true", help="save settings without building")
    parser.add_argument("--save-config", action="store_true", help="save settings after a successful deployment")
    parser.add_argument("--method", choices=("disabled", "uki", "efistub"))
    parser.add_argument("--output-root", type=Path)
    parser.add_argument("--efi-disk")
    parser.add_argument("--efi-partition", type=int)
    parser.add_argument("--firmware-directory")
    parser.add_argument("--label-prefix")
    parser.add_argument("--default-kernel")
    parser.add_argument("--install-fallback", choices=("yes", "no"))
    parser.add_argument("--cmdline", help="save kernel parameters with --configure-only")
    arguments = parser.parse_args()
    if arguments.cmdline is not None and not arguments.configure_only:
        parser.error("--cmdline requires --configure-only")
    return arguments


def main() -> int:
    arguments = parse_arguments()
    if arguments.hook:
        sys.stdin.read()
    configuration = load_configuration(arguments.config)
    overrides = {key: value for key, value in {
        "method": arguments.method, "esp_path": arguments.output_root,
        "efi_disk": arguments.efi_disk, "efi_partition": arguments.efi_partition,
        "label_prefix": arguments.label_prefix, "default_kernel": arguments.default_kernel,
        "firmware_directory": arguments.firmware_directory,
        "install_fallback": None if arguments.install_fallback is None else arguments.install_fallback == "yes",
    }.items() if value is not None}
    configuration = replace(configuration, **overrides)
    if not configuration.esp_path.is_absolute():
        raise FirmwareBootError("output root must be an absolute path")
    if bool(configuration.efi_disk) != (configuration.efi_partition > 0) or configuration.efi_partition < 0:
        raise FirmwareBootError("provide both a disk and a positive partition number")
    firmware_directory = Path(configuration.firmware_directory)
    if not firmware_directory.is_absolute() or ".." in firmware_directory.parts or "\\" in configuration.firmware_directory:
        raise FirmwareBootError("firmware directory must be an absolute path without parent traversal")
    for value in (configuration.firmware_directory, str(configuration.esp_path), configuration.efi_disk, configuration.label_prefix, configuration.default_kernel, arguments.cmdline or ""):
        if any(ord(character) < 32 or ord(character) == 127 for character in value):
            raise FirmwareBootError("settings must contain a single line")
    if arguments.describe:
        print(json.dumps(describe(configuration)))
        return 0
    if INSTALL_MARKER.exists() and not arguments.force:
        if arguments.hook:
            print("catos-firmware-boot: deferred while bootloadu is preparing the installed system")
            return 0
        raise FirmwareBootError("installation is currently in progress")
    LOCK_PATH.parent.mkdir(parents=True, exist_ok=True)
    with LOCK_PATH.open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if arguments.configure_only:
            if arguments.cmdline is not None:
                CMDLINE_PATH.parent.mkdir(parents=True, exist_ok=True)
                with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=CMDLINE_PATH.parent, delete=False) as stream:
                    temporary = Path(stream.name)
                    try:
                        stream.write(arguments.cmdline + "\n")
                        stream.flush()
                        os.fsync(stream.fileno())
                        os.replace(temporary, CMDLINE_PATH)
                    finally:
                        temporary.unlink(missing_ok=True)
            save_configuration(arguments.config, configuration)
        elif configuration.method != "disabled":
            if arguments.verify:
                verify(configuration, discover_kernels(), require_entries=bool(configuration.efi_disk) and not arguments.no_variables and not machine_secure_boot_managed(configuration))
            else:
                update(configuration, no_variables=arguments.no_variables)
                if arguments.save_config:
                    save_configuration(arguments.config, configuration)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (FirmwareBootError, OSError, ValueError, configparser.Error, subprocess.CalledProcessError) as error:
        print(f"catos-firmware-boot: {error}", file=sys.stderr)
        raise SystemExit(1)
