#!/bin/bash

# SPDX-License-Identifier: GPL-2.0-or-later
# Copyright (C) 2009-2016 Stephan Raue (stephan@openelec.tv)
# Copyright (C) 2016-present Team LibreELEC (https://libreelec.tv)

################################################################################
# variables such as ${ROOT} ${PATH} etc... that are required for this
# script to work must be passed via env ... in scripts/image
################################################################################

# set variables
IMG_TMP=$(mktemp -d -p ${TARGET_IMG})
SAVE_ERROR="${IMG_TMP}/save_error"

if [ "${ROOTFS_TYPE}" = "ext4" ]; then
  if [ -z "${BOOT_SIZE}" -o -z "${ROOT_SIZE}" -o -z "${SYSTEM_PART_START}" ]; then
    echo "mkimage: BOOT_SIZE, ROOT_SIZE, and SYSTEM_PART_START must be configured for ext4!"
    exit 1
  fi
  DISK_START_PADDING=$(((${SYSTEM_PART_START} + 2048 - 1) / 2048))
  DISK_GPT_PADDING=1
  DISK_SIZE=$((${DISK_START_PADDING} + ${BOOT_SIZE} + ${ROOT_SIZE} + ${STORAGE_SIZE} + ${DISK_GPT_PADDING}))
else
  if [ -z "${SYSTEM_SIZE}" -o -z "${SYSTEM_PART_START}" ]; then
    echo "mkimage: SYSTEM_SIZE and SYSTEM_PART_START must be configured!"
    exit 1
  fi
  DISK_START_PADDING=$(((${SYSTEM_PART_START} + 2048 - 1) / 2048))
  DISK_GPT_PADDING=1
  DISK_SIZE=$((${DISK_START_PADDING} + ${SYSTEM_SIZE} + ${STORAGE_SIZE} + ${DISK_GPT_PADDING}))
fi
DISK_BASENAME="${TARGET_IMG}/${IMAGE_NAME}"
if [ -n "${SUBDEVICE}" ]; then
  DISK_BASENAME="${DISK_BASENAME}-${SUBDEVICE}"
fi
DISK="${DISK_BASENAME}.img"

# functions
cleanup() {
  echo -e "image: cleanup...\n"
  rm -rf "${IMG_TMP}"
}

show_error() {
  echo "image: An error has occurred..."
  echo
  if [ -s "${SAVE_ERROR}" ]; then
    cat "${SAVE_ERROR}"
  else
    echo "Folder ${IMG_TMP} might be out of free space..."
  fi
  echo
  cleanup
  exit 1
}

trap cleanup SIGINT

# create an image
echo -e "\nimage: creating sparse file for disk image ${DISK##*/}..."
dd if=/dev/zero of="${DISK}" bs=1M count="${DISK_SIZE}" conv=fsync >"${SAVE_ERROR}" 2>&1 || show_error

if [ "${BOOTLOADER}" = "syslinux" ]; then
  DISK_LABEL=gpt
else
  DISK_LABEL=msdos
fi

# write a disklabel
echo "image: creating ${DISK_LABEL} partition table..."
parted -s "${DISK}" mklabel ${DISK_LABEL}
sync

# create partitions
echo "image: creating partitions..."

if [ "${ROOTFS_TYPE}" = "ext4" ]; then
  # 3-partition layout: boot FAT32 + root ext4 + storage ext4
  BOOT_PART_START=${SYSTEM_PART_START}
  BOOT_PART_END=$((${BOOT_PART_START} + (${BOOT_SIZE} * 1024 * 1024 / 512) - 1))
  ROOT_PART_START=$((${BOOT_PART_END} + 1))
  ROOT_PART_END=$((${ROOT_PART_START} + (${ROOT_SIZE} * 1024 * 1024 / 512) - 1))
  STORAGE_PART_START=$((${ROOT_PART_END} + 1))
  STORAGE_PART_END=$((${STORAGE_PART_START} + (${STORAGE_SIZE} * 1024 * 1024 / 512) - 1))

  if [ "${DISK_LABEL}" = "gpt" ]; then
    parted -s "${DISK}" -a min unit s mkpart boot fat32 ${BOOT_PART_START} ${BOOT_PART_END}
    parted -s "${DISK}" -a min unit s mkpart root ext4 ${ROOT_PART_START} ${ROOT_PART_END}
    parted -s "${DISK}" -a min unit s mkpart storage ext4 ${STORAGE_PART_START} ${STORAGE_PART_END}
    parted -s "${DISK}" set 1 legacy_boot on
  else
    parted -s "${DISK}" -a min unit s mkpart primary fat32 ${BOOT_PART_START} ${BOOT_PART_END}
    parted -s "${DISK}" -a min unit s mkpart primary ext4 ${ROOT_PART_START} ${ROOT_PART_END}
    parted -s "${DISK}" -a min unit s mkpart primary ext4 ${STORAGE_PART_START} ${STORAGE_PART_END}
    parted -s "${DISK}" set 1 boot on
  fi

  # Set SYSTEM_PART_START for boot partition (used by bootloader dd and part1 merge)
  SYSTEM_PART_END=${BOOT_PART_END}
else
  # Original 2-partition layout: system FAT32 + storage ext4
  SYSTEM_PART_END=$((${SYSTEM_PART_START} + (${SYSTEM_SIZE} * 1024 * 1024 / 512) - 1))
  STORAGE_PART_START=$((${SYSTEM_PART_END} + 1))
  STORAGE_PART_END=$((${STORAGE_PART_START} + (${STORAGE_SIZE} * 1024 * 1024 / 512) - 1))

  if [ "${DISK_LABEL}" = "gpt" ]; then
    parted -s "${DISK}" -a min unit s mkpart system fat32 ${SYSTEM_PART_START} ${SYSTEM_PART_END}
    parted -s "${DISK}" -a min unit s mkpart storage ext4 ${STORAGE_PART_START} ${STORAGE_PART_END}
    parted -s "${DISK}" set 1 legacy_boot on
  else
    parted -s "${DISK}" -a min unit s mkpart primary fat32 ${SYSTEM_PART_START} ${SYSTEM_PART_END}
    parted -s "${DISK}" -a min unit s mkpart primary ext4 ${STORAGE_PART_START} ${STORAGE_PART_END}
    parted -s "${DISK}" set 1 boot on
  fi
fi
sync

if [ "${BOOTLOADER}" = "syslinux" ]; then
  # write mbr
  echo "image: writing mbr..."
  dd bs=440 count=1 conv=fsync,notrunc if="${TOOLCHAIN}/share/syslinux/gptmbr.bin" of="${DISK}" >"${SAVE_ERROR}" 2>&1 || show_error
fi

# create part2 to format and copy files
echo "image: creating sparse file for part2..."
STORAGE_PART_COUNT=$((${STORAGE_PART_END} - ${STORAGE_PART_START} + 1))
sync
dd if=/dev/zero of="${IMG_TMP}/part2.ext4" bs=512 count="${STORAGE_PART_COUNT}" conv=fsync >"${SAVE_ERROR}" 2>&1 || show_error

# create filesystem on part2
echo "image: creating filesystem on part2..."
mke2fs -F -q -t ext4 -T ext4 -O ^orphan_file -m 0 "${IMG_TMP}/part2.ext4"
tune2fs -L "${DISTRO_DISKLABEL}" -U ${UUID_STORAGE} "${IMG_TMP}/part2.ext4" >"${SAVE_ERROR}" 2>&1 || show_error
e2fsck -n "${IMG_TMP}/part2.ext4" >"${SAVE_ERROR}" 2>&1 || show_error
sync

# add resize mark
mkdir "${IMG_TMP}/part2.fs"
touch "${IMG_TMP}/part2.fs/.please_resize_me"
echo "image: populating filesystem on part2..."
populatefs -U -d "${IMG_TMP}/part2.fs" "${IMG_TMP}/part2.ext4" >"${SAVE_ERROR}" 2>&1 || show_error
sync
e2fsck -n "${IMG_TMP}/part2.ext4" >"${SAVE_ERROR}" 2>&1 || show_error

# merge part2 into disk image
echo "image: merging part2 into disk image..."
dd if="${IMG_TMP}/part2.ext4" of="${DISK}" bs=512 seek="${STORAGE_PART_START}" conv=fsync,notrunc >"${SAVE_ERROR}" 2>&1 || show_error

# populate root ext4 partition (ext4 layout only)
if [ "${ROOTFS_TYPE}" = "ext4" ]; then
  echo "image: preparing root ext4 partition..."
  # Set label and UUID on the ext4 rootfs image
  tune2fs -L "${DISTRO_ROOTLABEL}" -U ${UUID_ROOT} "${RELEASE_DIR}/target/SYSTEM" >"${SAVE_ERROR}" 2>&1 || show_error

  # Write root ext4 image into disk at ROOT_PART_START
  echo "image: merging root ext4 into disk image..."
  dd if="${RELEASE_DIR}/target/SYSTEM" of="${DISK}" bs=512 seek="${ROOT_PART_START}" conv=fsync,notrunc >"${SAVE_ERROR}" 2>&1 || show_error
fi

# create disk image for virtual appliance
if [ "${PROJECT}" = "Generic" ]; then
  echo "image: creating open virtual appliance..."
  # duplicate ${DISK} so anything we do to it directly doesn't effect original
  dd if="${DISK}" of="${DISK_BASENAME}.tmp" bs=1M conv=fsync >"${SAVE_ERROR}" 2>&1 || show_error
fi

# create part1 to format and copy files
echo "image: creating sparse file for part1..."
SYSTEM_PART_COUNT=$((${SYSTEM_PART_END} - ${SYSTEM_PART_START} + 1))
sync
dd if=/dev/zero of="${IMG_TMP}/part1.fat" bs=512 count="${SYSTEM_PART_COUNT}" conv=fsync >"${SAVE_ERROR}" 2>&1 || show_error

shopt -s expand_aliases # enables alias expansion in script
alias mcopy='mcopy -i "${IMG_TMP}/part1.fat" -o'
alias mmd='mmd -i "${IMG_TMP}/part1.fat"'

# create filesystem on part1
echo "image: creating filesystem on part1..."

if [ "${BOOTLOADER}" = "syslinux" -o "${BOOTLOADER}" = "bcm2835-bootloader" -o "${BOOTLOADER}" = "u-boot" -o "${BOOTLOADER}" = "arm-efi" -o "${BOOTLOADER}" = "qcom-abl" ]; then
  mkfs.vfat -F 32 -S 512 -s 32 -i "${UUID_SYSTEM//-/}" -n "${DISTRO_BOOTLABEL}" "${IMG_TMP}/part1.fat" >"${SAVE_ERROR}" 2>&1 || show_error
fi
sync

if [ "${BOOTLOADER}" = "syslinux" ]; then
  # create bootloader configuration
  echo "image: creating bootloader configuration..."
  cat <<EOF >"${IMG_TMP}/syslinux.cfg"
SAY Wait for installer mode to start automatically in 5 seconds...
SAY
SAY Options
SAY =======
SAY installer: permanently install ${DISTRO} to HDD/SSD
SAY live: boot ${DISTRO} using RAM for temporary storage
SAY run: boot ${DISTRO} using this USB memory device for storage
SAY
DEFAULT installer
TIMEOUT 50
PROMPT 1

LABEL installer
  KERNEL /${KERNEL_NAME}
  APPEND boot=UUID=${UUID_SYSTEM} installer quiet systemd.debug_shell vga=current

LABEL live
  KERNEL /${KERNEL_NAME}
  APPEND boot=UUID=${UUID_SYSTEM} live quiet vga=current

LABEL run
  KERNEL /${KERNEL_NAME}
  APPEND boot=UUID=${UUID_SYSTEM} disk=UUID=${UUID_STORAGE} portable quiet
EOF

  cat <<EOF >"${IMG_TMP}/grub.cfg"
set timeout="25"
set default="Installer"
menuentry "Installer" {
	search --set -f /KERNEL
	linux /KERNEL boot=UUID=${UUID_SYSTEM} installer quiet systemd.debug_shell vga=current
}
menuentry "Live" {
	search --set -f /KERNEL
	linux /KERNEL boot=UUID=${UUID_SYSTEM} grub_live quiet vga=current
}
menuentry "Run" {
	search --set -f /KERNEL
	linux /KERNEL boot=UUID=${UUID_SYSTEM} disk=UUID=${UUID_STORAGE} grub_portable quiet
}
EOF

  mcopy "${IMG_TMP}/syslinux.cfg" :: >"${SAVE_ERROR}" 2>&1 || show_error

  # install syslinux
  echo "image: installing syslinux to part1..."
  syslinux.mtools -i "${IMG_TMP}/part1.fat" >"${SAVE_ERROR}" 2>&1 || show_error

  # copy files
  echo "image: copying files to part1..."
  mcopy "${TARGET_IMG}/${BUILD_NAME}.kernel" "::/${KERNEL_NAME}" >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${RELEASE_DIR}/target/KERNEL.md5" "::/${KERNEL_NAME}.md5" >"${SAVE_ERROR}" 2>&1 || show_error
  if [ "${ROOTFS_TYPE}" != "ext4" ]; then
    mcopy "${TARGET_IMG}/${BUILD_NAME}.system" ::/SYSTEM >"${SAVE_ERROR}" 2>&1 || show_error
    mcopy "${RELEASE_DIR}/target/SYSTEM.md5" ::/SYSTEM.md5 >"${SAVE_ERROR}" 2>&1 || show_error
  fi

  mmd EFI EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${TOOLCHAIN}/share/syslinux/bootx64.efi" ::/EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${TOOLCHAIN}/share/syslinux/ldlinux.e64" ::/EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${TOOLCHAIN}/share/grub/bootia32.efi" ::/EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${IMG_TMP}/grub.cfg" ::/EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error

elif [ "${BOOTLOADER}" = "bcm2835-bootloader" ]; then
  # create bootloader configuration
  echo "image: creating bootloader configuration..."
  cat <<EOF >"${IMG_TMP}/cmdline.txt"
boot=UUID=${UUID_SYSTEM} disk=UUID=${UUID_STORAGE} quiet ${EXTRA_CMDLINE}
EOF

  mcopy "${IMG_TMP}/cmdline.txt" :: >"${SAVE_ERROR}" 2>&1 || show_error

  # copy files
  echo "image: copying files to part1..."
  mcopy "${TARGET_IMG}/${BUILD_NAME}.kernel" "::/${KERNEL_NAME}" >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${RELEASE_DIR}/target/KERNEL.md5" "::/${KERNEL_NAME}.md5" >"${SAVE_ERROR}" 2>&1 || show_error
  if [ "${ROOTFS_TYPE}" != "ext4" ]; then
    mcopy "${TARGET_IMG}/${BUILD_NAME}.system" ::/SYSTEM >"${SAVE_ERROR}" 2>&1 || show_error
    mcopy "${RELEASE_DIR}/target/SYSTEM.md5" ::/SYSTEM.md5 >"${SAVE_ERROR}" 2>&1 || show_error
  fi

  for f in bootcode.bin fixup.dat start.elf; do
    if [ -f "${RELEASE_DIR}/3rdparty/bootloader/$f" ]; then
      mcopy "${RELEASE_DIR}/3rdparty/bootloader/$f" :: >"${SAVE_ERROR}" 2>&1 || show_error
    fi
  done

  mcopy "${RELEASE_DIR}/3rdparty/bootloader/config.txt" :: >"${SAVE_ERROR}" 2>&1 || show_error
  for distro in "${RELEASE_DIR}/3rdparty/bootloader/distroconfig"*.txt; do
    if [ -f "${distro}" ]; then
      mcopy "${distro}" ::/"${distro##*/}" >"${SAVE_ERROR}" 2>&1 || show_error
    fi
  done

  for dtb in "${RELEASE_DIR}/3rdparty/bootloader/"*.dtb; do
    if [ -f "${dtb}" ]; then
      mcopy "${dtb}" ::/"${dtb##*/}" >"${SAVE_ERROR}" 2>&1 || show_error
    fi
  done

  if [ -d "${RELEASE_DIR}/3rdparty/bootloader/overlays" ]; then
    mcopy -s "${RELEASE_DIR}/3rdparty/bootloader/overlays" :: >"${SAVE_ERROR}" 2>&1 || show_error
  fi

elif [ "${BOOTLOADER}" = "u-boot" ]; then
  # create bootloader configuration
  echo "image: creating bootloader configuration... (u-boot)"
  if [ -f "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage"
  elif [ -f "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage"
  else
    echo "image: skipping u-boot. no mkimage script found"
  fi

  echo "image: copying files to part1..."
  mcopy "${TARGET_IMG}/${BUILD_NAME}.kernel" "::/${KERNEL_NAME}" >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${RELEASE_DIR}/target/KERNEL.md5" "::/${KERNEL_NAME}.md5" >"${SAVE_ERROR}" 2>&1 || show_error
  if [ "${ROOTFS_TYPE}" != "ext4" ]; then
    mcopy "${TARGET_IMG}/${BUILD_NAME}.system" ::/SYSTEM >"${SAVE_ERROR}" 2>&1 || show_error
    mcopy "${RELEASE_DIR}/target/SYSTEM.md5" ::/SYSTEM.md5 >"${SAVE_ERROR}" 2>&1 || show_error
  fi

elif [ "${BOOTLOADER}" = "arm-efi" ]; then
  # create bootloader configuration
  echo "image: creating bootloader configuration... (arm-efi)"
  if [ -f "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage"
  elif [ -f "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage"
  else
    echo "image: skipping arm-efi. no mkimage script found"
  fi

  # copy files
  echo "image: copying files to part1..."
  mcopy "${TARGET_IMG}/${BUILD_NAME}.kernel" "::/${KERNEL_NAME}"
  mcopy "${RELEASE_DIR}/target/KERNEL.md5" "::/${KERNEL_NAME}.md5"
  if [ "${ROOTFS_TYPE}" != "ext4" ]; then
    mcopy "${TARGET_IMG}/${BUILD_NAME}.system" ::/SYSTEM
    mcopy "${RELEASE_DIR}/target/SYSTEM.md5" ::/SYSTEM.md5
  fi

elif [ "${BOOTLOADER}" = "qcom-abl" ]; then
  # create bootloader configuration
  echo "image: creating bootloader configuration... (qcom-abl)"
  if [ -f "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/devices/${DEVICE}/bootloader/mkimage"
  elif [ -f "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage" ]; then
    . "${PROJECT_DIR}/${PROJECT}/bootloader/mkimage"
  else
    echo "image: skipping qcom-abl. no mkimage script found"
  fi

  # copy files
  echo "image: copying files to part1..."
  mcopy "${TARGET_IMG}/${BUILD_NAME}.kernel" "::/${KERNEL_NAME}"
  mcopy "${RELEASE_DIR}/target/KERNEL.md5" "::/${KERNEL_NAME}.md5"
  if [ "${ROOTFS_TYPE}" != "ext4" ]; then
    mcopy "${TARGET_IMG}/${BUILD_NAME}.system" ::/SYSTEM
    mcopy "${RELEASE_DIR}/target/SYSTEM.md5" ::/SYSTEM.md5
  fi

fi # bootloader

# run fsck
echo "image: checking filesystem on part1..."
sync
fsck.fat -n "${IMG_TMP}/part1.fat" >"${SAVE_ERROR}" 2>&1 || show_error

# merge part1 into disk image
echo "image: merging part1 into disk image..."
dd if="${IMG_TMP}/part1.fat" of="${DISK}" bs=512 seek="${SYSTEM_PART_START}" conv=fsync,notrunc >"${SAVE_ERROR}" 2>&1 || show_error

# finalize virtual appliance
if [ "${PROJECT}" = "Generic" ]; then
  # change syslinux default to 'run'
  echo "image: modifying files on part1 for open virtual appliance..."
  sed -e "/DEFAULT/ s/installer/run/" -i "${IMG_TMP}/syslinux.cfg"
  sed -e "/set default=/s/\"Installer\"/\"Run\"/" -i "${IMG_TMP}/grub.cfg"
  mcopy "${IMG_TMP}/syslinux.cfg" :: >"${SAVE_ERROR}" 2>&1 || show_error
  mcopy "${IMG_TMP}/grub.cfg" ::/EFI/BOOT >"${SAVE_ERROR}" 2>&1 || show_error
  sync
  # run fsck
  echo "image: checking filesystem on part1..."
  fsck.fat -n "${IMG_TMP}/part1.fat" >"${SAVE_ERROR}" 2>&1 || show_error
  # merge modified part1 into tmp disk image
  echo "image: merging part1 into open virtual appliance..."
  dd if="${IMG_TMP}/part1.fat" of="${DISK_BASENAME}.tmp" bs=512 seek="${SYSTEM_PART_START}" conv=fsync,notrunc >"${SAVE_ERROR}" 2>&1 || show_error
  # create vmdk from tmp ${DISK}
  echo "image: creating vmdk for open virtual appliance..."
  qemu-img convert -O vmdk -o subformat=streamOptimized "${DISK_BASENAME}.tmp" "${DISK_BASENAME}.vmdk"
  # generate ovf from template
  sed -e "s,@DISTRO@,${DISTRO},g" \
    -e "s,@DISTRO_HOME_URL@,${DISTRO_HOME_URL},g" \
    -e "s,@DISK@,${IMAGE_NAME},g" \
    -e "s,@OVA_SIZE@,$((${OVA_SIZE} * 1024 * 1024)),g" \
    "${PROJECT_DIR}/${PROJECT}/config/ovf.template" >"${DISK_BASENAME}.ovf"
  # combine ovf and vmdk into official ova
  tar -C "${TARGET_IMG}" -cf "${DISK_BASENAME}.ova" "${IMAGE_NAME}.ovf" "${IMAGE_NAME}.vmdk"
  # create sha256 checksum of ova image
  (
    cd "${TARGET_IMG}"
    sha256sum "${IMAGE_NAME}.ova" >"${IMAGE_NAME}.ova.sha256"
  )
  echo "image: cleaning up open virtual appliance..."
  # remove tmp ${DISK}, vmdk and ovf
  rm "${DISK_BASENAME}.tmp" "${DISK_BASENAME}.vmdk" "${DISK_BASENAME}.ovf"
fi

# gzip
echo "image: compressing..."
pigz --best --force "${DISK}"

# create sha256 checksum of image
(
  cd "${TARGET_IMG}"
  sha256sum "${DISK##*/}.gz" >"${DISK##*/}.gz.sha256"
)

# cleanup
cleanup
exit
