#!/bin/sh

set -e

TMP_DIR=""
TMP_IMAGE=""
IMAGE_SIZE=${IMAGE_SIZE:-0}
MEGABYTE=$((1024 * 1024))
BOOT_MOUNT=""
LOOP=""
NAME="$(basename "$0")"

# Define these environment variables to override mkfs options
SECTOR_SIZE=${SECTOR_SIZE:-512}
ROOT_DIR_ENTRIES=${ROOT_DIR_ENTRIES:-256}

# Add 64k to the size calculation to reserve some space for the FAT,
# directory entries and rounding up files to cluster sizes.
FAT_OVERHEAD=${FAT_OVERHEAD:-64}

cleanup() {
   unmount_image

   if [ -d "${TMP_DIR}" ]; then
      rm -rf "${TMP_DIR}"
   fi
}

die() {
   echo "$@" >&2
   exit 1
}

createfs() {
   size_mb="$1"
   image="$2"

   volume_label="BOOT"
   if [ -n "${SECTORS_PER_CLUSTER}" ]; then
      SECTORS_PER_CLUSTER="-s ${SECTORS_PER_CLUSTER}"
   fi

   if [ -n "${FAT_SIZE}" ]; then
      fat_size="-F ${FAT_SIZE}"
   fi

   sectors=$((size_mb * MEGABYTE / SECTOR_SIZE))
   sectors=$((sectors / 2))
   /sbin/mkfs.fat -C -f 1 \
      ${SECTORS_PER_CLUSTER} -n "${volume_label}" \
      ${fat_size} \
      -S "${SECTOR_SIZE}" -r "${ROOT_DIR_ENTRIES}" "${image}" ${sectors} || \
      die "Failed to create FAT filesystem"
}

mountfs() {
   image="$1"

   LOOP=$(losetup -f)
   losetup "${LOOP}" "${image}"
   [ -e "${LOOP}" ] ||  die "Failed to create loop device ${LOOP}"

   BOOT_MOUNT=$(mktemp -d)
   mount "${LOOP}" "${BOOT_MOUNT}"
   [ -d "${BOOT_MOUNT}" ] || die "Failed to mount bootfs @ ${BOOT_MOUNT}"

   echo "Mounted ${LOOP} @ ${BOOT_MOUNT}"
}

unmount_image() {
   if [ -d "${BOOT_MOUNT}" ]; then
       umount "${BOOT_MOUNT}" > /dev/null 2>&1 || true
       rmdir "${BOOT_MOUNT}"
       BOOT_MOUNT=""
   fi

   if [ -n "${LOOP}" ]; then
      losetup -d "${LOOP}"
      LOOP=""
   fi
}


createstaging() {
   source_dir="$1"
   staging="$2"
   board="$3"

   mkdir -p "${staging}" || die "Failed to create ${staging}"
   cp -r "${source_dir}/"* "${staging}"

   # Remove files for previous hardware version
   if [ "${board}" = "pi4" ] || [ "${board}" = "pi400" ] || [ "${board}" = "cm4" ]; then
      (
         cd "${staging}"
         rm -f kernel.img kernel7.img bootcode.bin
         rm -f start.elf fixup.dat start_cd.elf fixup_cd.dat start_db.elf fixup_db.dat start_x.elf fixup_x.dat
         rm -f start4cd.elf fixup4cd.dat
         rm -f start4db.elf fixup4db.dat
         rm -f bcm2708* bcm2709* bcm2710*
         rm -f bootcode.bin
      )
   fi

   if [ "${ARCH}" = 32 ]; then
      rm -f "${staging}/kernel8.img"
   elif [ "${ARCH}" = 64 ]; then
      rm -f "${staging}/kernel7l.img"
   fi

   if [ "${board}" = pi400 ]; then
      rm -f "${staging}/start4x.elf"
      rm -f "${staging}/fixup4x.dat"
   fi
   # Estimate the size of the image in KBs
   content="${TMP_DIR}/content.tar"
   echo "$(cd "${staging}"; ls -R)"
   tar -cf "${content}" "${staging}" > /dev/null 2>&1
   if [ "${IMAGE_SIZE}" = 0 ]; then
      IMAGE_SIZE=$(stat --printf "%s" "${content}")
      IMAGE_SIZE=$(((IMAGE_SIZE + 1023) / 1024))
      rm -f "${content}"

      # Add a little padding for FAT etc and convert to megabytes
      IMAGE_SIZE=$((IMAGE_SIZE + FAT_OVERHEAD))
      IMAGE_SIZE=$(((IMAGE_SIZE + 1023) / 1024))
   fi

   echo "Using IMAGE_SIZE of ${IMAGE_SIZE}"

   if [ "${IMAGE_SIZE}" -gt 20 ]; then
      echo "Warning: Large image size detected. Try removing unused files."
   fi
}

checkDependencies() {
   if [ ! -f /sbin/mkfs.fat ]; then
       die "mkfs.fat is required. Run this script on Linux"
   fi
}

usage() {
cat <<EOF
sudo ${NAME} -d SOURCE_DIR -o OUTPUT

Options:
   -a Select 32 or 64 bit kernel
   -b Optionally prune the files to those required for the given board type.
   -d The directory containing the files to include in the boot image.
   -o The filename for the boot image.  -h Display help text and exit

Examples:
# Include all files in bootfs/
sudo ${NAME} -d bootfs/ -o boot.img

# Include only the files from bootfs/ required by Pi 4B
sudo ${NAME} -b pi4 -d bootfs/ -o boot.img

Environment variables:
The following environment variables may be specified to optionally override mkfs.vfat
arguments to help minimise the size of the boot image.


Name               mkfs.vfat parameter
SECTOR_SIZE        -S
ROOT_DIR_ENTRIES   -r
FAT_SIZE           -F

EOF
exit 0
}

SOURCE_DIR=""
OUTPUT=""
ARCH=32
while getopts a:b:d:ho: option; do
   case "${option}" in
   a) ARCH="${OPTARG}"
      ;;
   b) BOARD="${OPTARG}"
      ;;
   d) SOURCE_DIR="${OPTARG}"
      ;;
   o) OUTPUT="${OPTARG}"
      ;;
   h) usage
      ;;
   *) echo "Unknown argument \"${option}\""
      usage
      ;;
   esac
done

checkDependencies

[ -d "${SOURCE_DIR}" ] || usage
[ -n "${OUTPUT}" ] || usage
[ "$(id -u)" = "0" ] || die "\"${NAME}\" must be run as root. Try \"sudo ${NAME}\""
[ -n "${SUDO_UID}" ] || die "SUDO_UID not defined"
[ -n "${SUDO_GID}" ] || die "SUDO_GID not defined"

trap cleanup EXIT
TMP_DIR="$(mktemp -d)"
STAGING="${TMP_DIR}/staging"
rm -f "${OUTPUT}"
echo "Processing source files"
createstaging  "${SOURCE_DIR}" "${STAGING}" "${BOARD}"

echo "Creating FAT file system"
TMP_IMAGE="${TMP_DIR}/boot.img"
createfs ${IMAGE_SIZE} "${TMP_IMAGE}"

echo "Copying files to temporary mount ${BOOT_MOUNT}"
mountfs "${TMP_IMAGE}"
cp -rv "${staging}"/* "${BOOT_MOUNT}"
sync

echo "Sync"
sync

echo "Unmount"
unmount_image

cp -f "${TMP_IMAGE}" "${OUTPUT}"
chown "${SUDO_UID}:${SUDO_GID}" "${OUTPUT}"

echo "Created image ${OUTPUT}"
file "${OUTPUT}"