#!/usr/bin/env python
# SPDX-License-Identifier: GPL-2.0-only
# This scripts use bitbake api, so it has to be under the GPL-2.0 license

"""
    devtool wrapper for the coreos-device script
"""

import subprocess
import textwrap
import os
import sys
import subprocess

from devtool import setup_tinfoil, parse_recipe
from dataclasses import dataclass
from pathlib import Path
from shlex import quote
from typing import List

def register_commands(subparsers, context):
    """Register devtool subcommands from this plugin"""
    subparsers.add_subparser_group("efibootguard", "Developers tools for efibootguard")
    cmd_uki = subparsers.add_parser(
        "generate-uki",
        help="Generate signed UKI kernel images",
        description="""
        Helper to generate Unified Kernel Image using bg_generate_uki. The UKI is
        signed and contain a EFI bootstraper, a kernel image, the kernel command
        line and one or multiple device tree
        
        Before calling this function, theses recipes should be build by bitbake:
        `virtual/kernel efibootguard efibootguard-native`
        """,
        group="efibootguard",
    )
    cmd_uki.add_argument("imagename", help="Image recipe used to get the configuration")
    cmd_uki.add_argument(
        "--cmdline-append",
        default="",
        help="string to append to the kernel command line",
    )

    cmd_uki.add_argument(
        "--cmdline",
        default=None,
        help="""
        Replace the common part of the cmdline. root= parameter will be added
        automatically if --cmdline-no-root is not used
        """,
    )

    cmd_uki.add_argument(
        "--cmdline-no-root",
        default=False,
        action="store_true",
        help="Don't automatically append root= if --cmdline is used",
    )

    cmd_uki.add_argument(
        "--dtb",
        default=None,
        help="Embedded the given dtb instead of the default set",
    )

    cmd_uki.add_argument(
        "--kernel",
        default=None,
        help="Use a custom kernel binary",
    )

    cmd_uki.set_defaults(
        func=efibootguard_generate_uki, fixed_setup=context.fixed_setup
    )


def run(cmd):
    return subprocess.run(cmd).returncode


class InvalidImage(Exception):
    pass

@dataclass
class UKIGeneratorArgs:
    """
    Class used to generate and run commands needed to build and sign
    the UKI image
    """

    kernel: str
    output: str
    cmdline: str
    dtb: str
    stub: list
    root: str
    build_binary: str
    keydir: str

    def process_args(self, args):
        """
        This apply args passed to the program to override or modify the
        parameters that we received from the bitbake context
        """
        if args.dtb:
            self.dtb = [args.dtb]

        if args.kernel:
            self.kernel = args.kernel

        if args.cmdline:
            self.cmdline = args.cmdline
            if not args.cmdline_no_root:
                self.cmdline = f"{self.cmdline} root={self.root}"

        if args.cmdline_append:
            self.cmdline = f"{self.cmdline} {args.cmdline_append}"

    def __str__(self):
        return (
            f"stub: {self.stub}\n"
            f"kernel: {self.kernel}\n"
            f"cmdline: {self.cmdline}\n"
            f"dtb: {self.dtb}\n"
            f"output: {self.output}\n"
            f"keydir: {self.keydir}\n"
        )

    def get_buid_cmd(self) -> List[str]:
        """
        Return the command needed to build the UKI as an array
        """
        args = [
            self.build_binary,
            self.stub,
            self.kernel,
            self.output,
            "--cmdline",
            self.cmdline,
        ]
        for dtb in self.dtb:
            args += ["--dtb", dtb]
        return args

    def get_sign_cmd(self) -> List[str]:
        """
        Return the command needed to sign the UKI in place as an array
        """
        return [
            "sbsign",
            "--key",
            os.path.join(self.keydir, "db.key"),
            "--cert",
            os.path.join(self.keydir, "db.crt"),
            self.output,
            "--output",
            self.output,
        ]

    def print_cmd(self, cmd):
        """
        Pretty print the command in the terminal with indentation and line
        splitting using standard \ keyword
        """
        first = True
        previous_was_an_option = False
        for index, arg in enumerate(cmd):
            # Quote the string as needed, so that copy past from the printed
            # command works.
            arg = quote(arg)
            end = " \\\n"
            if index == len(cmd) - 1:
                # For the last
                end = " \n"

            if first:
                # only print the command name not the full path
                print(os.path.basename(arg), end=end)
                first = False
                continue

            if previous_was_an_option:
                print(arg, end=end)
                previous_was_an_option = False
                continue

            if arg.startswith("-"):
                previous_was_an_option = True
                printi(arg, 1, end=" ")
            else:
                printi(arg, 1, end=end)

    def build_and_sign(self) -> int:
        """
        Run the sign and build command and return the sum of both commands
        exit status
        """
        build_cmd = self.get_buid_cmd()
        sign_cmd = self.get_sign_cmd()

        self.print_cmd(build_cmd)
        r = run(build_cmd)

        self.print_cmd(sign_cmd)
        r += run(sign_cmd)

        return r


def printi(txt, indent=1, *, iprefix="  ", end="\n"):
    """
    Helper function to print multiples lines with some indentation
    """
    print(textwrap.indent(txt, iprefix * indent), end=end)


def dtb_to_list(dtb, deploy_dir) -> List[str]:
    """
    Convert the string array from bitbake into a python list. Only the basename
    of the dtb is taken and DEPLOY_DIR is appended
    """
    if dtb:
        return list(
            map(
                lambda s: os.path.join(deploy_dir, os.path.basename(s)),
                filter(lambda s: s.strip() != "", dtb.split(" ")),
            )
        )
    else:
        return []


def efibootguard_generate_uki(args, config, basepath, workspace):
    """
    Entrypoint for devtool generate-uki
    """
    image = args.imagename
    tinfoil = setup_tinfoil(basepath=basepath)

    rd = parse_recipe(config, tinfoil, image, True)
    if not rd:
        # Error already shown
        return 1

    if not bb.data.inherits_class("coreos-image", rd):
        raise InvalidImage("Image should be based on the coreos-image class")

    if rd.getVar("COREOS_IMAGE_GENERATE_UKI") != "1":
        raise InvalidImage("COREOS_IMAGE_GENERATE_UKI should be set to '1'")

    deploy_dir = os.path.relpath(rd.getVar("DEPLOY_DIR_IMAGE"))
    kernel = os.path.relpath(rd.getVar("COREOS_UKI_PART_KERNEL"))
    dtb = dtb_to_list(rd.getVar("KERNEL_DEVICETREE"), deploy_dir)

    rd_efibootguard = parse_recipe(config, tinfoil, "efibootguard", True)

    stub = os.path.relpath(
        os.path.join(
            rd_efibootguard.getVar("WORKDIR"),
            "package/usr/lib/efibootguard",
            rd.getVar("COREOS_UKI_PART_STUB_FILENAME"),
        )
    )

    build_binary = os.path.join(
        rd.getVar("COMPONENTS_DIR"),
        rd.getVar("BUILD_ARCH"),
        "efibootguard-native",
        rd.getVar("bindir_native").lstrip(os.path.sep),
        "bg_gen_unified_kernel",
    )

    keydir = os.path.relpath(rd.getVar("COREOS_EFI_SECUREBOOT_KEYDIR"))

    uki0 = UKIGeneratorArgs(
        kernel=kernel,
        output=os.path.relpath(rd.getVar("COREOS_KERNEL0")),
        cmdline=rd.getVar("COREOS_KERNEL0_CMDLINE"),
        dtb=dtb,
        stub=stub,
        root=rd.getVar("COREOS_ROOTFS0_ROOT"),
        build_binary=build_binary,
        keydir=keydir,
    )
    uki1 = UKIGeneratorArgs(
        kernel=kernel,
        output=os.path.relpath(rd.getVar("COREOS_KERNEL1")),
        cmdline=rd.getVar("COREOS_KERNEL1_CMDLINE"),
        dtb=dtb,
        stub=stub,
        root=rd.getVar("COREOS_ROOTFS1_ROOT"),
        build_binary=build_binary,
        keydir=keydir,
    )

    print(f"Applying passed parameters...")
    uki0.process_args(args)
    uki1.process_args(args)

    print(f"KERNEL0 image will be generated with the following settings:")
    printi(f"{uki0}", 1)
    print()
    print(f"KERNEL1 image will be generated with the following settings:")
    printi(f"{uki1}", 1)
    print()

    print(f"Generating the files...")
    r = uki0.build_and_sign()
    r += uki1.build_and_sign()

    return r
