#!/bin/sh
# SPDX-License-Identifier: GPL-3.0+
# Copyright 2023-2025 Johannes Schauer Marin Rodrigues <josch@mister-muffin.de>

set -eu

usage() {
  echo "Usage: " >&2
  echo "  reform-flash-bootloader [--offline] [--force] [DEVICE...]" >&2
  echo >&2
  echo "Download and flash the bootloader for the current platform to DEVICE. Unless" >&2
  echo "the --offline option is given, download latest bootloader to /boot/flash.bin." >&2
  echo "If one or more DEVICE is given, flash /boot/flash.bin with the correct" >&2
  echo "offset to DEVICE. The short-hands 'sd' and 'emmc' can be used to flash" >&2
  echo "the bootloader to the SD-card or eMMC, respectively." >&2
  echo >&2
  echo "Options:" >&2
  echo "  DEVICE       One or more block device(s) to flash or short-hands 'sd' or 'emmc'" >&2
  echo "  -h, --help   Display this help and exit." >&2
  echo "  -i IMAGE, --image=IMAGE" >&2
  echo "               Custom image (default: /boot/flash.bin). Implies --offline and" >&2
  echo "               disables checksum verification." >&2
  echo "  --offline    Do not download latest bootloader to /boot/flash.bin." >&2
  echo "  -f, --force  No user interaction and flash to devices marked as 'warn'" >&2
  echo "  --zero       Fill the space where the bootloader would've been flashed to with zeroes" >&2
  echo "  --machine=MACHINE" >&2
  echo "               Force a different machine than the one this script is running on" >&2
  echo "               using the values from /proc/device-tree/model or the basenames of" >&2
  echo "               config files in /usr/share/reform-tools/machines/. This will change" >&2
  echo "               flashing offsets. Implies --offline. Requires a custom bootloader binary" >&2
  echo "               via --image. Device short-hands are disabled. Flashing to emmc is" >&2
  echo "               disabled." >&2
  echo "  --verbose    Print with higher verbosity." >&2
  echo "  --dry-run    Do not perform the actual flashing." >&2
  echo "" >&2
}

nth_arg() {
  shift "$1"
  printf "%s" "$1"
}

OFFLINE=
FORCE=
IMAGE=
ZERO=
MACHINE=
VERBOSE=
DRYRUN=
while getopts :hfi:-: OPTCHAR; do
  case "$OPTCHAR" in
    h)
      usage
      exit 0
      ;;
    f) FORCE=yes ;;
    i) IMAGE="$OPTARG" ;;
    -)
      case "$OPTARG" in
        help)
          usage
          exit 0
          ;;
        force) FORCE=yes ;;
        image)
          if [ "$OPTIND" -gt "$#" ]; then
            echo "E: missing argument for --image" >&2
            exit 1
          fi
          IMAGE="$(nth_arg "$OPTIND" "$@")"
          OPTIND=$((OPTIND + 1))
          OFFLINE=yes
          ;;
        image=*)
          IMAGE="${OPTARG#*=}"
          OFFLINE=yes
          ;;
        machine)
          if [ "$OPTIND" -gt "$#" ]; then
            echo "E: missing argument for --machine" >&2
            exit 1
          fi
          MACHINE="$(nth_arg "$OPTIND" "$@")"
          OPTIND=$((OPTIND + 1))
          OFFLINE=yes
          ;;
        machine=*)
          MACHINE="${OPTARG#*=}"
          OFFLINE=yes
          ;;
        offline) OFFLINE=yes ;;
        zero) ZERO=yes ;;
        verbose) VERBOSE=yes ;;
        dry-run) DRYRUN=yes ;;
        *)
          echo "E: unrecognized option: --$OPTARG" >&2
          exit 1
          ;;
      esac
      ;;
    ':')
      echo "E: missing argument for -$OPTARG" >&2
      exit 1
      ;;
    '?')
      echo "E: unrecognized option -$OPTARG" >&2
      exit 1
      ;;
    *)
      echo "E: error parsing options" >&2
      exit 1
      ;;
  esac
done
shift "$((OPTIND - 1))"

if [ "$(id -u)" -ne 0 ] && [ "$DRYRUN" != "yes" ]; then
  echo "reform-flash-bootloader has to be run as root / using sudo."
  exit 1
fi

if ! command -v parted >/dev/null && [ "$DRYRUN" = "yes" ]; then
  # parted was not found in $PATH
  # try adding /bin/sbin and /sbin to $PATH
  PATH="$PATH:/usr/sbin:/sbin"
  export PATH
fi

if ! command -v parted >/dev/null; then
  echo "E: unable to find parted utility" >&2
  exit 1
fi

# Even with --machine is used to override the machine config to load, we are
# loading the native machine config to be able to forbid flashing to eMMC.
# Since eMMC might be named differently in the config passed via --machine
# we need to load the native one for that information first.
# shellcheck source=/dev/null
if [ -e "./machines/$(cat /proc/device-tree/model).conf" ]; then
  . "./machines/$(cat /proc/device-tree/model).conf"
elif [ -e "/usr/share/reform-tools/machines/$(cat /proc/device-tree/model).conf" ]; then
  . "/usr/share/reform-tools/machines/$(cat /proc/device-tree/model).conf"
else
  echo "E: unable to find config for $(cat /proc/device-tree/model)" >&2
  exit 1
fi

if [ -n "$MACHINE" ]; then
  if [ -z "$IMAGE" ]; then
    echo "E: using the --machine option requires a custom image passed via the --image option" >&2
    exit 1
  fi
  for dev in "$@"; do
    case $dev in
      "/dev/${DEV_MMC}"*)
        echo "E: flashing to eMMC not supported with --machine" >&2
        exit 1
        ;;
      "sd" | "emmc")
        echo "E: short-hands not supported with --machine. Supply the device path explicitly." >&2
        exit 1
        ;;
    esac
  done
  case $MACHINE in
    *"/"*)
      echo "E: invalid machine name (contains a slash)" >&2
      exit 1
      ;;
    *.conf)
      echo "E: invalid machine name (ends with .conf)" >&2
      exit 1
      ;;
  esac
  # shellcheck source=/dev/null
  if [ -e "./machines/$MACHINE.conf" ]; then
    . "./machines/$MACHINE.conf"
  elif [ -e "/usr/share/reform-tools/machines/$MACHINE.conf" ]; then
    . "/usr/share/reform-tools/machines/$MACHINE.conf"
  else
    echo "E: unable to find config for $MACHINE" >&2
    exit 1
  fi
  # Make sure that the given device name can never be interpreted as meaning
  # the eMMC. This is because the other machine config might indicate a device
  # to be eMMC which it is not on this platform.
  DEV_MMC="i-do-not-exist"
fi

for dev in "$@"; do
  case $dev in
    emmc | "/dev/${DEV_MMC}"*)
      if [ "$EMMC_BOOT" = false ]; then
        echo "E: writing bootloader to eMMC not supported on $(cat /proc/device-tree/model)" >&2
        exit 1
      fi
      ;;
    sd | "/dev/${DEV_SD}"*)
      if [ "$SD_BOOT" = false ]; then
        printf "E: Refusing to write bootloader for this platform to SD-Card because this SoM is unable to load the bootloader from SD-Card: " >&2
        if [ -z "$MACHINE" ]; then
          printf "%s\n" "$(cat /proc/device-tree/model)" >&2
        else
          echo "$MACHINE" >&2
        fi
        exit 1
      fi
      ;;
  esac
done

if [ "$OFFLINE" != "yes" ]; then
  if echo "$BOOTLOADER_SHA1  /boot/flash.bin" | sha1sum --strict --check >/dev/null 2>&1; then
    echo "I: /boot/flash.bin is up-to-date and has expected sha1sum -- not downloading it again" >&2
  else
    echo "I: Downloading the bootloader to /boot/flash.bin and comparing checksum" >&2
    bootloaderurl="https://source.mnt.re/reform/${BOOTLOADER_PROJECT}/-/jobs/artifacts/${BOOTLOADER_TAG}/raw/$(basename "$DTBPATH" .dtb)-flash.bin?job=build"
    /usr/lib/apt/apt-helper -oAPT::Sandbox::User=root download-file "$bootloaderurl" "/boot/flash.bin" "SHA1:$BOOTLOADER_SHA1"
  fi
  # download mhdpfw.bin on ls1028a
  case "$(cat /proc/device-tree/model)" in "MNT Reform 2 with LS1028A Module")
    if echo "fa96b9aa59d7c1e9e6ee1c0375d0bcc8f8e5b78c  /boot/ls1028a-mhdpfw.bin"; then
      echo "I: /boot/ls1028a-mhdpfw.bin is up-to-date -- not downloading it again" >&2
    else
      echo "I: Downloading LS1028A MHDP firmware to /boot/ls1028a-mhdpfw.bin and comparing checksum" >&2
      /usr/lib/apt/apt-helper -oAPT::Sandbox::User=root download-file \
        "https://source.mnt.re/reform/reform-ls1028a-uboot/-/raw/main/ls1028a-mhdpfw.bin" \
        "/boot/ls1028a-mhdpfw.bin" \
        "SHA1:fa96b9aa59d7c1e9e6ee1c0375d0bcc8f8e5b78c"
    fi
    ;;
  esac
fi

if [ "$#" -eq 0 ]; then
  echo "I: No device unto which to flash the bootloader provided. Exiting." >&2
  exit 0
fi

if [ -z "$IMAGE" ]; then
  if [ ! -e /boot/flash.bin ]; then
    echo "E: /boot/flash.bin does not exist" >&2
    exit 1
  fi
  if ! echo "$BOOTLOADER_SHA1  /boot/flash.bin" | sha1sum --strict --check >/dev/null 2>&1; then
    echo "Incorrect checksum for /boot/flash.bin" >&2
    echo "Either flash a custom image with --image, or run without --offline to download the latest bootloader version" >&2
    exit 1
  fi
  IMAGE="/boot/flash.bin"
fi

if [ ! -e "$IMAGE" ]; then
  echo "E: $IMAGE does not exist" >&2
  exit 1
fi

bootloadersize=$(stat --format=%s "$IMAGE")

if [ "$((bootloadersize % 512))" -ne 0 ]; then
  echo "E: the size of the bootloader must be a multiple of 512 bytes" >&2
fi

# check if there is enough free space at the beginning of the disk
for dev in "$@"; do
  case $dev in
    emmc | "/dev/${DEV_MMC}"*)
      if [ "$DEV_MMC_BOOT0" = true ]; then
        # there are no partitions on boot0, so no need to check here
        continue
      else
        realdev=/dev/${DEV_MMC}
      fi
      ;;
    sd)
      realdev=/dev/${DEV_SD}
      ;;
    *) realdev="$dev" ;;
  esac

  if [ ! -e "$realdev" ]; then
    if [ "$DRYRUN" = "yes" ]; then
      echo "S: $realdev does not exist but that's okay because this is a simulation" >&2
      continue
    else
      echo "E: $realdev does not exist" >&2
      exit 1
    fi
  elif [ ! -b "$realdev" ]; then
    if [ "$DRYRUN" = "yes" ]; then
      echo "S: $realdev is not a block device but that's okay because this is a simulation" >&2
      continue
    else
      echo "E: $realdev is not a block device" >&2
      exit 1
    fi
  elif [ ! -r "$realdev" ]; then
    if [ "$DRYRUN" = "yes" ]; then
      echo "S: $realdev is not readable but that's okay because this is a simulation" >&2
      continue
    else
      echo "E: $realdev is not readable" >&2
      exit 1
    fi
  fi

  disk_label=$(parted --json --script "$realdev" unit B print 2>/dev/null | jq --raw-output '.disk.label')
  # no further tests for disks without a partition table
  case $disk_label in
    unknown)
      echo "I: No partition table found on $realdev" >&2
      continue
      ;;
    msdos)
      if [ "$BOOTLOADER_OFFSET" -lt 512 ]; then
        echo "E: The bootloader would be flashed with an offset of $BOOTLOADER_OFFSET which would overwrite parts of the" >&2
        echo "E: MBR partition table which requires 512 bytes of space at the beginning of $realdev" >&2
        exit 1
      fi
      ;;
    gpt)
      if [ "$BOOTLOADER_OFFSET" -lt 17408 ]; then
        echo "E: The bootloader would be flashed with an offset of $BOOTLOADER_OFFSET which would overwrite parts of the" >&2
        echo "E: GUID partition table which requires 17408 bytes of space at the beginning of $realdev" >&2
        exit 1
      fi
      ;;
  esac

  num_parts=$(parted --json --script "$realdev" unit B print | jq --raw-output '.disk.partitions | length')
  if [ "$num_parts" -eq 0 ]; then
    echo "I: No partition was found on $realdev" >&2
    continue
  fi

  firstpartstart=$(parted --json --script "$realdev" unit B print | jq --raw-output '.disk.partitions[0].start')
  # strip off trailing B
  firstpartstart=${firstpartstart%B}
  if [ "$((BOOTLOADER_OFFSET - FLASHBIN_OFFSET + bootloadersize))" -ge "$firstpartstart" ]; then
    echo "E: The first partition on $realdev starts at $firstpartstart and would be overwritten by the bootloader" >&2
    echo "E: make sure that the first $((BOOTLOADER_OFFSET - FLASHBIN_OFFSET + bootloadersize)) bytes are free on $realdev" >&2
    exit 1
  fi
done

# rk3588 and a311d u-boot version 2026-01-11 starts preferring NVMe over eMMC, so warn the
# user about it in case they have a bootflow on NVMe but want to use eMMC
# for their /boot instead
if [ "$ZERO" != "yes" ]; then
  case $(findmnt --noheadings --evaluate --mountpoint /boot --output SOURCE) in "/dev/${DEV_MMC}"*)
    case "$(cat /proc/device-tree/model)" in
      "MNT Pocket Reform with BPI-CM4 Module" | \
        "MNT Pocket Reform with RCORE RK3588 Module" | \
        "MNT Reform 2 with BPI-CM4 Module" | \
        "MNT Reform 2 with RCORE-DSI RK3588 Module" | \
        "MNT Reform 2 with RCORE RK3588 Module" | \
        "MNT Reform Next with RCORE RK3588 Module")
        if dpkg --compare-versions "$BOOTLOADER_TAG" "ge" "2026-01-11"; then
          echo "W: u-boot version 2026-01-11 for A311D and RK3588 moved the bootflow for NVMe to be preferred over eMMC." >&2
          echo "W: Make sure that you do not have a boot.scr or extlinux.conf on your NVMe if you were" >&2
          echo "W: using eMMC for your /boot before" >&2
          if [ "$FORCE" = "yes" ]; then
            echo "Proceeding without user interaction because of --force" >&2
            response="y"
          else
            printf "Are you sure you want to proceed? [y/N] "
            read -r response
          fi
          if [ "$response" != "y" ]; then
            echo "Exiting."
            exit
          fi
        fi
        ;;
    esac
    ;;
  esac
fi

if [ "$EMMC_BOOT" = warn ] && [ "$FORCE" != "yes" ]; then
  for dev in "$@"; do
    case $dev in
      emmc | "/dev/${DEV_MMC}"*)
        echo "W: Flashing the bootloader to eMMC on $(cat /proc/device-tree/model) is not without risk." >&2
        echo "W: If you flash the wrong bootloader or if the flashing process goes wrong, it is" >&2
        echo "W: possible to soft-brick your board. Restoring it might need some extra hardware." >&2
        echo "W: Please only proceed if you are sure that the benefits outweigh the risks for you." >&2
        printf "Are you sure you want to proceed? [y/N] "
        read -r response

        if [ "$response" != "y" ]; then
          echo "Exiting."
          exit
        fi

        break
        ;;
    esac
  done
fi

# do the flashing
for dev in "$@"; do
  case $dev in
    emmc | "/dev/${DEV_MMC}"*)
      if [ "$DEV_MMC_BOOT0" = true ]; then
        realdev=/dev/${DEV_MMC}boot0
      else
        realdev=/dev/${DEV_MMC}
      fi
      ;;
    sd)
      realdev=/dev/${DEV_SD}
      ;;
    *) realdev="$dev" ;;
  esac

  if [ "$DRYRUN" = "yes" ]; then
    if [ "$ZERO" != "yes" ]; then
      echo "I: Simulate writing $IMAGE to $realdev" >&2
    else
      echo "I: Simulate overwriting the bootloader on $realdev with zeroes" >&2
    fi
  else
    if [ "$ZERO" != "yes" ]; then
      echo "I: Writing $IMAGE to $realdev" >&2
    else
      echo "I: Overwriting the bootloader on $realdev with zeroes" >&2
    fi
  fi

  case $dev in
    emmc | "/dev/${DEV_MMC}"*)
      if [ "$DEV_MMC_BOOT0" = true ]; then
        echo 0 >"/sys/class/block/${DEV_MMC}boot0/force_ro"
      fi
      ;;
  esac
  set -- of="$realdev" bs=512 seek="$((BOOTLOADER_OFFSET / 512))" skip="$((FLASHBIN_OFFSET / 512))" conv=fdatasync
  if [ "$ZERO" != "yes" ]; then
    # write bootloader binary
    set -- if="$IMAGE" "$@"
  else
    # write zeroes instead of the bootloader image
    set -- if="/dev/zero" "$@" count="$(((bootloadersize - FLASHBIN_OFFSET) / 512))"
  fi
  if [ "$DRYRUN" = "yes" ]; then
    echo "S: Simulating: dd $*" >&2
  else
    if [ "$VERBOSE" = "yes" ]; then
      echo "D: Running: dd $*" >&2
    fi
    dd "$@"
  fi
  case $dev in
    emmc | "/dev/${DEV_MMC}"*)
      if [ "$DEV_MMC_BOOT0" = true ]; then
        echo 1 >"/sys/class/block/${DEV_MMC}boot0/force_ro"
      fi
      ;;
  esac

  if [ "$DRYRUN" = "yes" ]; then
    echo "S: Skipping verification because this is a simulation" >&2
    continue
  elif [ "$ZERO" != "yes" ]; then
    echo "I: Reading bootloader from $realdev and comparing it to $IMAGE" >&2
  else
    echo "I: Reading bootloader from $realdev and making sure it's all zeroed out" >&2
  fi

  # For the extra paranoid, read back what got written and compare.
  set -- --bytes="$((bootloadersize - FLASHBIN_OFFSET))"
  if [ "$ZERO" != "yes" ]; then
    set -- "$@" --ignore-initial="0:$FLASHBIN_OFFSET" - "$IMAGE"
  else
    set -- "$@" --bytes="$((bootloadersize - FLASHBIN_OFFSET))" - /dev/zero
  fi
  # We use "dd iflag=direct" (which uses O_DIRECT) in the hopes that we read
  # from the device directly instead of re-using existing caches.
  if dd if="$realdev" bs=512 skip="$((BOOTLOADER_OFFSET / 512))" count="$(((bootloadersize - FLASHBIN_OFFSET) / 512))" iflag=direct \
    | cmp "$@"; then
    echo "I: Successfully read the same bytes from $realdev as were written before" >&2
  else
    echo "E: Did not get the same bytes back which were written"
    exit 1
  fi

  # compare checksum if no custom image was provided and if not just zeroes
  # were written
  if [ "$IMAGE" = "/boot/flash.bin" ] && [ "$ZERO" != "yes" ]; then
    echo "I: Reading bootloader from $realdev and comparing sha1sum" >&2
    trap 'rm -f $tmpbootloader' EXIT INT TERM
    tmpbootloader="$(mktemp)"
    dd if="$realdev" of="$tmpbootloader" bs=512 seek=$((FLASHBIN_OFFSET / 512)) skip="$((BOOTLOADER_OFFSET / 512))" count="$(((bootloadersize - FLASHBIN_OFFSET) / 512))"
    if echo "$BOOTLOADER_SHA1  $tmpbootloader" | sha1sum --strict --check >/dev/null 2>&1; then
      echo "I: sha1sum check successful" >&2
    else
      echo "E: sha1sum hash sum mismatch: bootloader on $realdev has an unexpected checksum" >&2
      exit 1
    fi
    rm "$tmpbootloader"
    trap - EXIT INT TERM
  fi
done

# inform about the DIP switch position only on imx8mq
if [ -z "$MACHINE" ]; then
  case "$(cat /proc/device-tree/model)" in "MNT Reform 2" | "MNT Reform 2 HDMI")
    for dev in "$@"; do
      case $dev in
        emmc | "/dev/${DEV_MMC}"*)
          echo "I: For the i.MX8MQ to load u-boot from MMC, make sure" >&2
          echo "I: that your DIP switch is set to OFF." >&2
          continue
          ;;
        sd | "/dev/${DEV_SD}"*)
          echo "I: For the i.MX8MQ to load u-boot from SD-Card, make sure" >&2
          echo "I: that your DIP switch is set to ON." >&2
          continue
          ;;
      esac
    done
    ;;
  esac
fi