#!/bin/bash
SCRIPT_NAME=rpi-keyboard-fw-update
FIRMWARE_DIR=/usr/lib/firmware/raspberrypi/keyboard
FIRMWARE_PI500=pi500-keyboard-fw.uf2
FIRMWARE_PI500PLUS_ISO=pi500plus-keyboard-fw-iso.uf2
FIRMWARE_PI500PLUS_ANSI=pi500plus-keyboard-fw-ansi.uf2
FIRMWARE_PI500PLUS_JIS=pi500plus-keyboard-fw-jis.uf2
FIRMWARE_UF2=unknown
MODEL_STRING="Raspberry Pi 500"
USB_VID_RASPBERRYPI=2e8a
USB_PID_RP2040_BOOTSEL=0003
USB_PID_PI500_KEYBOARD=0010
USB_PID_PI500PLUS_KEYBOARD=0011
RP2040_BOOTSEL_DRIVENAME=RPI-RP2
KEYB_BOOTSEL=118
KEYB_RUN=121
DT_ROOT=/proc/device-tree
UDEV_FILE=/etc/udev/rules.d/80-temporarily-ignore-rp2040.rules
SILENT=0
DEBUG=${DEBUG:-0}
COMPARE_VERSIONS=1
VERSION_CHECK_ONLY=0
WIPE_FIRST=0

die() {
  if [[ $SILENT -eq 0 ]]; then
    echo -e "$@" >&2
  fi
  exit 1
}

debug() {
  if [[ $DEBUG -eq 1 ]]; then
    echo "$@" >&2
  fi
}

cleanup() {
  # Remove "ignore RP2040 BOOTSEL" udev rule
  if [[ -f "$UDEV_FILE" ]]; then
    rm -f "$UDEV_FILE"
    udevadm control --reload-rules
  fi
}
trap cleanup EXIT

usage() {
  cat <<EOF
sudo $SCRIPT_NAME [options]

Update the firmware in a Pi 500 or Pi 500+ keyboard.

Options:
   -f <FILE> Use the specified UF2 file (skips model auto-detection).
   -h Display this usage message.
   -i Ignore version checks and always update firmware.
   -s Silent mode.
   -v Display keyboard type and firmware version (don't update anything).
   -w Wipe entire flash before updating firmware.
EOF
  exit 0
}

force_reset () {
  RUN_GPIO=$1
  pinctrl set $RUN_GPIO op pn dl
  pinctrl set $RUN_GPIO dh
}

force_bootsel_mode () {
  RUN_GPIO=$1
  BOOTSEL_GPIO=$2
  pinctrl set $RUN_GPIO op pn dl
  pinctrl set $BOOTSEL_GPIO op pn dh
  pinctrl set $RUN_GPIO dh
  pinctrl set $BOOTSEL_GPIO dl
}

wait_for_bootsel_drive () {
  debug "Waiting for BOOTSEL drive"
  # Wait up to 15 seconds for a new BOOTSEL drive to appear
  LOOP_COUNT=30
  FOUND_BOOTSEL_DRIVE=0
  # Store the list of xxxx1 partitions
  OLD_FIRST_PARTS=()
  while read -r PART; do
    OLD_FIRST_PARTS+=("$PART")
  done <<< "$(lsblk -n -o PATH | grep '1$')"
  while [[ $LOOP_COUNT -gt 0 ]]; do
    sleep 0.5
    LOOP_COUNT=$((LOOP_COUNT - 1))
    ALL_FIRST_PARTS=()
    while read -r PART; do
      ALL_FIRST_PARTS+=("$PART")
    done <<< "$(lsblk -n -o PATH | grep '1$')"
    if [[ ${#ALL_FIRST_PARTS[@]} -ne ${#OLD_FIRST_PARTS[@]} ]]; then
      # Wait for udevadm to settle
      udevadm settle
      # Look for any xxxx1 partitions which weren't in OLD_FIRST_PARTS
      NEW_FIRST_PARTS=()
      for PART in "${ALL_FIRST_PARTS[@]}"; do
        FOUND=0
        for OLD_PART in "${OLD_FIRST_PARTS[@]}"; do
          if [[ "$PART" == "$OLD_PART" ]]; then
            FOUND=1
            break
          fi
        done
        if [[ $FOUND -eq 0 ]]; then
          NEW_FIRST_PARTS+=("$PART")
        fi
      done
      # Update the list of xxxx1 partitions
      OLD_FIRST_PARTS=("${ALL_FIRST_PARTS[@]}")
      for DRIVE in "${NEW_FIRST_PARTS[@]}"; do
        debug "Candidate drive: $DRIVE"
        # Check that the new drive is an RP2040 in BOOTSEL mode
        DRIVE_INFO="$(udevadm info "$DRIVE")"
        if echo "$DRIVE_INFO" | grep -q "ID_USB_VENDOR_ID=$USB_VID_RASPBERRYPI" && echo "$DRIVE_INFO" | grep -q "ID_USB_MODEL_ID=$USB_PID_RP2040_BOOTSEL"; then
          debug "Found BOOTSEL drive $DRIVE"
          FOUND_BOOTSEL_DRIVE=1
          echo "$DRIVE"
          break
        else
          debug "$DRIVE isn't an RP2040 in BOOTSEL mode"
        fi
      done
      if [[ $FOUND_BOOTSEL_DRIVE -eq 1 ]]; then
        break
      fi
    fi
  done
  if [[ $LOOP_COUNT -eq 0 ]]; then
    die "Timed out waiting for BOOTSEL drive to appear"
  fi
}

copy_uf2_to_drive() {
  UF2_FILE=$1
  DRIVE=$2
  MOUNTPATH=$(mktemp -d)
  mount "$DRIVE" "$MOUNTPATH"
  cp "$UF2_FILE" "$MOUNTPATH"
  sync
  umount "$MOUNTPATH"
  rmdir "$MOUNTPATH"
}

while getopts f:hisvw option; do
   case "${option}" in
   f) FIRMWARE_UF2="${OPTARG}"
      if [[ "${FIRMWARE_UF2##*.}" != "uf2" ]]; then
        die "$FIRMWARE_UF2 must be a uf2 file"
      fi
      ;;
   h) usage
      ;;
   i) COMPARE_VERSIONS=0
      ;;
   s) SILENT=1
      ;;
   v) VERSION_CHECK_ONLY=1
      DEBUG=1
      ;;
   w) WIPE_FIRST=1
      ;;
   *) echo "Unknown argument \"${option}\""
      usage
      ;;
   esac
done

# Detect Pi model and variant
DT_MODEL=$(tr -d '\000' < "$DT_ROOT/model")
debug "DT_MODEL: $DT_MODEL"
if [[ "$DT_MODEL" != "$MODEL_STRING"* ]]; then
  die "This script can only run on a $MODEL_STRING"
else
  DT_BOARDREV_EXT=$(od -v -An -t x1 "$DT_ROOT/chosen/rpi-boardrev-ext" | tr -d ' \n')
  if [[ $(((0x$DT_BOARDREV_EXT >> 29) & 1)) -eq 1 ]]; then
    MODEL_VARIANT=pi500plus
  else
    MODEL_VARIANT=pi500
  fi
fi
debug "MODEL_VARIANT: $MODEL_VARIANT"

if [[ "$FIRMWARE_UF2" == "unknown" ]]; then
  # Attempt to auto-detect which UF2 file is needed
  # Detect Pi's country-code, and convert this to a keyboard layout
  DT_COUNTRY_CODE=$(od -v -An -j3 -N1 -i "$DT_ROOT/chosen/rpi-country-code" | tr -d ' \n')
  # see https://github.com/raspberrypi-ui/piwiz/blob/master/src/piwiz.c#L278 for reference 
  debug "DT_COUNTRY_CODE: $DT_COUNTRY_CODE"
  case $DT_COUNTRY_CODE in 
    4) # US
      CC_LAYOUT=ANSI
      ;;
    7) # JP
      CC_LAYOUT=JIS
      ;;
    14) # IL
      CC_LAYOUT=ANSI
      ;;
    16) # KR
      CC_LAYOUT=ANSI
      ;;
    *) # Everything else
      CC_LAYOUT=ISO
      ;;
  esac
  debug "CC_LAYOUT: $CC_LAYOUT"
  # Detect which USB keyboard is connected
  DETECTED_PI500_KEYBOARD=0
  DETECTED_PI500PLUS_KEYBOARD=0
  DETECTED_PI500PLUS_KEYBOARD_VARIANT=unknown
  DETECTED_KEYBOARD_FIRMWARE_VERSION=unknown
  if lsusb -d "$USB_VID_RASPBERRYPI:$USB_PID_PI500_KEYBOARD" &> /dev/null; then
    DETECTED_PI500_KEYBOARD=1
    DETECTED_KEYBOARD_FIRMWARE_VERSION=$(lsusb -v -d "$USB_VID_RASPBERRYPI:$USB_PID_PI500_KEYBOARD" 2> /dev/null | grep bcdDevice | tr -s " " | cut -d" " -f3)
  fi
  debug "DETECTED_PI500_KEYBOARD: $DETECTED_PI500_KEYBOARD"
  if lsusb -d "$USB_VID_RASPBERRYPI:$USB_PID_PI500PLUS_KEYBOARD" &> /dev/null; then
    DETECTED_PI500PLUS_KEYBOARD=1
    DETECTED_KEYBOARD_FIRMWARE_VERSION=$(lsusb -v -d "$USB_VID_RASPBERRYPI:$USB_PID_PI500PLUS_KEYBOARD" 2> /dev/null | grep bcdDevice | tr -s " " | cut -d" " -f3)
    KEYBOARD_PRODUCT=$(lsusb -v -d "$USB_VID_RASPBERRYPI:$USB_PID_PI500PLUS_KEYBOARD" 2> /dev/null | grep iProduct | tr -s " " | cut -d" " -f4-)
    for VARIANT in ISO ANSI JIS; do
      if [[ "$KEYBOARD_PRODUCT" == *"($VARIANT)" ]]; then
        DETECTED_PI500PLUS_KEYBOARD_VARIANT=$VARIANT
        break
      fi
    done
    if [[ "$DETECTED_PI500PLUS_KEYBOARD_VARIANT" == "unknown" ]]; then
      die "Unexpected Pi 500+ keyboard variant: $KEYBOARD_PRODUCT"
    fi
  fi
  debug "DETECTED_PI500PLUS_KEYBOARD: $DETECTED_PI500PLUS_KEYBOARD"
  debug "DETECTED_PI500PLUS_KEYBOARD_VARIANT: $DETECTED_PI500PLUS_KEYBOARD_VARIANT"
  debug "DETECTED_KEYBOARD_FIRMWARE_VERSION: $DETECTED_KEYBOARD_FIRMWARE_VERSION"
  if [[ "$MODEL_VARIANT" == "pi500" ]] && [[ $DETECTED_PI500PLUS_KEYBOARD -eq 1 ]]; then
    die "Pi firmware thinks this is a Pi 500, but the keyboard thinks it's a Pi 500+\nPlease use the -f option to manually select a UF2 file"
  fi
  if [[ "$MODEL_VARIANT" == "pi500plus" ]] && [[ $DETECTED_PI500_KEYBOARD -eq 1 ]]; then
    die "Pi firmware thinks this is a Pi 500+, but the keyboard thinks it's a Pi 500\nPlease use the -f option to manually select a UF2 file"
  fi
  if [[ "$MODEL_VARIANT" == "pi500plus" ]] && [[ $DETECTED_PI500PLUS_KEYBOARD -eq 1 ]]; then
    if [[ "$CC_LAYOUT" != "$DETECTED_PI500PLUS_KEYBOARD_VARIANT" ]]; then
      die "Pi firmware thinks this is a Pi 500+ with a $CC_LAYOUT keyboard, but the keyboard thinks it has a $DETECTED_PI500PLUS_KEYBOARD_VARIANT layout\nPlease use the -f option to manually select a UF2 file"
    fi
  fi
  if [[ "$MODEL_VARIANT" == "pi500" ]]; then
    AVAILABLE_FIRMWARE_FILE=$FIRMWARE_PI500
  elif [[ "$MODEL_VARIANT" == "pi500plus" ]]; then
    if [[ $CC_LAYOUT == "ISO" ]]; then
      AVAILABLE_FIRMWARE_FILE=$FIRMWARE_PI500PLUS_ISO
    elif [[ $CC_LAYOUT == "ANSI" ]]; then
      AVAILABLE_FIRMWARE_FILE=$FIRMWARE_PI500PLUS_ANSI
    elif [[ $CC_LAYOUT == "JIS" ]]; then
      AVAILABLE_FIRMWARE_FILE=$FIRMWARE_PI500PLUS_JIS
    else
      die "Unexpected keyboard layout: $CC_LAYOUT"
    fi
  fi
  AVAILABLE_FIRMWARE_FILE_VERSION=unknown
  if [[ -f "$FIRMWARE_DIR/versions.txt" ]]; then
    if AVAILABLE_FIRMWARE_VERSION_LINE=$(grep "$AVAILABLE_FIRMWARE_FILE" "$FIRMWARE_DIR/versions.txt"); then
      AVAILABLE_FIRMWARE_FILE_VERSION=$(echo "$AVAILABLE_FIRMWARE_VERSION_LINE" | cut -d"=" -f2)
    fi
  fi
  debug "AVAILABLE_FIRMWARE_FILE_VERSION: $AVAILABLE_FIRMWARE_FILE_VERSION"
  if [[ $VERSION_CHECK_ONLY -eq 1 ]]; then
    exit 0
  fi
  if [[ $COMPARE_VERSIONS -eq 1 ]] && [[ "$DETECTED_KEYBOARD_FIRMWARE_VERSION" != "unknown" ]] && [[ "$AVAILABLE_FIRMWARE_FILE_VERSION" != "unknown" ]] && [[ "$DETECTED_KEYBOARD_FIRMWARE_VERSION" == "$AVAILABLE_FIRMWARE_FILE_VERSION" ]]; then
    if [[ $SILENT -eq 0 ]]; then
      echo "Your keyboard firmware is already up to date"
    fi
    exit 0
  fi
  FIRMWARE_UF2="$FIRMWARE_DIR/$AVAILABLE_FIRMWARE_FILE"
fi

debug "FIRMWARE_UF2: $FIRMWARE_UF2"
if [[ ! -f "$FIRMWARE_UF2" ]]; then
  die "Required firmware file $FIRMWARE_UF2 is missing"
else
  # Everything before this line is just gathering information, and runs fine without root access
  if [[ $(id -u) -ne 0 ]]; then
    die "This script must be run as root. Try 'sudo $SCRIPT_NAME'"
  fi
  # Add "ignore RP2040 BOOTSEL" udev rule to stop drive-automounting and prevent GUI popups
  echo "ENV{ID_FS_TYPE}==\"vfat\", ENV{ID_FS_LABEL}==\"$RP2040_BOOTSEL_DRIVENAME\", ENV{UDISKS_IGNORE}=\"1\"" > "$UDEV_FILE"
  udevadm control --reload-rules
  # Trigger keyboard's BOOTSEL mode
  force_bootsel_mode $KEYB_RUN $KEYB_BOOTSEL
  # Wait for BOOTSEL drive to appear
  NEW_DRIVE=$(wait_for_bootsel_drive)
  if [[ -z "$NEW_DRIVE" ]]; then
    die "No BOOTSEL drive appeared"
  fi
  if [[ $WIPE_FIRST -eq 1 ]]; then
    # Use flash_nuke to wipe flash
    debug "Erasing flash"
    copy_uf2_to_drive "$FIRMWARE_DIR/flash_nuke.uf2" "$NEW_DRIVE"
    sleep 0.5
    # Wait for BOOTSEL drive to re-appear (which indicates the flash has been fully erased)
    NEW_DRIVE=$(wait_for_bootsel_drive)
    if [[ -z "$NEW_DRIVE" ]]; then
      die "No BOOTSEL drive appeared after erasing flash"
    fi
  fi
  # Copy new firmware to the drive
  debug "Writing new firmware"
  copy_uf2_to_drive "$FIRMWARE_UF2" "$NEW_DRIVE"
  # Remove "ignore RP2040 BOOTSEL" udev rule
  cleanup
  if [[ $SILENT -eq 0 ]]; then
    echo "Your keyboard firmware has been updated with $FIRMWARE_UF2"
  fi
fi
