#!/bin/bash

reboot_pi () {
  umount /boot
  mount / -o remount,ro
  sync
  if [ "$NOOBS" = "1" ]; then
    if [ "$NEW_KERNEL" = "1" ]; then
      reboot -f "$BOOT_PART_NUM"
      sleep 5
    else
      echo "$BOOT_PART_NUM" > "/sys/module/${BCM_MODULE}/parameters/reboot_part"
    fi
  fi
  reboot -f
  sleep 5
  exit 0
}

check_noobs () {
  if [ "$BOOT_PART_NUM" = "1" ]; then
    NOOBS=0
  else
    NOOBS=1
  fi
}

get_variables () {
  ROOT_PART_DEV=$(findmnt / -o source -n)
  ROOT_PART_NAME=$(echo "$ROOT_PART_DEV" | cut -d "/" -f 3)
  ROOT_DEV_NAME=$(echo /sys/block/*/"${ROOT_PART_NAME}" | cut -d "/" -f 4)
  ROOT_DEV="/dev/${ROOT_DEV_NAME}"
  ROOT_PART_NUM=$(cat "/sys/block/${ROOT_DEV_NAME}/${ROOT_PART_NAME}/partition")

  BOOT_PART_DEV=$(findmnt /boot -o source -n)
  BOOT_PART_NAME=$(echo "$BOOT_PART_DEV" | cut -d "/" -f 3)
  BOOT_DEV_NAME=$(echo /sys/block/*/"${BOOT_PART_NAME}" | cut -d "/" -f 4)
  BOOT_PART_NUM=$(cat "/sys/block/${BOOT_DEV_NAME}/${BOOT_PART_NAME}/partition")

  OLD_DISKID=$(fdisk -l "$ROOT_DEV" | sed -n 's/Disk identifier: 0x\([^ ]*\)/\1/p')

  check_noobs

  ROOT_DEV_SIZE=$(cat "/sys/block/${ROOT_DEV_NAME}/size")
  TARGET_END=$((ROOT_DEV_SIZE - 1))

  PARTITION_TABLE=$(parted -m "$ROOT_DEV" unit s print | tr -d 's')

  LAST_PART_NUM=$(echo "$PARTITION_TABLE" | tail -n 1 | cut -d ":" -f 1)

  ROOT_PART_LINE=$(echo "$PARTITION_TABLE" | grep -e "^${ROOT_PART_NUM}:")
  ROOT_PART_START=$(echo "$ROOT_PART_LINE" | cut -d ":" -f 2)
  ROOT_PART_END=$(echo "$ROOT_PART_LINE" | cut -d ":" -f 3)

  if [ "$NOOBS" = "1" ]; then
    EXT_PART_LINE=$(echo "$PARTITION_TABLE" | grep ":::;" | head -n 1)
    EXT_PART_NUM=$(echo "$EXT_PART_LINE" | cut -d ":" -f 1)
    EXT_PART_START=$(echo "$EXT_PART_LINE" | cut -d ":" -f 2)
    EXT_PART_END=$(echo "$EXT_PART_LINE" | cut -d ":" -f 3)
  fi
}

fix_partuuid() {
  mount -o remount,rw "$ROOT_PART_DEV"
  mount -o remount,rw "$BOOT_PART_DEV"
  DISKID="$(dd if=/dev/hwrng bs=4 count=1 status=none | od -An -tx4 | cut -c2-9)"
  fdisk "$ROOT_DEV" > /dev/null <<EOF
x
i
0x$DISKID
r
w
EOF
  if [ "$?" -eq 0 ]; then
    sed -i "s/${OLD_DISKID}/${DISKID}/g" /etc/fstab
    sed -i "s/${OLD_DISKID}/${DISKID}/" /boot/cmdline.txt
    sync
  fi

  mount -o remount,ro "$ROOT_PART_DEV"
  mount -o remount,ro "$BOOT_PART_DEV"
}

fix_wpa() {
  if [ -e /boot/firstrun.sh ] \
     && ! grep -q 'imager_custom set_wlan' /boot/firstrun.sh \
     && grep -q wpa_supplicant.conf /boot/firstrun.sh; then
    mount -o remount,rw "$ROOT_PART_DEV"
    modprobe rfkill
    REGDOMAIN=$(sed -n 's/^\s*country=\(..\)$/\1/p' /boot/firstrun.sh)
    [ -n "$REGDOMAIN" ] && raspi-config nonint do_wifi_country "$REGDOMAIN"
    if systemctl -q is-enabled NetworkManager; then
      systemctl disable NetworkManager
    fi
    mount -o remount,ro "$ROOT_PART_DEV"
  fi
}

check_variables () {
  if [ "$NOOBS" = "1" ]; then
    if [ "$EXT_PART_NUM" -gt 4 ] || \
       [ "$EXT_PART_START" -gt "$ROOT_PART_START" ] || \
       [ "$EXT_PART_END" -lt "$ROOT_PART_END" ]; then
      FAIL_REASON="Unsupported extended partition\n$FAIL_REASON"
      return 1
    fi
  fi

  if [ "$ROOT_PART_NUM" -ne "$LAST_PART_NUM" ]; then
    # Skip resize if root partition is not last
    return 1
  fi

  if [ "$ROOT_PART_END" -gt "$TARGET_END" ]; then
    FAIL_REASON="Root partition runs past the end of device\n$FAIL_REASON"
    return 1
  fi

  if [ ! -b "$ROOT_DEV" ] || [ ! -b "$ROOT_PART_DEV" ] || [ ! -b "$BOOT_PART_DEV" ] ; then
    FAIL_REASON="Could not determine partitions\n$FAIL_REASON"
    return 1
  fi
  if [ "$ROOT_PART_END" -eq "$TARGET_END" ]; then
    # Root partition already the expected size
    return 1
  fi
}

check_kernel () {
  MAJOR="$(uname -r | cut -f1 -d.)"
  MINOR="$(uname -r | cut -f2 -d.)"
  if [ "$MAJOR" -eq "4" ] && [ "$MINOR" -lt "9" ]; then
    return 0
  fi
  if [ "$MAJOR" -lt "4" ]; then
    return 0
  fi
  NEW_KERNEL=1
}

do_resize () {
  check_kernel

  if [ "$NOOBS" = "1" ] && [ "$NEW_KERNEL" != "1" ]; then
    BCM_MODULE=$(grep -e "^Hardware" /proc/cpuinfo | cut -d ":" -f 2 | tr -d " " | tr '[:upper:]' '[:lower:]')
    if ! modprobe "$BCM_MODULE"; then
      FAIL_REASON="Couldn't load BCM module $BCM_MODULE\n$FAIL_REASON"
      return 1
    fi
  fi

  whiptail --infobox "Resizing root filesystem...\n\nDepending on storage size and speed, this may take a while." 20 60
  if [ "$NOOBS" = "1" ]; then
    if ! printf "resizepart %s\nyes\n%ss\n" "$EXT_PART_NUM" "$TARGET_END" | parted "$ROOT_DEV" ---pretend-input-tty; then
      FAIL_REASON="Extended partition resize failed\n$FAIL_REASON"
      return 1
    fi
  fi

  if ! parted -m "$ROOT_DEV" u s resizepart "$ROOT_PART_NUM" "$TARGET_END"; then
    FAIL_REASON="Partition table resize of the root partition ($ROOT_PART_DEV) failed\n$FAIL_REASON"
    return 1
  fi

  mount -o remount,rw /
  resize2fs "$ROOT_PART_DEV" > /dev/null 2>&1
  RET="$?"
  if [ "$RET" -ne 0 ]; then
    FAIL_REASON="Root partition resize failed\n$FAIL_REASON"
  fi

  mount -o remount,ro /
  return "$RET"
}

regenerate_ssh_host_keys () {
  mount -o remount,rw /
  /usr/lib/raspberrypi-sys-mods/regenerate_ssh_host_keys
  RET="$?"
  mount -o remount,ro /
  return "$RET"
}

apply_custom () {
  CONFIG_FILE="$1"
  mount -o remount,rw /
  mount -o remount,rw /boot
  if ! python3 -c "import toml" 2> /dev/null; then
    FAIL_REASON="custom.toml provided, but python3-toml is not installed\n$FAIL_REASON"
  else
    set -o pipefail
    /usr/lib/raspberrypi-sys-mods/init_config "$CONFIG_FILE" |& tee /run/firstboot.log | while read -r line; do
        MSG="$MSG\n$line"
        whiptail --infobox "$MSG" 20 60
    done
    if [ "$?" -ne 0 ]; then
      mv /run/firstboot.log /var/log/firstboot.log
      FAIL_REASON="Failed to apply customisations from custom.toml\n\nLog file saved as /var/log/firstboot.log\n$FAIL_REASON"
    fi
    set +o pipefail
  fi
  rm -f "$CONFIG_FILE"
  mount -o remount,ro /boot
  mount -o remount,ro /
}

main () {
  get_variables

  if check_variables; then
    do_resize
  fi

  # Switch to dhcpcd here if Imager < v1.7.3 was used to generate firstrun.sh
  fix_wpa > /dev/null 2>&1

  whiptail --infobox "Generating SSH keys..." 20 60
  regenerate_ssh_host_keys

  if [ -f "/boot/custom.toml" ]; then
    MSG="Applying customisations from custom.toml...\n"
    whiptail --infobox "$MSG" 20 60
    apply_custom "/boot/custom.toml"
  fi

  whiptail --infobox "Fix PARTUUID..." 20 60
  fix_partuuid

  return 0
}

mount -t proc proc /proc
mount -t sysfs sys /sys
mount -t tmpfs tmp /run
mkdir -p /run/systemd

mount /boot
mount / -o remount,ro

sed -i 's| init=/usr/lib/raspberrypi-sys-mods/firstboot||' /boot/cmdline.txt
sed -i 's| sdhci\.debug_quirks2=4||' /boot/cmdline.txt

if ! grep -q splash /boot/cmdline.txt; then
  sed -i "s/ quiet//g" /boot/cmdline.txt
fi
mount /boot -o remount,ro
sync

main

if [ -z "$FAIL_REASON" ]; then
  whiptail --infobox "Rebooting in 5 seconds..." 20 60
  sleep 5
else
  whiptail --msgbox "Failed running firstboot:\n${FAIL_REASON}" 20 60
fi

reboot_pi
