#!/bin/bash -eux

case $ARCH in
    amd64|arm64|armhf|riscv64)
        ;;
    *)
        echo "We don't create EFI images for $ARCH."
        exit 0
        ;;
esac

case ${PROJECT:-} in
    ubuntu)
        IMAGE_STR="# DESKTOP_IMG: This file was created/modified by the Desktop Image build process"
        FS_LABEL="desktop-rootfs"
        IMAGE_SIZE=12884901888 # 12G
        ;;
    *)
        IMAGE_STR="# CLOUD_IMG: This file was created/modified by the Cloud Image build process"
        FS_LABEL="cloudimg-rootfs"
        ;;
esac

# Change image size for preinstalled generic images & all preinstalled riscv64 images
if [ -n "${SUBARCH:-}" ]; then
    if [ "${SUBARCH:-}" = "generic" ] || [ "$ARCH" = "riscv64" ]; then
        IMAGE_SIZE=3758096384 # bump to 3.5G (3584*1024**2), due to linux-generic instead of virtual
fi
fi

. config/binary

. config/functions

create_partitions() {
    disk_image="$1"
    sgdisk "${disk_image}" --zap-all
    case $ARCH in
        arm64|armhf)
            if [ "${SUBARCH:-}" = "generic" ]; then
                sgdisk "${disk_image}" \
                       --new=15:0:204800 \
                       --typecode=15:ef00 \
                       --attributes=15:set:2 \
                       --new=14::+4M \
                       --change-name=14:CIDATA \
                       --new=1:
            else
                sgdisk "${disk_image}" \
                       --new=15:0:204800 \
                       --typecode=15:ef00 \
                       --new=1:
            fi
            ;;
        riscv64)
            # same as arm64/armhf, but set bit 2 legacy bios bootable
            # on the first partition for uboot
            # and have two loader partitions of uboot SPL & real one
            # and have CIDATA partition for preinstalled image
            if [ -z "${SUBARCH:-}" ]; then
                # cloud-image
                sgdisk "${disk_image}" \
                       --set-alignment=2 \
                       --new=15::+106M \
                       --typecode=15:ef00 \
                       --new=1:: \
                       --attributes=1:set:2
            else
                # preinstalled server, currently FU540
                # FU740 too in the future
                sgdisk "${disk_image}" \
                       --set-alignment=2 \
                       --new=13:34:2081 \
                       --change-name=13:loader1 \
                       --typecode=13:5B193300-FC78-40CD-8002-E86C45580B47 \
                       --new=14:2082:10273 \
                       --change-name=14:loader2 \
                       --typecode=14:2E54B353-1271-4842-806F-E436D6AF6985 \
                       --new=15::+106M \
                       --typecode=15:ef00 \
                       --new=12::+4M \
                       --change-name=12:CIDATA \
                       --new=1:: \
                       --attributes=1:set:2
            fi
            ;;
        amd64)
            if [ "${SUBARCH:-}" = "generic" ]; then
                sgdisk "${disk_image}" \
                       --new=14::+4M \
                       --typecode=14:ef02 \
                       --attributes=14:set:2 \
                       --new=15::+106M \
                       --typecode=15:ef00 \
                       --new=13::+4M \
                       --change-name=13:CIDATA \
                       --new=1::
            else
                sgdisk "${disk_image}" \
                       --new=14::+4M \
                       --new=15::+106M \
                       --new=1::
                sgdisk "${disk_image}" \
                       -t 14:ef02 \
                       -t 15:ef00
            fi
            ;;
    esac
    sgdisk "${disk_image}" \
        --print
}

create_and_mount_uefi_partition() {
    uefi_dev="/dev/mapper${loop_device///dev/}p15"
    mountpoint="$1"
    mkfs.vfat -F 32 -n UEFI "${uefi_dev}"

    mkdir -p "${mountpoint}"/boot/efi
    mount "${uefi_dev}" "$mountpoint"/boot/efi

    cat << EOF >> "mountpoint/etc/fstab"
LABEL=UEFI	/boot/efi	vfat	umask=0077	0 1
EOF
}

install_grub() {
    mkdir mountpoint
    mount_partition "${rootfs_dev_mapper}" mountpoint

    create_and_mount_uefi_partition mountpoint

    echo "(hd0) ${loop_device}" > mountpoint/tmp/device.map
    mkdir -p mountpoint/etc/default/grub.d
    efi_boot_dir="/boot/efi/EFI/BOOT"
    chroot mountpoint mkdir -p "${efi_boot_dir}"

    chroot mountpoint apt-get -y update

    # UEFI GRUB modules are meant to be used equally by Secure Boot and
    # non-Secure Boot systems. If you need an extra module not already
    # provided or run into "Secure Boot policy forbids loading X" problems,
    # please file a bug against grub2 to include the affected module.
    case $ARCH in
        arm64)
            chroot mountpoint apt-get -qqy install --no-install-recommends shim-signed grub-efi-arm64-signed
            efi_target=arm64-efi
            if [ "${SUBARCH:-}" = "generic" ]; then
                # Server preinstalled image
                # Setup cidata sample data & nocloud fallback
                # Allows login on first boot with or without metadata
                cidata_dev="/dev/mapper${loop_device///dev/}p14"
                setup_cidata "${cidata_dev}"
                setup_cinocloud mountpoint
            fi
            ;;
        armhf)
            chroot mountpoint apt-get -qqy install --no-install-recommends grub-efi-arm grub-efi-arm-bin
            efi_target=arm-efi
            if [ "${SUBARCH:-}" = "generic" ]; then
                # Server preinstalled image
                # Setup cidata sample data & nocloud fallback
                # Allows login on first boot with or without metadata
                cidata_dev="/dev/mapper${loop_device///dev/}p14"
                setup_cidata "${cidata_dev}"
                setup_cinocloud mountpoint
            fi
            ;;
        amd64)
            chroot mountpoint apt-get install -qqy grub-pc shim-signed
            efi_target=x86_64-efi
            if [ "${SUBARCH:-}" = "generic" ]; then
                # Server preinstalled image
                # Setup cidata sample data & nocloud fallback
                # Allows login on first boot with or without metadata
                cidata_dev="/dev/mapper${loop_device///dev/}p13"
                setup_cidata "${cidata_dev}"
                setup_cinocloud mountpoint
            fi
            ;;
        riscv64)
            # TODO grub-efi-riscv64 does not exist yet on riscv64
            chroot mountpoint apt-get install -qqy u-boot-menu #grub-efi-riscv64
            efi_target=riscv64-efi

            chroot mountpoint u-boot-update

            if [ -n "${SUBARCH:-}" ]; then
                u_boot_arch="${SUBARCH}"
                if [ "${u_boot_arch}" = "hifive" ]; then
                    u_boot_arch=sifive_fu540
                fi
                chroot mountpoint apt-get install -qqy u-boot-sifive
                # FSBL, which gets U-Boot SPL
                loader1="/dev/mapper${loop_device///dev/}p13"
                # The real U-Boot
                loader2="/dev/mapper${loop_device///dev/}p14"
                dd if=mountpoint/usr/lib/u-boot/${u_boot_arch}/u-boot-spl.bin of=$loader1
                dd if=mountpoint/usr/lib/u-boot/${u_boot_arch}/u-boot.itb of=$loader2
                # Provide end-user modifyable CIDATA
                cidata_dev="/dev/mapper${loop_device///dev/}p12"
                setup_cidata "${cidata_dev}"
                # Provide stock nocloud datasource
                # Allow interactive login on baremetal SiFive board,
                # without a cloud datasource.
                setup_cinocloud mountpoint
            fi
            ## TODO remove below once we have grub-efi-riscv64
            rm mountpoint/tmp/device.map
            umount mountpoint/boot/efi
            mount
            umount_partition mountpoint
            rmdir mountpoint
            return
            ##
            ;;
    esac

    chroot mountpoint apt-get autoremove --purge --assume-yes

    chroot mountpoint grub-install "${loop_device}" \
        --boot-directory=/boot \
        --efi-directory=/boot/efi \
        --target=${efi_target} \
        --uefi-secure-boot \
        --no-nvram

    if [ "$ARCH" = "amd64" ]; then
        # Install the BIOS/GPT bits. Since GPT boots from the ESP partition,
        # it means that we just run this simple command and we're done
        chroot mountpoint grub-install --target=i386-pc "${loop_device}"
    fi

    # Use the linux-kvm kernel for minimal images where available
    # linux-kvm currently only exists for amd64
    if [ "${SUBPROJECT:-}" = "minimized" ] && [ "$ARCH" = "amd64" ]; then
        replace_kernel mountpoint linux-kvm
    fi

    # This call to rewrite the debian package manifest is added here to capture
    # grub-efi packages that otherwise would not make it into the base
    # manifest. filesystem.packages is moved into place via symlinking to
    # livecd.ubuntu-cpc.manifest by live-build/auto/build after lb_binary runs
    # and at that time snaps are added to the manifest (create-manifest is
    # not called here as it calls snap-seed-parse, resulting in duplicate
    # snap listings)
    chroot mountpoint dpkg-query -W > binary/boot/filesystem.packages

    divert_grub mountpoint
    track_initramfs_boot_fallback mountpoint
    chroot mountpoint update-grub
    replace_grub_root_with_label mountpoint
    undivert_grub mountpoint

    chroot mountpoint apt-get -y clean

    rm mountpoint/tmp/device.map
    umount mountpoint/boot/efi
    mount
    umount_partition mountpoint
    rmdir mountpoint
}

disk_image=binary/boot/disk-uefi.ext4

create_empty_disk_image "${disk_image}"
create_partitions "${disk_image}"
mount_image "${disk_image}" 1

# Copy the chroot in to the disk
make_ext4_partition "${rootfs_dev_mapper}"
mkdir mountpoint
mount "${rootfs_dev_mapper}" mountpoint
cp -a chroot/* mountpoint/
umount mountpoint
rmdir mountpoint

install_grub

clean_loops
trap - EXIT