#!/bin/sh

set -eu

usage() {
  echo "Download latest u-boot for the current platform and update /boot/flash.bin" >&2
  echo "Then flash u-boot to DEVICE which can be one of the short-hands sd or emmc." >&2
  echo >&2
  echo "Usage: $0 [--help] [--offline] DEVICE" >&2
  echo >&2
  echo "Options:" >&2
  echo '  DEVICE           Target device to flash: either "sd" or "emmc"' >&2
  echo "  --help           Display this help and exit." >&2
  echo "  --offline        Instead of download u-boot use /boot/flash.bin." >&2
}

if [ "$#" -gt 0 ] && [ "$1" = "--help" ]; then
  usage
  exit 0
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

OFFLINE=
FORCE=
for dev in "$@"; do
  case $dev in
    --offline) OFFLINE=yes ;;
    --force) FORCE=yes ;;
    emmc)
      if [ "$EMMC_BOOT" = false ]; then
        echo "E: writing uboot to eMMC not supported on $(cat /proc/device-tree/model)" >&2
        exit 1
      fi
      ;;
    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
      ;;
    *)
      echo "E: unknown device target \"$dev\" given." >&2
      exit 1
      ;;
  esac
done

usage() {
  echo "Usage: " >&2
  echo "  reform-flash-uboot [--offline] [--force] [device...]" >&2
  echo "" >&2
  if [ "$EMMC_BOOT" != false ] && [ "$SD_BOOT" = true ]; then
    echo "Download and write recent uboot to eMMC or SD-Card." >&2
  elif [ "$SD_BOOT" = true ]; then
    echo "Download and write recent uboot to SD-Card." >&2
  elif [ "$EMMC_BOOT" != false ]; then
    echo "Download and write recent uboot to eMMC." >&2
  else
    echo "Download recent uboot, but update no devices." >&2
  fi
  echo "Without arguments only /boot/flash.bin will be updated unless --offline is given." >&2
  echo "Otherwise, the arguments are the devices onto which to install uboot." >&2
  echo "These device names are supported:" >&2
  echo "" >&2
  if [ "$EMMC_BOOT" != false ]; then
    if [ "$DEV_MMC_BOOT0" = true ]; then
      echo "   emmc (${DEV_MMC}boot0)" >&2
    else
      echo "   emmc (${DEV_MMC})" >&2
    fi
  fi
  if [ "$SD_BOOT" = true ]; then
    echo "   sd (${DEV_SD})" >&2
  fi
  echo "" >&2
  echo "The --force allows flashing to targets marked as 'warn' without " >&2
  echo "interactively asking the user." >&2
  echo "" >&2
}

if [ "$#" -gt 0 ] && [ "$1" = "--help" ] || [ "$#" -gt 0 ] && [ "$1" = "-h" ]; then
  usage
  exit 0
fi

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

if [ "$#" -eq 0 ] || [ "$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
elif [ "$#" -gt 0 ] && [ "$OFFLINE" = "yes" ]; then
  echo "Not downloading uboot but using /boot/flash.bin without updating it." >&2
  shift
  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 manually, or run without --offline to download the latest uboot version" >&2
    exit 1
  fi
fi

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

ubootsize=$(stat --format=%s "/boot/flash.bin")

# check if there is enough free space at the beginning of the disk
for dev in "$@"; do
  case $dev in
    emmc)
      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}
      ;;
    *) continue ;;
  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 $dev $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
    if [ "$dev" = "emmc" ]; then
      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
    fi
  done
fi

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

  echo "Writing /boot/flash.bin to $dev ($realdev)" >&2

  if [ "$DEV_MMC_BOOT0" = true ] && [ "$dev" = "emmc" ]; then
    echo 0 >"/sys/class/block/${DEV_MMC}boot0/force_ro"
  fi
  dd if=/boot/flash.bin of="$realdev" bs=512 seek="$((UBOOT_OFFSET / 512))" skip="$((FLASHBIN_OFFSET / 512))"
  if [ "$DEV_MMC_BOOT0" = true ] && [ "$dev" = "emmc" ]; then
    echo 1 >"/sys/class/block/${DEV_MMC}boot0/force_ro"
  fi

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