# This class generate a fitimage
# This class should ihnerited after the nwl-image class

# Configuration:
# ==============================================================================

# Set the initramfs to the current image being build by default
# This can be overriden if another image as to be included
INITRAMFS_IMAGE ??= "${IMAGE_BASENAME}" 
INITRAMFS_IMAGE_NAME ??= "${INITRAMFS_IMAGE}-${MACHINE}"

# Glue for kernel-fitimage
# ==============================================================================

kernel_do_deploy () {
    # The kernel-fitimage class appends the deployment to this task. We just have to provide the task.
    deployDir=${DEPLOY_DIR_IMAGE}
}

inherit kernel-fitimage
inherit kernel-arch
inherit ${@bb.utils.contains('MACHINE_FEATURES', 'fpga', 'kernel-fitimage-fpga', '', d)}


do_assemble_fitimage:prepend() {
    install -d ${S}/arch/${ARCH}/boot/
    install -m 0644 ${DEPLOY_DIR_IMAGE}/${KERNEL_IMAGETYPE}-${KERNEL_IMAGE_LINK_NAME}${KERNEL_IMAGE_BIN_EXT} ${S}/arch/${ARCH}/boot/vmlinuz.bin

    for DTB in ${KERNEL_DEVICETREE}; do
        # We have to watch out for when the device-tree is located in its own subdirectory in the kernel sources.
        # DTB_SUBDIR will contain this directory.
        DTB_DEPLOY=`basename ${DTB}`
        DTB_SUBDIR=`dirname ${DTB}`
        install -d arch/${ARCH}/boot/${DTB_SUBDIR}
        install -m 0644 ${DEPLOY_DIR_IMAGE}/${DTB_DEPLOY} ${S}/arch/${ARCH}/boot/${DTB}
    done
}

KERNEL_CLASSES = "kernel-fitimage"
KERNEL_IMAGETYPES += "fitImage"

# We want the initramfs inside the fitimage as a separate part and not bundle
# inside the kernel part
INITRAMFS_IMAGE_BUNDLE = "0"


do_image_fitimage() {
    kernel_do_deploy
}


addtask image_fitimage after do_image before do_image_complete
addtask assemble_fitimage_initramfs after do_image before do_image_fitimage
addtask assemble_fitimage after do_image before do_image_fitimage

do_assemble_fitimage_initramfs[depends] += "virtual/kernel:do_deploy virtual/${TARGET_PREFIX}binutils:do_populate_sysroot"
do_assemble_fitimage[depends] += "virtual/kernel:do_deploy virtual/${TARGET_PREFIX}binutils:do_populate_sysroot"


# a fitimage image can be embedded into a WIC image
do_image_wic[recrdeptask] += "do_image_fitimage"

python __anonymous () {
        initramfs = d.getVar('INITRAMFS_IMAGE')
        image = d.getVar('IMAGE_BASENAME')

        if image == initramfs:

            # If the initramfs is the same as the image being build, we should
            # not depend on do_image_complete but on do_image_${FSTYPE_WITHOUT_EXT}
            fstype = d.getVar('IMAGE_FSTYPES').split('.')[0]
            d.appendVarFlag('do_assemble_fitimage_initramfs', 'recrdeptask', f' do_image_{fstype}')
            depends = d.getVarFlag('do_assemble_fitimage_initramfs', 'depends')
            depends = depends.replace(f'{initramfs}:do_image_complete', '')
            d.setVarFlag('do_assemble_fitimage_initramfs', 'depends', depends)


            # If the initramfs is the same as the image being build, it should
            # come from IMGDEPLOYDIR instead of DEPLOY_DIR_IMAGE
            # This is done by using an ugly monkey patching of the
            # fitimage_assemble function
            imgdeploydir = d.getVar('IMGDEPLOYDIR')
            deploy_dir_image = d.getVar('DEPLOY_DIR_IMAGE')
            task = d.getVar('fitimage_assemble')
            task = task.replace(f'initramfs_path="{deploy_dir_image}', f'initramfs_path="{imgdeploydir}')
            d.setVar('fitimage_assemble', task)
}

# Glue for coreos-image-ci.bbclass
# ==============================================================================

def get_nwl_fitimage_ci_artifacts(d):
    bundle = d.getVar('INITRAMFS_IMAGE_BUNDLE')
    initramfs = d.getVar('INITRAMFS_IMAGE_NAME')

    # We only support fitimage with INITRAMFS_IMAGE_BUNDLE set to 0
    if bundle == "1":
        bb.warn(f"Adding a fitimage built with INITRAMFS_IMAGE_BUNDLE into COREOS_CI_DEPLOY_ARTIFACTS is currently not supported")
        return ""

    # If an initramfs is used, publish the fitImage that contains it
    # otherwise the image without an initramfs
    if initramfs.strip() != "":
        return "fitImage-${INITRAMFS_IMAGE_NAME}-${KERNEL_FIT_LINK_NAME}"
    else:
        return "fitImage-linux.bin-${KERNEL_FIT_LINK_NAME}"

# Add the generated fitImage to the list of artifacts to publish in the CI
COREOS_CI_DEPLOY_ARTIFACTS += "${@get_nwl_fitimage_ci_artifacts(d)}"