#! /bin/bash

VESION="v5.0"

set -e

die ()
{
    echo "ERROR: $1"
    exit 1
}

usage ()
{
  echo "
Usage: `basename $1` --machine [us01|us02|us03|ns01|ns02] --device <device>

Mandatory options:
  --device				SD block device node (e.g /dev/sdd)
  --machine				Machine [us01,us02,us03,ns01,ns02]

Optional options:
  --rootfs <path/to/rootfs.tar.gz>	Use provided rootfs tarball
  --kernel <path/to/kernel.tar.gz>	Use provided kernel tarball
  --uboot <path/to/u-boot.tar.gz>	Use provided u-boot tarball
  --xloader <path/to/xloader.tar.gz>	Use provided u-boot tarball
  --help				Print this help message.
"
}

BASE_DIR=$(dirname $0)

# Process command line...
while [ $# -gt 0 ]; do
  case $1 in
    --help | -h)
      usage $0
      exit 0
      ;;
    --device) shift; device=$1; shift; ;;
    --machine) shift; machine=$1; shift; ;;
    --kernel) shift; kernel=$1; shift; ;;
    --uboot) shift; uboot=$1; shift; ;;
    --rootfs) shift; rootfs=$1; shift; ;;
    --xloader) shift; xloader=$1; shift; ;;
    *) shift;;
  esac
done

[ -z "$device" -o -z "$machine" ] && usage $0 && exit 1

echo mkSDTool $VERSION
echo

case "$machine" in
"us01")
   KERNEL_URL="http://download.exorembedded.net:8080/Public/usom01/kernel/us01-kit-kernel-${VERSION}.tar.gz"
   ROOTFS_URL="http://download.exorembedded.net:8080/Public/usom01/rootfs/core-image-exor-x11-us01-kit-${VERSION}.tar.gz"
   UBOOT_URL="http://download.exorembedded.net:8080/Public/usom01/bootloader/us01-kit-uboot-${VERSION}.tar.gz"
   XLOADER_URL="http://download.exorembedded.net:8080/Public/usom01/xloader/us01-kit-xloader-${VERSION}.tar.gz"
   ;;
"us02")
   KERNEL_URL="http://download.exorembedded.net:8080/Public/usom02/kernel/us02-kit-kernel-${VERSION}.tar.gz"
   ROOTFS_URL="http://download.exorembedded.net:8080/Public/usom02/rootfs/core-image-exor-x11-us02-kit-${VERSION}.tar.gz"
   UBOOT_URL="http://download.exorembedded.net:8080/Public/usom02/bootloader/us02-kit-uboot-${VERSION}.tar.gz"
   ;;
"us03")
   KERNEL_URL="http://download.exorembedded.net:8080/Public/usom03/kernel/us03-kit-kernel-${VERSION}.tar.gz"
   ROOTFS_URL="http://download.exorembedded.net:8080/Public/usom03/rootfs/core-image-exor-x11-us03-kit-${VERSION}.tar.gz"
   UBOOT_URL="http://download.exorembedded.net:8080/Public/usom03/bootloader/us03-kit-uboot-${VERSION}.tar.gz"
   ;;
"ns01")
   KERNEL_URL="http://download.exorembedded.net:8080/Public/nsom01/kernel/ns01-kit-kernel-${VERSION}.tar.gz"
   ROOTFS_URL="http://download.exorembedded.net:8080/Public/nsom01/rootfs/core-image-exor-x11-ns01-kit-${VERSION}.tar.gz"
   UBOOT_URL="http://download.exorembedded.net:8080/Public/nsom01/bootloader/ns01-kit-uboot-${VERSION}.tar.gz"
   ;;
"ns02")
   KERNEL_URL="http://download.exorembedded.net:8080/Public/nsom02/kernel/ns02-kit-kernel-${VERSION}.tar.gz"
   ROOTFS_URL="http://download.exorembedded.net:8080/Public/nsom02/rootfs/core-image-exor-x11-ns02-kit-${VERSION}.tar.gz"
   UBOOT_URL="http://download.exorembedded.net:8080/Public/nsom02/bootloader/ns02-kit-uboot-${VERSION}.tar.gz"
   XLOADER_URL="http://download.exorembedded.net:8080/Public/nsom02/xloader/ns02-kit-xloader-${VERSION}.tar.gz"
   ;;
*)
   echo "Uknown machine: $machine"
   usage $0
   exit 1
   ;;
esac

[ ! -b $device ] && die "$device is not a block device file"

# Check privileges
[ ! "$(whoami)" = "root" ] && die "this script should be executed with root privileges"

FILES_DIR="$BASE_DIR/mkSDTool_files"
DOWNLOAD_DIR="$FILES_DIR/$machine"
mkdir -p "$DOWNLOAD_DIR/$machine"

# Check necessary files
if [ -n "$kernel" -a ! -f "$kernel" ]; then
   die "specified kernel tarball not found"
else
   kernel="$DOWNLOAD_DIR/kernel.tar.gz"
   if [ ! -f "$kernel" ]; then
      echo Downloading kernel from: $KERNEL_URL
      wget $KERNEL_URL -O "$kernel"
   fi
fi

echo "Using kernel $kernel"

if [ -n "$uboot" -a ! -f "$uboot" ]; then
   die "specified uboot tarball not found"
else
   uboot="$DOWNLOAD_DIR/uboot.tar.gz"
   if [ ! -f "$uboot" ]; then
      echo Downloading uboot from: $UBOOT_URL...
      wget $UBOOT_URL -O "$uboot"
   fi
fi

echo "Using uboot $uboot"

if [ -n "$rootfs" -a ! -f "$rootfs" ]; then
   die "specified rootfs tarball not found"
else
   rootfs="$DOWNLOAD_DIR/rootfs.tar.gz"
   if [ ! -f "$rootfs" ]; then
      echo Downloading rootfs from: $ROOTFS_URL...
      wget $ROOTFS_URL -O "$rootfs"
   fi
fi

echo "Using rootfs $rootfs"

if [ "$machine" == "us01" -o "$machine" == "ns02" ]; then

   if [ -n "$xloader" -a ! -f "$xloader" ]; then
      die "specified xloader tarball not found"
   else
      xloader="$DOWNLOAD_DIR/xloader.tar.gz"
      if [ ! -f "$xloader" ]; then
         echo Downloading xloader from: $XLOADER_URL...
         wget $XLOADER_URL -O "$xloader"
      fi
   fi

   echo "Using xloader $xloader"

fi

echo
echo "************************************************************"
echo "*         THIS WILL DELETE ALL THE DATA ON $device        *"
echo "*                                                          *"
echo "*         WARNING! Make sure your computer does not go     *"
echo "*                  in to idle mode while this script is    *"
echo "*                  running. The script will complete,      *"
echo "*                  but your SD card may be corrupted.      *"
echo "*                                                          *"
echo "*         Press <ENTER> to confirm....                     *"
echo "************************************************************"
read junk

for i in `ls -1 $device? 2>/dev/null`; do
 echo "Unmounting device '$i'"
 umount $i 2>/dev/null || true
done
echo

SECTORS="$( fdisk -l $device | grep -m1 "Disk" | awk '{print $7}' )"
SECTOR_SIZE="$( fdisk -l $device | grep -m1 "Sector size" | awk '{print $7}' )"
DISK_SIZE_MB=$(( ( $SECTOR_SIZE * $SECTORS ) / 1024 / 1024 ))
MB=$(( (1024 * 1024) /$SECTOR_SIZE ))

# Default rootfs partition size in mb
ROOTFS_SIZE=2048

# Default boot partition size in mb
BOOT_SIZE=40

if [ $DISK_SIZE_MB -lt 2500 ]; then
    ROOTFS_SIZE=1024
elif [ $DISK_SIZE_MB -lt 1500 ]; then
    die "please use an SDcard of at least 2GB in size"
fi

echo "NOTE: rootfs partition size will be ${ROOTFS_SIZE} MB"
echo

echo "Creating SD card layout..."

sudo dd if=/dev/zero of=${device} bs=1M count=10 &>/dev/null

case "$machine" in
"us01|us03|ns01")

    sfdisk --force $device << EOF
        $(( $MB )),$(( $BOOT_SIZE * $MB ))
        $(( ( $BOOT_SIZE + 1 ) *$MB )),$(( $ROOTFS_SIZE * $MB ))
EOF
    [ $? -ne 0 ] && die "failed to format media"
    ;;

"us02")

    sfdisk --force $device << EOF
        $(( $MB )),$(( 40*$MB ))
        $(( ( $BOOT_SIZE + 1 ) * $MB )),$(( $ROOTFS_SIZE * $MB ))
        $(( ( $BOOT_SIZE + $ROOTFS_SIZE + 1 ) * $MB )),,a2
EOF
    [ $? -ne 0 ] && die "failed to format media"
    ;;

"ns02")

    sgdisk -Z ${device} &>/dev/null || true
    sgdisk -n 0:0:+50M -t 0:8300 -c 0:"BOOT" ${device} || die "failed to format media"
    sgdisk -n 0:0:+${ROOTFS_SIZE}M -t 0:8300 -c 0:"ROOTFS" ${device} || die "failed to format media"
    sgdisk -n 0:0:+1M -t 0:8300 -c 0:"fsbl1" ${device} || die "failed to format media"
    sgdisk -n 0:0:+1M -t 0:8300 -c 0:"fsbl2" ${device} || die "failed to format media"
    ;;

esac

# Some systems may try to re-mount partitions
sleep 5
umount -l $device* &>/dev/null || true

# Format partitions
echo "Formating BOOT partition..."
mkfs.vfat -F 32 -n "BOOT" ${device}1 || die "failed to format boot partition"

echo "Formating ROOT partition ..."
mkfs.ext4 -F -L "ROOT" ${device}2 || die "failed to format rootfs partitions"

# Mount partitions
mkdir -p /media/mk_BOOT
mount ${device}1 /media/mk_BOOT
mkdir -p /media/mk_ROOT
mount ${device}2 /media/mk_ROOT

# Deploy files
echo "Copying files to SD..."
tar xzf $rootfs -C /media/mk_ROOT || die "failed extracting rootfs"

case "$machine" in
"us01")

   mkdir -p /media/mk_BOOT/boot
   tar xzf $kernel --no-same-owner -C /media/mk_BOOT/boot || die "failed extracting kernel files"
   tar xzf $uboot --no-same-owner -C /media/mk_BOOT || die "failed extracting uboot"
   tar xzf $xloader --no-same-owner -C /media/mk_BOOT || die "failed extracting xloader"
   ;;

"us02")

   tar xzf $kernel --no-same-owner -C /media/mk_BOOT/ || die "failed extracting kernel files"
   tar xzf $uboot -C $FILES_DIR
   dd if=$FILES_DIR/u-boot.img of=${device}3 bs=64k seek=4 || die "failed installing uboot"
   rm $FILES_DIR/u-boot.img
   ;;

"us03|ns01")
   mkdir -p /media/mk_BOOT/boot
   tar xzf $kernel --no-same-owner -C /media/mk_BOOT/boot || die "failed extracting kernel files"
   tar xzf $uboot -C $FILES_DIR
   dd if=$FILES_DIR/u-boot.imx of=${device} bs=1k seek=1 || die "failed installing uboot"
   rm $FILES_DIR/u-boot.imx
   ;;

"ns02")
   mkdir -p /media/mk_BOOT/boot
   tar xzf $kernel --no-same-owner -C /media/mk_BOOT/boot || die "failed extracting kernel files"
   tar xzf $xloader -C $FILES_DIR
   dd if=$FILES_DIR/u-boot-spl.stm32 of=${device}3 || "failed installing xloader"
   dd if=$FILES_DIR/u-boot-spl.stm32 of=${device}4 || "failed installing xloader"
   rm $FILES_DIR/u-boot-spl.stm32
   tar xzf $uboot -C $FILES_DIR
   dd if=$FILES_DIR/u-boot.img of=${device} bs=1K seek=33 || "failed installing uboot"
   rm $FILES_DIR/u-boot.img
   ;;

esac

echo
echo Syncing...
sync

umount -l ${device}1 || true
umount -l ${device}2 || true

rm -rf /media/mk_BOOT
rm -rf /media/mk_ROOT

echo
echo "Completed!"
