# Copyright The Monogon Project Authors.
# SPDX-License-Identifier: Apache-2.0

"""
Rules for building Linux kernel images.

This currently performs the build in a fully hermetic manner, using
make/gcc/... from the toolchain bundle, and is only slightly better than a genrule. This
should be replaced by a hermetic build that just uses cc_library targets.
"""

load("@rules_cc//cc:find_cc_toolchain.bzl", "CC_TOOLCHAIN_ATTRS", "find_cpp_toolchain", "use_cc_toolchain")
load("@rules_cc//cc/common:cc_common.bzl", "cc_common")
load("@rules_cc//cc/common:cc_info.bzl", "CcInfo")
load("//build/utils:detect_root.bzl", "detect_root", "detect_roots")
load("//build/utils:foreign_build.bzl", "generate_foreign_build_env", "merge_env")
load("//build/utils:target_info.bzl", "TargetInfo")

def _linux_image_impl_resources(_os, _ninputs):
    """
    Configures linux build resources.

    See `resource_set` documentation in builtins.actions Bazel docs.
    """

    # 16 threads seems about right - this fits well in both our build machines and
    # development machines.
    cpu = 16

    # In MB. Picked based on observing build in htop.
    mb_per_cpu = 256
    return {
        "cpu": cpu,
        "memory": cpu * mb_per_cpu,
        "local_test": 0,
    }

TOOLCHAINS = [
    "//build/toolchain/toolchain-bundle:make_toolchain",
    "//build/toolchain/toolchain-bundle:flex_toolchain",
    "//build/toolchain/toolchain-bundle:bison_toolchain",
    "//build/toolchain/toolchain-bundle:m4_toolchain",
    "//build/toolchain/toolchain-bundle:busybox_toolchain",
    "//build/toolchain/toolchain-bundle:bc_toolchain",
    "//build/toolchain/toolchain-bundle:diff_toolchain",
    "//build/toolchain/toolchain-bundle:perl_toolchain",
    "//build/toolchain/toolchain-bundle:lz4_toolchain",
]

def _linux_image_impl(ctx):
    # Tuple containing information about how to build and access the resulting
    # image.
    # The first element (target) is the make target to build, the second
    # (image_source) is the resulting file to be copied and the last
    # (image_name) is the name of the image that will be generated by this
    # rule.
    (target, image_source, image_name) = {
        "vmlinux": ("vmlinux modules", "vmlinux", "vmlinux"),
        "Image": ("all modules", "arch/" + ctx.attr._target_arch[TargetInfo].value + "/boot/" + ctx.attr._image_name[TargetInfo].value, "Image"),
    }[ctx.attr.image_format]

    ssl_src, ssl_gen = detect_roots(ctx.attr._ssl[CcInfo].compilation_context.direct_public_headers)
    crypto_src, crypto_gen = detect_roots(ctx.attr._crypto[CcInfo].compilation_context.direct_public_headers)
    extra_env = {
        "HOSTLDFLAGS": " -L ".join(
            [
                "",  # First element empty, for force a the join prefix
                detect_root(ctx.attr._zstd.files.to_list()).rsplit("/", 1)[0],
                detect_root(ctx.attr._zlib.files.to_list()).rsplit("/", 1)[0],
                detect_root(ctx.attr._libelf.files.to_list()).rsplit("/", 1)[0],
                detect_root(ctx.attr._ssl.files.to_list()).rsplit("/", 1)[0],
                detect_root(ctx.attr._crypto.files.to_list()).rsplit("/", 1)[0],
            ],
        ),
        "HOSTCFLAGS": " -I ".join(
            [
                "",  # First element empty, for force a the join prefix
                detect_root(ctx.attr._libelf[CcInfo].compilation_context.direct_public_headers),
                ssl_src + "/../",
                ssl_gen + "/../include/",
                crypto_src + "/../",
                crypto_gen + "/../include/",
            ],
        ),
    }

    inputs = depset(
        ctx.files.kernel_config +
        ctx.files.kernel_src +
        ctx.files._libelf +
        ctx.attr._libelf[CcInfo].compilation_context.direct_public_headers +
        ctx.files._zlib +
        ctx.files._ssl +
        ctx.attr._ssl[CcInfo].compilation_context.direct_public_headers +
        ctx.files._crypto +
        ctx.attr._crypto[CcInfo].compilation_context.direct_public_headers,
    )

    # Setup the environment for the foreign build.
    toolchain_env, toolchain_inputs, toolchain_cmd = generate_foreign_build_env(
        ctx = ctx,
        target_toolchain = find_cpp_toolchain(ctx),
        exec_toolchain = ctx.attr._exec_toolchain[cc_common.CcToolchainInfo],
        toolchain_bundle_tools = TOOLCHAINS,
    )

    image = ctx.actions.declare_file(image_name)
    modinfo = ctx.actions.declare_file("modules.builtin.modinfo")
    modules = ctx.actions.declare_directory("modules")
    ctx.actions.run_shell(
        outputs = [image, modinfo, modules],
        inputs = depset(transitive = [inputs, toolchain_inputs]),
        resource_set = _linux_image_impl_resources,
        env = merge_env(toolchain_env, extra_env),
        progress_message = "Building Linux Kernel: {}".format(ctx.label.name),
        mnemonic = "BuildLinux",
        command = toolchain_cmd + """
            export BISON_PKGDATADIR=$(realpath $(dirname $BISON))/../share/bison
            builddir=$(mktemp -d)

            mkdir {kernel_src}/.bin
            cp {kconfig} $builddir/.config
            (
                cd {kernel_src} &&
                make -j 16 \
                \
                CC="$CC" CXX="$CXX" LD="$LD" AR="$AR" NM="$NM" STRIP="$STRIP" \
                OBJCOPY="$OBJCOPY" OBJDUMP="$OBJDUMP" READELF="$READELF" \
                CFLAGS="$CFLAGS" LDFLAGS="$LDFLAGS" \
                \
                HOSTCC="$HOSTCC" HOSTCXX="$HOSTCXX" HOSTLD="$HOSTLD" \
                HOSTAR="$HOSTAR" HOSTNM="$HOSTNM" HOSTSTRIP="$HOSTSTRIP" \
                HOSTOBJCOPY="$HOSTOBJCOPY" HOSTOBJDUMP="$HOSTOBJDUMP" \
                HOSTREADELF="$HOSTREADELF" HOSTCFLAGS="$HOSTCFLAGS" \
                HOSTLDFLAGS="$HOSTLDFLAGS" \
                \
                KBUILD_OUTPUT="$builddir" \
                ARCH="{target_arch}" \
                olddefconfig {target}
            ) > /dev/null

            cp "$builddir"/{image_source} {image}
            cp "$builddir"/modules.builtin.modinfo {modinfo}
            # Not using modules_install as it tries to run depmod and friends
            for f in $(find "$builddir" -name '*.ko' -type f -printf "%P\n" ); do
                install -D "$builddir/$f" "{modules}/$f"
            done
            rm -Rf "$builddir"
            """.format(
            kconfig = ctx.file.kernel_config.path,
            target = target,
            image_source = image_source,
            kernel_src = detect_root(ctx.attr.kernel_src.files.to_list()),
            image = image.path,
            modinfo = modinfo.path,
            modules = modules.path,
            target_arch = ctx.attr._target_arch[TargetInfo].value,
        ),
        use_default_shell_env = True,
    )

    return [
        DefaultInfo(
            files = depset([image]),
            runfiles = ctx.runfiles(files = [image]),
        ),
        OutputGroupInfo(
            modinfo = depset([modinfo]),
            modules = depset([modules]),
        ),
    ]

linux_image = rule(
    doc = """
        Build Linux kernel image hermetically in a given format.
    """,
    implementation = _linux_image_impl,
    attrs = {
        "kernel_config": attr.label(
            doc = """
                Linux kernel configuration file to build this kernel image with.
            """,
            allow_single_file = True,
        ),
        "kernel_src": attr.label(
            doc = """
                Filegroup containing Linux kernel sources.
            """,
        ),
        "image_format": attr.string(
            doc = """
                Format of generated Linux image, one of 'vmlinux' or 'Image',
            """,
            values = [
                "vmlinux",
                "Image",
            ],
            default = "Image",
        ),
        "_libelf": attr.label(
            default = "//third_party:libelf_elf",
            cfg = "exec",
        ),
        "_zstd": attr.label(
            default = "//third_party:zstd_zstd",
            cfg = "exec",
        ),
        "_zlib": attr.label(
            default = "//third_party:zlib_z",
            cfg = "exec",
        ),
        "_ssl": attr.label(
            default = "//third_party:openssl_ssl",
            cfg = "exec",
        ),
        "_crypto": attr.label(
            default = "//third_party:openssl_crypto",
            cfg = "exec",
        ),
        "_exec_toolchain": attr.label(
            default = "@rules_cc//cc:current_cc_toolchain",
            cfg = "exec",
        ),
        "_image_name": attr.label(
            default = "//third_party/linux:image_name",
        ),
        "_target_arch": attr.label(
            default = "//third_party/linux:target_arch",
        ),
    } | CC_TOOLCHAIN_ATTRS,
    fragments = ["cpp"],
    toolchains = TOOLCHAINS + use_cc_toolchain(),
)