# SPDX-License-Identifier: LGPL-2.1-or-later

import argparse
import contextlib
import dataclasses
import logging
import os
import platform
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Optional, cast

import mkosi.resources
from mkosi.config import DocFormat, OutputFormat
from mkosi.documentation import show_docs
from mkosi.log import log_notice, log_setup
from mkosi.run import find_binary, run, uncaught_exception_handler
from mkosi.sandbox import __version__, umask
from mkosi.tree import copy_tree
from mkosi.util import PathString, mandatory_variable, resource_path


@dataclasses.dataclass(frozen=True)
class KernelInstallContext:
    command: str
    kernel_version: str
    entry_dir: Path
    kernel_image: Path
    initrds: list[Path]
    staging_area: Path
    layout: str
    image_type: str
    initrd_generator: Optional[str]
    uki_generator: Optional[str]
    verbose: bool

    @staticmethod
    def parse(*, name: str, description: str) -> "KernelInstallContext":
        parser = argparse.ArgumentParser(
            description=description,
            allow_abbrev=False,
            usage=f"{name} COMMAND KERNEL_VERSION ENTRY_DIR KERNEL_IMAGE…",
        )

        parser.add_argument(
            "command",
            metavar="COMMAND",
            help="The action to perform. Only 'add' is supported.",
        )
        parser.add_argument(
            "kernel_version",
            metavar="KERNEL_VERSION",
            help="Kernel version string",
        )
        parser.add_argument(
            "entry_dir",
            metavar="ENTRY_DIR",
            type=Path,
            nargs="?",
            help="Type#1 entry directory (ignored)",
        )
        parser.add_argument(
            "kernel_image",
            metavar="KERNEL_IMAGE",
            type=Path,
            nargs="?",
            help="Kernel image",
        )
        parser.add_argument(
            "initrds",
            metavar="INITRD…",
            type=Path,
            nargs="*",
            help="Initrd files",
        )
        parser.add_argument(
            "--version",
            action="version",
            version=f"mkosi {__version__}",
        )

        args = parser.parse_args()

        return KernelInstallContext(
            command=args.command,
            kernel_version=args.kernel_version,
            entry_dir=args.entry_dir,
            kernel_image=args.kernel_image,
            initrds=args.initrds,
            staging_area=Path(mandatory_variable("KERNEL_INSTALL_STAGING_AREA")),
            layout=mandatory_variable("KERNEL_INSTALL_LAYOUT"),
            image_type=mandatory_variable("KERNEL_INSTALL_IMAGE_TYPE"),
            initrd_generator=os.getenv("KERNEL_INSTALL_INITRD_GENERATOR"),
            uki_generator=os.getenv("KERNEL_INSTALL_UKI_GENERATOR"),
            verbose=int(os.getenv("KERNEL_INSTALL_VERBOSE", 0)) > 0,
        )


def process_crypttab(staging_dir: str) -> list[str]:
    cmdline = []

    # Generate crypttab with all the x-initrd.attach entries
    if Path("/etc/crypttab").exists():
        try:
            crypttab = [
                line
                for line in Path("/etc/crypttab").read_text().splitlines()
                if (
                    len(entry := line.split()) >= 4
                    and not entry[0].startswith("#")
                    and "x-initrd.attach" in entry[3]
                )
            ]
            if crypttab:
                with (Path(staging_dir) / "crypttab").open("w") as f:
                    f.write("# Automatically generated by mkosi-initrd\n")
                    f.write("\n".join(crypttab))
                cmdline += ["--extra-tree", f"{staging_dir}/crypttab:/etc/crypttab"]
        except PermissionError:
            logging.warning("Permission denied to access /etc/crypttab, the initrd may be unbootable")

    return cmdline


def initrd_finalize(staging_dir: str, output: str, output_dir: str) -> None:
    if output_dir:
        with umask(~0o700) if os.getuid() == 0 else cast(umask, contextlib.nullcontext()):
            Path(output_dir).mkdir(parents=True, exist_ok=True)
    else:
        output_dir = str(Path.cwd())

    log_notice(f"Copying {staging_dir}/{output} to {output_dir}/{output}")
    # mkosi symlinks the expected output image, so dereference it
    copy_tree(
        Path(f"{staging_dir}/{output}").resolve(),
        Path(f"{output_dir}/{output}"),
    )


def initrd_common_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--kernel-version",
        metavar="KERNEL_VERSION",
        help="Kernel version string",
        default=platform.uname().release,
    )
    parser.add_argument(
        "-O",
        "--output-dir",
        metavar="DIR",
        help="Output directory",
        default="",
    )
    parser.add_argument(
        "--debug",
        help="Turn on debugging output",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--debug-shell",
        help="Spawn debug shell if a sandboxed command fails",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "-D",
        "--show-documentation",
        help="Show the man page",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"mkosi {__version__}",
    )


def include_system_config(name: str) -> list[str]:
    cmdline = []

    for d in ("/usr/lib", "/usr/local/lib", "/run", "/etc"):
        p = Path(d) / name
        if p.exists():
            cmdline += ["--include", os.fspath(p)]

    return cmdline


@uncaught_exception_handler()
def main() -> None:
    log_setup()

    parser = argparse.ArgumentParser(
        prog="mkosi-initrd",
        description="Build initrds or unified kernel images for the current system using mkosi",
        allow_abbrev=False,
        usage="mkosi-initrd [options...]",
    )
    parser.add_argument(
        "-o",
        "--output",
        metavar="NAME",
        help="Output name",
        default="initrd",
    )
    parser.add_argument(
        "--kernel-image",
        metavar="KERNEL_IMAGE",
        help="Kernel image",
        type=Path,
    )
    parser.add_argument(
        "-t",
        "--format",
        choices=[str(OutputFormat.cpio), str(OutputFormat.uki), str(OutputFormat.directory)],
        help="Output format (CPIO archive, UKI or local directory)",
        default="cpio",
    )

    initrd_common_args(parser)

    args = parser.parse_args()

    if args.show_documentation:
        with resource_path(mkosi.resources) as r:
            show_docs("mkosi-initrd", DocFormat.all(), resources=r)
        return

    with (
        tempfile.TemporaryDirectory() as staging_dir,
        tempfile.TemporaryDirectory() as sandbox_tree,
    ):
        cmdline: list[PathString] = [
            "mkosi",
            "--force",
            "--directory", "",
            "--format", args.format,
            "--output", args.output,
            "--output-directory", staging_dir,
            "--extra-tree", f"/usr/lib/modules/{args.kernel_version}:/usr/lib/modules/{args.kernel_version}",
            "--extra-tree=/usr/lib/firmware:/usr/lib/firmware",
            "--remove-files=/usr/lib/firmware/*-ucode",
            "--kernel-modules-exclude=.*",
            "--kernel-modules-include=host",
            "--build-sources", "",
            "--include=mkosi-initrd",
        ]  # fmt: skip

        if args.kernel_image:
            cmdline += [
                "--extra-tree", f"{args.kernel_image}:/usr/lib/modules/{args.kernel_version}/vmlinuz",
            ]  # fmt: skip

        if args.debug:
            cmdline += ["--debug"]
        if args.debug_shell:
            cmdline += ["--debug-shell"]

        if os.getuid() == 0:
            cmdline += [
                "--workspace-dir=/var/tmp",
                "--package-cache-dir=/var",
                "--cache-only=metadata",
            ]
            if args.format != OutputFormat.directory.value:
                cmdline += ["--output-mode=600"]

        cmdline += include_system_config("mkosi-initrd")

        # Make sure we don't use any of mkosi's default repositories.
        for p in (
            "yum.repos.d/mkosi.repo",
            "apt/sources.list.d/mkosi.sources",
            "zypp/repos.d/mkosi.repo",
            "pacman.conf",
        ):
            (Path(sandbox_tree) / "etc" / p).parent.mkdir(parents=True, exist_ok=True)
            (Path(sandbox_tree) / "etc" / p).touch()

        # Copy in the host's package manager configuration.
        for p in (
            "dnf",
            "yum.repos.d/",
            "apt",
            "zypp",
            "pacman.conf",
            "pacman.d/",
        ):
            if not (Path("/etc") / p).exists():
                continue

            (Path(sandbox_tree) / "etc" / p).parent.mkdir(parents=True, exist_ok=True)
            if (Path("/etc") / p).resolve().is_file():
                shutil.copy2(Path("/etc") / p, Path(sandbox_tree) / "etc" / p)
            else:
                shutil.copytree(
                    Path("/etc") / p,
                    Path(sandbox_tree) / "etc" / p,
                    ignore=shutil.ignore_patterns("gnupg"),
                    dirs_exist_ok=True,
                )

        cmdline += ["--sandbox-tree", sandbox_tree]

        cmdline += process_crypttab(staging_dir)

        if Path("/etc/kernel/cmdline").exists():
            cmdline += ["--kernel-command-line", Path("/etc/kernel/cmdline").read_text()]

        # Resolve dnf binary to determine which version the host uses by default
        # (to avoid preferring dnf5 if the host uses dnf4)
        # as there's a much bigger chance that it has a populated dnf cache directory.
        run(
            cmdline,
            stdin=sys.stdin,
            stdout=sys.stdout,
            env={"MKOSI_DNF": dnf.resolve().name} if (dnf := find_binary("dnf")) else {},
        )

        initrd_finalize(staging_dir, args.output, args.output_dir)


if __name__ == "__main__":
    main()
