#!/bin/sh

set -e

TMP_DIR=""
TMP_IMAGE=""
IMAGE_SIZE=${IMAGE_SIZE:-0}
if [ "$IMAGE_SIZE" -ne 0 ]; then
   IMAGE_SIZE="$((IMAGE_SIZE * 1024))"
fi
BOOT_MOUNT=""
LOOP=""
NAME="$(basename "$0")"

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

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

HAVE_MCOPY=false

cleanup() {
   unmount_image

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

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

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

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

   if [ -n "${DISK_GEOMETRY}" ]; then
      DISK_GEOMETRY="-g ${DISK_GEOMETRY}"
   fi

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

   mkfs.fat -C -f 1 \
      ${SECTORS_PER_CLUSTER} -n "${volume_label}" \
      ${fat_size} ${DISK_GEOMETRY} \
      -S "${SECTOR_SIZE}" -r "${ROOT_DIR_ENTRIES}" "${image}" ${size_kib} || \
      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
}

copyfiles() {
   image="$1"
   shift
   if ${HAVE_MCOPY} ; then
      mcopy -i "${image}" -vsmpQ "$@" ::/
   else
      mountfs "${image}"
      cp -rpv "$@" "${BOOT_MOUNT}"
      sync

      echo "Sync"
      sync

      echo "Unmount"
      unmount_image
   fi
}

estimate_fat_clusters() {
   local cluster_size
   cluster_size="$1"
   local directory
   directory="$2"

   local clusters
   clusters="0"
   local subdir

   for subdir in $(find "$directory" -mindepth 1 -maxdepth 1 -type d); do
      local sd_clusters
      sd_clusters="$(estimate_fat_clusters "$cluster_size" "$subdir")"
      clusters="$((clusters + sd_clusters))"
   done

   # Determine number of clusters required for file contents
   local file_clusters
   file_clusters="$(find "$directory" -maxdepth 1 -type f -exec du --apparent-size --block-size="${cluster_size}" {} + | awk '{clust=clust+$1} END {print clust}')"
   clusters="$((clusters + file_clusters))"

   # Determine number of clusters required for directory entries
   # Two additional entries are required for "." and ".."
   local dir_entries
   dir_entries="2"
   local file_name
   local file_base_name

   for file_name in "${directory}"/*; do
      file_base_name="$(basename "${file_name}")"

      # Always at least one entry
      dir_entries="$((dir_entries + 1))"

      # MSDOS 8.3 Filename Requirements
      # A-Z
      # 0-9
      # Space
      # ! # $ % & ' ( ) - @ ^ _ ` { } ~
      # Values 128-255
      VALID_MSDOS_FILENAME_CHAR='[A-Z0-9 \!#\$%\&'"'"'\(\)\-@\^_`\{\}~\x80-\xff]'
      VALID_MSDOS_FILENAME="^${VALID_MSDOS_FILENAME_CHAR}{1,8}(?:\.${VALID_MSDOS_FILENAME_CHAR}{1,3})$"

      local msdos_compatible_filename
      msdos_compatible_filename="0"
      printf "%s" "${file_base_name}" | \
         grep --silent --perl-regexp "${VALID_MSDOS_FILENAME}" || msdos_compatible_filename=$?

      if [ "$msdos_compatible_filename" -ne "0" ]; then
         local ucs2_bytes
         ucs2_bytes="$(printf "%s" "${file_name}" | iconv --to-code=UCS2 | wc --bytes)"
         dir_entries="$((dir_entries + ((ucs2_bytes + 25)/26)))"
      fi
   done

   # 32-bytes are required for each entry
   clusters="$((clusters + ((dir_entries * 32) + cluster_size - 1)/cluster_size))"
   echo "$clusters"
}

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

   mkdir -p "${staging}" || die "Failed to create ${staging}"
   cp -rp "${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 start4x.elf fixup4x.dat
         rm -f bcm2708* bcm2709* bcm2710*
         rm -f kernel_2712.img initramfs_2712
         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

   if [ "${IMAGE_SIZE}" = 0 ]; then
      # Estimate the size of the image in clusters
      cluster_size="$((SECTOR_SIZE * SECTORS_PER_CLUSTER))"
      clusters="$(estimate_fat_clusters "${cluster_size}" "${staging}")"

      IMAGE_SIZE="$((clusters * cluster_size))"

      root_dir_sectors="$((((ROOT_DIR_ENTRIES * 32) + cluster_size - 1)/cluster_size))"
      # FAT32/FAT16 determined by number of clusters
      if [ "$clusters" -gt "65526" ]; then
         FAT_SIZE="32"
         fat_table_sectors="$(((((clusters + 2)*4) + SECTOR_SIZE - 1)/SECTOR_SIZE))"
      elif [ "$clusters" -gt "4085" ]; then
         FAT_SIZE="16"
         fat_table_sectors="$(((((clusters + 2)*2) + SECTOR_SIZE - 1)/SECTOR_SIZE))"
         # Add some sectors based on ROOT_DIR_ENTRIES
         fat_table_sectors="$((fat_table_sectors + root_dir_sectors))"
      else
         FAT_SIZE="12"
         #12 bits per cluster = 3 bytes for 2 clusters
         fat_table_sectors=$(((((clusters + 2)*3+1) / 2 + SECTOR_SIZE - 1)/SECTOR_SIZE))
         # Add some sectors based on ROOT_DIR_ENTRIES
         fat_table_sectors="$((fat_table_sectors + root_dir_sectors))"
      fi
      IMAGE_SIZE="$((IMAGE_SIZE + fat_table_sectors * SECTOR_SIZE))"
      IMAGE_SIZE="$(((IMAGE_SIZE + 1023)/1024))"

      # Add a little padding for FAT etc
      IMAGE_SIZE=$((IMAGE_SIZE + FAT_OVERHEAD))
   fi

   echo "Using IMAGE_SIZE of ${IMAGE_SIZE}"

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

checkDependencies() {
   if ! mkfs.fat --help > /dev/null 2> /dev/null ; then
       die "mkfs.fat is required. Run this script on Linux"
   fi
   if mcopy --help > /dev/null 2> /dev/null ; then
       HAVE_MCOPY=true
   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
$HAVE_MCOPY || [ "$(id -u)" = "0" ] || die "Install mtools or use \"sudo ${NAME}\""

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 file system image ${TMP_IMAGE}"
copyfiles "${TMP_IMAGE}" "${staging}"/*

cp -f "${TMP_IMAGE}" "${OUTPUT}"
if [ -n "${SUDO_UID}" ] && [ -n "${SUDO_GID}" ] ; then
   chown "${SUDO_UID}:${SUDO_GID}" "${OUTPUT}"
fi

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