#!/bin/bash

working_directory="$1"
firmware_directory="$2"
boot_link="$3"

# Check working directory existence.
if [ ! -n "$working_directory" ] || [ ! -d "$working_directory" ] || [ ! -e "$working_directory/boot" ] ; then
  echo "Invalid working directory: $working_directory"
  exit 1
fi

# Check firmware directory existence.
if [ ! -n "$firmware_directory" ] || [ ! -d "$firmware_directory" ] ; then
  echo "Invalid firmware directory: $firmware_directory"
  exit 1
fi

# Check boot link existence.
if [ ! -L "$boot_link" ] || [ ! -d "$boot_link" ] ; then
  echo "Invalid boot link: $boot_link"
  exit 1
fi

# Check boot link to point into firmware directory.
active_firmware=`readlink -f "$boot_link"`
if [ `dirname "$active_firmware"` != `readlink -f "$firmware_directory"` ] ; then
  echo "Invalid boot link destination: $boot_link"
  exit 1
fi

# Check boot loader versions
if [ ! -L "$boot_link/u-boot.imx" ] ; then
  echo "Cannot determine active boot loader."
  exit 1
fi
active_bootloader=`readlink "$boot_link/u-boot.imx"`
if [ ! -L "$working_directory/boot/u-boot.imx" ] ; then
  echo "Cannot determine new boot loader."
  exit 1
fi
new_firmware_bootloader=`readlink "$working_directory/boot/u-boot.imx"`

# Check root device
if [ ! -L "/dev/root" ] || [ ! -b "/dev/root" ] ; then
  echo "Cannot determine system root device."
  exit 1
fi

# Check boot partition
root_device=`readlink "/dev/root"`
boot_device=${root_device%p[0-9]}
active_boot_partition=`mmc extcsd read "/dev/$boot_device" | grep "Boot Partition [12] enabled" | sed "s/ *Boot Partition \([12]\) enabled.*/\1/"`
if [ ! -b "/dev/${boot_device}boot$(($active_boot_partition-1))" ] ; then
  echo "Cannot determine active system boot partition."
  exit 1
fi
new_firmware_boot_partition="${boot_device}boot$((1-($active_boot_partition-1)))"
if [ ! -b "/dev/$new_firmware_boot_partition" ] ; then
  echo "Cannot determine new system boot partition."
  exit 1
fi

# Determine new firmware instance directory.
active_firmware_instance=`basename "$active_firmware"`
new_firmware_instance=$(($active_firmware_instance+1))
while [ -e "$firmware_directory/$new_firmware_instance" ]; do
  new_firmware_instance=$(($new_firmware_instance+1))
done
new_firmware_instance_directory="$firmware_directory/$new_firmware_instance"

# Move working directory into new firmware directory location.
mv `readlink -f "$working_directory/boot"` "$new_firmware_instance_directory"
if [ $? -ne 0 ] ; then
  echo "Failed to move working directory into new firmware directory location."
  exit 1
fi

# Flush all data to flash.
sync

# Update boot loader if needed.
if [ "$active_bootloader" != "$new_firmware_bootloader" ] ; then
  echo 0 > /sys/block/$new_firmware_boot_partition/force_ro
  if ! dd if=$new_firmware_instance_directory/$new_firmware_bootloader of=/dev/$new_firmware_boot_partition bs=512 seek=2 2>/dev/null ; then
    echo 1 > /sys/block/$new_firmware_boot_partition/force_ro
    echo "Failed to update boot loader."
    exit 1
  fi
  echo 1 > /sys/block/$new_firmware_boot_partition/force_ro
  i=0
  while true; do
    # Sync to ensure completion of all pending command activity to eMMC device (otherwise switching boot partition may silently fail).
    sync
    if ! mmc bootpart enable $((1-($active_boot_partition-1)+1)) 1 "/dev/$boot_device" ; then
      echo "Failed to switch boot partition."
      exit 1
    fi
    sync
    # Check that switching boot partition succeeded.
    new_boot_partition=`mmc extcsd read "/dev/$boot_device" | grep "Boot Partition [12] enabled" | sed "s/ *Boot Partition \([12]\) enabled.*/\1/"`
    if [ "$new_boot_partition" -eq "$active_boot_partition" ] ; then
      # Retry switching boot partition up to 5 times since it occasionally fails silently.
      if [ $i -lt 5 ]; then
        i=$(($i+1))
        continue
      fi
      echo "Could not switch boot partition."
      exit 1
    fi
    break
  done
fi

# Update boot link.
active_firmware=`readlink "$boot_link"`
active_firmware_base=`dirname "$active_firmware"`
new_firmware="$active_firmware_base/$new_firmware_instance"
ln -sfn "$new_firmware" "$boot_link"
if [ $? -ne 0 ] ; then
  echo "Failed update boot link."
  exit 1
else
  rm -f /var/boot/.bootmsg*
fi
