#!/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       Block device to flash or short-hands 'sd' or 'emmc'" >&2
  echo "  -h, --help   Display this help and exit." >&2
  echo "  -i, --image  Custom image (default: /boot/flash.bin). Implies --offline." >&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 "" >&2
}

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

OFFLINE=
FORCE=
IMAGE=
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
          ;;
        offline) OFFLINE=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

# 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

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 ! 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

  echo "Writing $IMAGE to $realdev" >&2

  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
  dd if="$IMAGE" of="$realdev" bs=512 seek="$((UBOOT_OFFSET / 512))" skip="$((FLASHBIN_OFFSET / 512))"
  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
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
