#!/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}
# Calculated from the source files unless set
ROOT_DIR_ENTRIES=${ROOT_DIR_ENTRIES:-0}

# mkfs.fat derives the FAT width from the number of clusters and rejects a -F
# option that disagrees with it. Cluster counts between MAX_CLUSTERS_FAT12 and
# MIN_CLUSTERS_FAT16 are accepted at neither width.
MAX_CLUSTERS_FAT12=4084
MIN_CLUSTERS_FAT16=4087
MAX_CLUSTERS_FAT16=65524

# Spare space, in KiB, to leave free in the filesystem
FAT_OVERHEAD=${FAT_OVERHEAD:-0}

# mkfs.fat refuses to create a filesystem whose data region is much smaller
# than this, whatever the sector and cluster sizes in use.
MIN_DATA_KIB=64

HAVE_MCOPY=false

cleanup() {
   unmount_image

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

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

# FAT32 reserves room for the FSInfo sector and for the backup boot sector,
# which must be at sector 6. These match the mkfs.fat defaults.
reserved_sectors() {
   if [ "$1" = "32" ]; then
      echo 32
   else
      echo 1
   fi
}

# Sectors occupied by a FAT holding the given number of clusters
fat_sectors() {
   fs_entries="$(($1 + 2))"

   case "$2" in
   12) fs_bytes="$(((fs_entries * 3 + 1) / 2))" ;;
   16) fs_bytes="$((fs_entries * 2))" ;;
   32) fs_bytes="$((fs_entries * 4))" ;;
   *)  die "Unsupported FAT size $2" ;;
   esac

   echo "$(((fs_bytes + SECTOR_SIZE - 1) / SECTOR_SIZE))"
}

# Total sectors needed for the reserved sectors, the FAT, the root directory
# and the given number of data clusters
layout_sectors() {
   ls_clusters="$1"
   ls_width="$2"
   ls_root_entries="$3"

   ls_sectors="$(reserved_sectors "${ls_width}")"
   ls_sectors="$((ls_sectors + $(fat_sectors "${ls_clusters}" "${ls_width}")))"

   # FAT32 holds the root directory in the data region
   if [ "${ls_width}" != "32" ]; then
      ls_sectors="$((ls_sectors + ls_root_entries * 32 / SECTOR_SIZE))"
   fi

   echo "$((ls_sectors + ls_clusters * SECTORS_PER_CLUSTER))"
}

# Largest number of data clusters whose layout fits in the given sectors. The
# size of the FAT depends on the number of clusters it describes, so search
# for the largest count that still fits rather than solving directly.
max_clusters() {
   mc_lo=1
   mc_hi="$(($1 / SECTORS_PER_CLUSTER + 1))"

   while [ "${mc_lo}" -lt "${mc_hi}" ]; do
      mc_mid="$(((mc_lo + mc_hi + 1) / 2))"
      if [ "$(layout_sectors "${mc_mid}" "$2" "$3")" -le "$1" ]; then
         mc_lo="${mc_mid}"
      else
         mc_hi="$((mc_mid - 1))"
      fi
   done

   echo "${mc_lo}"
}

# Narrowest FAT that mkfs.fat will accept for a filesystem of the given number
# of sectors. Left to itself mkfs.fat picks the width from the size of the
# device alone, which fails whenever the cluster size in use puts the cluster
# count outside the range permitted for that width.
select_fat_width() {
   sf_clusters="$(max_clusters "$1" 12 "$2")"
   if [ "${sf_clusters}" -le "${MAX_CLUSTERS_FAT12}" ]; then
      echo 12
      return
   fi

   sf_clusters="$(max_clusters "$1" 16 "$2")"
   if [ "${sf_clusters}" -ge "${MIN_CLUSTERS_FAT16}" ] && \
      [ "${sf_clusters}" -le "${MAX_CLUSTERS_FAT16}" ]; then
      echo 16
      return
   fi

   echo 32
}

createfs() {
   image_sectors="$1"
   image="$2"
   root_dir_entries="$3"

   volume_label="BOOT"

   rm -f "${image}"
   truncate -s "$((image_sectors * SECTOR_SIZE))" "${image}" || \
      die "Failed to allocate ${image}"

   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

   # -a suppresses alignment of the FAT and root directory to cluster
   # boundaries, which would otherwise consume sectors not budgeted for above.
   mkfs.fat -a -f 1 \
      ${SECTORS_PER_CLUSTER} -n "${volume_label}" \
      ${fat_size} ${DISK_GEOMETRY} \
      -S "${SECTOR_SIZE}" -r "${root_dir_entries}" "${image}" || \
      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
}

# Extract directory entry calculation into a separate function
estimate_directory_entries() {
   local directory
   directory="$1"
   local is_root_dir
   is_root_dir="${2:-0}"  # Default to 0 (not root directory)

   # Determine number of directory entries required
   local dir_entries
   if [ "$is_root_dir" = "1" ]; then
      # The root directory has no "." or ".." entries but does hold the
      # volume label
      dir_entries="1"
   else
      # Two additional entries are required for "." and ".." in subdirectories
      dir_entries="2"
   fi

   local file_name
   local file_base_name

   for file_name in "${directory}"/*; do
      # An unmatched glob expands to the pattern itself. A dangling symbolic
      # link still occupies a directory entry.
      [ -e "${file_name}" ] || [ -L "${file_name}" ] || continue

      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_base_name}" | iconv --to-code=UCS2 | wc --bytes)"
         dir_entries="$((dir_entries + ((ucs2_bytes + 25)/26)))"
      fi
   done

   echo "$dir_entries"
}

# Specific function for root directory entries in FAT12/16
estimate_root_directory_entries() {
   local directory
   directory="$1"
   estimate_directory_entries "$directory" 1
}

estimate_fat_clusters() {
   local cluster_size
   cluster_size="$1"
   local directory
   directory="$2"
   local include_root_dir_entries
   include_root_dir_entries="${3:-1}"  # Default to 1 (include) for backwards compatibility

   local clusters
   clusters="0"
   local subdir

   for subdir in $(find "$directory" -mindepth 1 -maxdepth 1 -type d); do
      local sd_clusters
      # Subdirectories always include directory entries regardless of FAT type
      sd_clusters="$(estimate_fat_clusters "$cluster_size" "$subdir" 1)"
      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))"

   # Only include directory entry clusters if requested (FAT32 or non-root directories)
   if [ "$include_root_dir_entries" = "1" ]; then
      local dir_entries
      dir_entries="$(estimate_directory_entries "$directory")"

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

   echo "$clusters"
}

# Entry point for FAT32 (includes all directory entries as clusters)
estimate_fat_clusters_fat32() {
   local cluster_size
   cluster_size="$1"
   local directory
   directory="$2"

   estimate_fat_clusters "$cluster_size" "$directory" 1
}

# Entry point for FAT12/FAT16 (excludes root directory entries from cluster calculation)
estimate_fat_clusters_fat12_16() {
   local cluster_size
   cluster_size="$1"
   local directory
   directory="$2"

   estimate_fat_clusters "$cluster_size" "$directory" 0
}

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

   cluster_size="$((SECTOR_SIZE * SECTORS_PER_CLUSTER))"

   # Calculate actual root directory entries needed for FAT12/16. The Linux
   # vfat driver refuses to mount a filesystem whose root directory does not
   # span a whole number of sectors.
   required_root_entries="$(estimate_root_directory_entries "${staging}")"
   estimated_root_entries="${required_root_entries}"
   if [ "${ROOT_DIR_ENTRIES}" != 0 ]; then
      estimated_root_entries="${ROOT_DIR_ENTRIES}"
      if [ "${ROOT_DIR_ENTRIES}" -lt "${required_root_entries}" ]; then
         echo "Warning: ROOT_DIR_ENTRIES=${ROOT_DIR_ENTRIES} is fewer than the" \
            "${required_root_entries} entries the source files require"
      fi
   fi
   entries_per_sector="$((SECTOR_SIZE / 32))"
   ROOT_ENTRIES="$(((estimated_root_entries + entries_per_sector - 1) \
      / entries_per_sector * entries_per_sector))"

   if [ "${IMAGE_SIZE}" = 0 ]; then
      # Estimate the size of the image in clusters
      overhead_clusters="$(((FAT_OVERHEAD * 1024 + cluster_size - 1)/cluster_size))"

      # First, estimate assuming FAT32 to determine potential FAT type
      clusters="$(($(estimate_fat_clusters_fat32 "${cluster_size}" "${staging}") \
         + overhead_clusters))"

      # FAT32/FAT16 determined by number of clusters
      if [ "$clusters" -gt "${MAX_CLUSTERS_FAT16}" ]; then
         FAT_SIZE="32"
      else
         # Will be FAT12 or FAT16, so recalculate without the root directory
         # entries, which live outside the data region at those widths
         clusters="$(($(estimate_fat_clusters_fat12_16 "${cluster_size}" "${staging}") \
            + overhead_clusters))"

         if [ "$clusters" -gt "${MAX_CLUSTERS_FAT12}" ]; then
            FAT_SIZE="16"
            if [ "$clusters" -lt "${MIN_CLUSTERS_FAT16}" ]; then
               clusters="${MIN_CLUSTERS_FAT16}"
            fi
         else
            FAT_SIZE="12"
         fi
      fi

      # mkfs.fat refuses to create a filesystem with a very small data region
      min_clusters="$(((MIN_DATA_KIB * 1024 + cluster_size - 1)/cluster_size))"
      if [ "$clusters" -lt "$min_clusters" ]; then
         clusters="$min_clusters"
      fi

      IMAGE_SECTORS="$(layout_sectors "${clusters}" "${FAT_SIZE}" "${ROOT_ENTRIES}")"
      IMAGE_SIZE="$(((IMAGE_SECTORS * SECTOR_SIZE + 1023)/1024))"
   else
      IMAGE_SECTORS="$((IMAGE_SIZE * 1024 / SECTOR_SIZE))"
      if [ -z "${FAT_SIZE}" ]; then
         FAT_SIZE="$(select_fat_width "${IMAGE_SECTORS}" "${ROOT_ENTRIES}")"
      fi
   fi

   echo "Using IMAGE_SIZE of ${IMAGE_SIZE} with ${ROOT_ENTRIES} root directory entries"

   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 ! truncate --help > /dev/null 2> /dev/null ; then
       die "truncate 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 / Description
SECTOR_SIZE        -S
ROOT_DIR_ENTRIES   -r (calculated from the source files if unset)
FAT_SIZE           -F
FAT_OVERHEAD       Spare space to leave free in the filesystem, in KiB

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_SECTORS}" "${TMP_IMAGE}" "${ROOT_ENTRIES}"

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}"
