# Copyright 2021 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""MiniOS build library."""

import logging
import os
import pathlib

from chromite.lib import build_target_lib
from chromite.lib import cros_build_lib
from chromite.lib import image_lib
from chromite.lib import kernel_builder


CROS_DEBUG_FLAG = "cros_debug"
MINIOS_KERNEL_IMAGE = "minios_vmlinuz.image"
KERNEL_FLAGS = [
    "minios",
    "minios_ramfs",
    "tpm",
    "i2cdev",
    "vfat",
    "kernel_compress_xz",
    "pcserial",
    "-kernel_afdo",
]
PART_NAME = "MINIOS-A"
BLOCK_SIZE = 512


class Error(Exception):
    """Base error class for the module."""


class MiniOsError(Error):
    """Raised when failing to build Mini OS image."""


def CreateMiniOsKernelImage(
    board: str,
    version: str,
    work_dir: str,
    keys_dir: str,
    public_key: str,
    private_key: str,
    keyblock: str,
    serial: str,
    build_kernel: bool = True,
    developer_mode: bool = False,
) -> str:
    """Creates the MiniOS kernel image.

    And puts it in the work directory.

    Args:
      board: The board to build the kernel for.
      version: The chromeos version string.
      work_dir: The directory for keeping intermediary files.
      keys_dir: The path to kernel keys directories.
      public_key: Filename to the public key whose private part signed the
                  keyblock.
      private_key: Filename to the private key whose public part is baked into
                   the keyblock.
      keyblock: Filename to the kernel keyblock.
      serial: Serial port for the kernel console (e.g. printks).
      build_kernel: Build a new kernel from source.
      developer_mode: Add developer mode flags to the kernel image.

    Returns:
      The path to the generated kernel image.
    """
    install_root = os.path.join(
        (build_target_lib.get_default_sysroot_path(board)), "factory-root"
    )
    kb = kernel_builder.Builder(board, work_dir, install_root)
    if build_kernel:
        # MiniOS ramfs cannot be built with multiple conflicting `_ramfs` flags.
        kb.CreateCustomKernel(
            KERNEL_FLAGS,
            [
                x
                for x in os.environ.get("USE", "").split()
                if not x.endswith("_ramfs")
            ],
        )
    kernel = os.path.join(work_dir, MINIOS_KERNEL_IMAGE)
    assert " " not in version, f"bad version: {version}"
    boot_args = f"noinitrd panic=60 cros_minios_version={version} cros_minios"
    if developer_mode:
        boot_args += f" {CROS_DEBUG_FLAG}"
    kb.CreateKernelImage(
        kernel,
        boot_args=boot_args,
        serial=serial,
        keys_dir=keys_dir,
        public_key=public_key,
        private_key=private_key,
        keyblock=keyblock,
    )
    return kernel


def InsertMiniOsKernelImage(image: str, kernel: str):
    """Writes the MiniOS into MINIOS-A partition of the Chromium OS image.

    It assumes the Chromium OS image has enough space in MINIOS-A partition,
    otherwise it will fail.

    Args:
      image: The path to the Chromium OS image.
      kernel: The path to the kernel image.
    """
    with image_lib.LoopbackPartitions(image) as devs:
        part_info = devs.GetPartitionInfo(PART_NAME)
        if pathlib.Path(kernel).stat().st_size > part_info.size:
            raise MiniOsError(
                "MiniOS kernel image is larger than the MINIOS-A partition."
            )

        device = devs.GetPartitionDevName(PART_NAME)
        logging.debug("Mini OS loopback partition is %s", device)

        logging.info(
            "Writing the MiniOS kernel image %s into the Chromium OS "
            "image %s",
            kernel,
            image,
        )
        # First zero out the partition. This generally would help with update
        # payloads so we don't have to compress junk bytes. The target file is a
        # loopback dev device of the MINIOS-A partition.
        cros_build_lib.sudo_run(
            [
                "dd",
                "if=/dev/zero",
                f"of={device}",
                f"bs={BLOCK_SIZE}",
                f"count={part_info.size // BLOCK_SIZE}",
            ]
        )
        # Write the actual MiniOS kernel into the MINIOS-A partition.
        cros_build_lib.sudo_run(
            ["dd", f"if={kernel}", f"of={device}", f"bs={BLOCK_SIZE}"]
        )