#!/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-uboot [--offline] [--force] [DEVICE...]" >&2
  echo >&2
  echo "Download and flash u-boot for the current platform to DEVICE. Unless" >&2
  echo "the --offline option is given, download latest u-boot 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 "u-boot 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 u-boot 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 u-boot 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 u-boot binary" >&2
  echo "               via --image. Device short-hands are disabled. Flashing to emmc is" >&2
  echo "               disabled." >&2
  echo "" >&2
}

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

OFFLINE=
FORCE=
IMAGE=
ZERO=
MACHINE=
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 ;;
        *)
          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 ]; then
  echo "reform-flash-uboot has to be run as root / using sudo."
  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
      ;;
    *.*)
      echo "E: invalid machine name (contains a dot)" >&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 uboot to eMMC not supported on $(cat /proc/device-tree/model)" >&2
        exit 1
      fi
      ;;
    sd | "/dev/${DEV_SD}"*)
      if [ "$SD_BOOT" = false ]; then
        echo "E: writing uboot to SD-Card not supported on $(cat /proc/device-tree/model)" >&2
        exit 1
      fi
      ;;
  esac
done

if [ "$OFFLINE" != "yes" ]; then
  if echo "$UBOOT_SHA1  /boot/flash.bin" | sha1sum --strict --check >/dev/null 2>&1; then
    echo "/boot/flash.bin is up-to-date -- not downloading it again" >&2
  else
    echo "Downloading uboot to /boot/flash.bin and comparing checksum" >&2
    ubooturl="https://source.mnt.re/reform/${UBOOT_PROJECT}/-/jobs/artifacts/${UBOOT_TAG}/raw/$(basename "$DTBPATH" .dtb)-flash.bin?job=build"
    /usr/lib/apt/apt-helper -oAPT::Sandbox::User=root download-file "$ubooturl" "/boot/flash.bin" "SHA1:$UBOOT_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 "/boot/ls1028a-mhdpfw.bin is up-to-date -- not downloading it again" >&2
    else
      echo "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 "No device unto which to flash uboot 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 "$UBOOT_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 uboot version" >&2
    exit 1
  fi
  IMAGE="/boot/flash.bin"
fi

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

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

# 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

  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
  if [ "$disk_label" = "unknown" ]; then
    echo "I: no partition table found on $realdev" >&2
    continue
  fi

  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 [ "$((UBOOT_OFFSET - FLASHBIN_OFFSET + ubootsize))" -ge "$firstpartstart" ]; then
    echo "the first partition on $realdev starts at $firstpartstart and would be overwritten by uboot" >&2
    echo "make sure that the first $((UBOOT_OFFSET - FLASHBIN_OFFSET + ubootsize)) bytes are free on $realdev" >&2
    exit 1
  fi
done

if [ "$EMMC_BOOT" = warn ] && [ "$FORCE" != "yes" ]; then
  for dev in "$@"; do
    case $dev in
      emmc | "/dev/${DEV_MMC}"*)
        echo "W: Flashing u-boot to eMMC on $(cat /proc/device-tree/model) is not without risk." >&2
        echo "W: If you flash the wrong u-boot 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 [ "$ZERO" != "yes" ]; then
    echo "I: Writing $IMAGE to $realdev" >&2
  else
    echo "I: Overwriting u-boot on $realdev with zeroes" >&2
  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
  if [ "$ZERO" != "yes" ]; then
    # write u-boot binary
    dd if="$IMAGE" of="$realdev" bs=512 seek="$((UBOOT_OFFSET / 512))" skip="$((FLASHBIN_OFFSET / 512))"
  else
    # write zeroes instead of the u-boot image
    dd if="/dev/zero" of="$realdev" bs=512 seek="$((UBOOT_OFFSET / 512))" count="$(((ubootsize - FLASHBIN_OFFSET) / 512))"
  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

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 "For the i.MX8MQ to load u-boot from MMC, make sure" >&2
          echo "that your DIP switch is set to OFF." >&2
          continue
          ;;
        sd | "/dev/${DEV_SD}"*)
          echo "For the i.MX8MQ to load u-boot from SD-Card, make sure" >&2
          echo "that your DIP switch is set to ON." >&2
          continue
          ;;
      esac
    done
    ;;
  esac
fi
