#!/bin/bash
#===============================================================================
# Copyright (c) 2012-2015 Wind River Systems, Inc.
# The right to copy, distribute, modify, or otherwise
# make use of this software may be licensed only pursuant
# to the terms of an applicable Wind River license agreement.
#===============================================================================
MYNAME=$(basename $0)
DEPLOYTOOL="/sbin/deploytool"
CAPSULEUPDATETOOL="/sbin/capsule_update"
VFAT_AUTOMOUNT="/boot/efi"

# option to be used by deploytool
DIOPTS=""
DEBUG=0

#===============================================================================
# Functions
#===============================================================================
SHOW() {
    local level=$1 c p e="\n"

    case $level in
        ERROR)   c=31; p="[  ERROR] " ;;
        WARNING) c=32; p="[WARNING] " ;;
        INFO)    c=00; p="[   INFO] " ;;
        DEBUG)   c=00; p="[  DEBUG] " ; [ $DEBUG = 0 ] && return 0 ;;
        *) return 0 ;;
    esac

    shift
    echo -ne "\033[1;${c}m${PREFIX:-$p}$@${END:-$e}\033[0m"
}

EXIT() {
    SHOW "$@"
    SHOW INFO "Exit ..."
    exit 1
}

OPTAPPEND() {
    DIOPTS="$DIOPTS $@"
}

get_user_answer() {
    local ans=""

    echo -ne "\033[1;34m[ ANSWER] $1 [Y/n]? " && read ans
    echo ""
    [ -z "$ans" -o "$ans" = "y" -o "$ans" = "Y" ] && return 0 || return 1
}

sys_get_current_vfat() {
    local curr_vfat=$1
    local root_uuid=$( grep -o '\<root=UUID=[^ ]*'  /proc/cmdline | cut -d "=" -f 3 | sed 's/"//g')
    local root_label=$(grep -o '\<root=LABEL=[^ ]*' /proc/cmdline | cut -d "=" -f 3 | sed 's/"//g')
    local root_lvm=$(  grep -o '\<root=LVM=[^ ]*'   /proc/cmdline | cut -d "=" -f 3 | sed 's/"//g')
    local root_sfr=$(  grep -o '\<root=SFR=[^ ]*'   /proc/cmdline | cut -d "=" -f 3 | sed 's/"//g')
    local root_dev=$(  grep -o '\<root=[^ ]*'       /proc/cmdline | cut -d "=" -f 2 | sed 's/"//g')
    local root_node vg

    [ -z "$root_lvm" ] && root_lvm=$root_sfr
    SHOW DEBUG "Boot root: $(grep -o '\<root=[^ ]*' /proc/cmdline)"

    if [ -n "$root_uuid" ] ; then
        root_node=$(blkid -U $root_uuid)
    elif [ -n "$root_label" ] ; then
        root_node=$(blkid -L $root_label)
    elif [ -n "$root_lvm" ] ; then
        vg=$(basename $root_lvm | awk -F'-' '{print $1}')
        root_node=$(pvs 2>/dev/null | grep "$vg" | tail -n1 | awk '{print $1}')
    elif [ -n "$root_dev" ] ; then
        root_node=$root_dev
    fi

    [ -z "$root_node" ] && return 1
    eval $curr_vfat="${root_node:0:-1}1"
}

show_help()
{
    cat <<EOF
usage:  [environment variables ...] $MYNAME [options]

environment variables:
tgt=<device-name>:      Set the target device node name to be deploy,
                        such as /dev/sda
fmt=<format-type>:      How to create/format partitions.
                            vfat+ext3:  vfat for UEFI boot files and ext3 for rootfs.
                            vfat+ext4:  vfat for UEFI boot files and ext4 for rootfs.
lvm=[0/1]:              Set whether format the rootfs as Logical partition, default 0
lvm_size=<size>:        Set the default size of the logical volume based on the size
                        of the volume group (remaining space is left for a snapshot).
                        N%FREE or N%VG: N=1~100, or N[KMGTPE], N < physical partition size.
                        Default is 45%FREE, which means reserve enough space for a snapshot.
Options:
-p:                     Specify the size of rootfs partition, default to use all space.
-v:                     Enable debug mode.
-h:                     Show this help menu.

EOF

## NOTE ##
#There is a new tool /sbin/deploytool in the system, you are suggested to
#use this new version of deploytool directly instead of ${MYNAME}.
#deploytool provides more function and the uniform interafces in both host
#and target. This $MYNAME tool may will be discarded in the near future.
}

#===============================================================================
# Start Main
#===============================================================================
trap "EXIT ERROR Canceled" SIGINT

if [ ! -x $DEPLOYTOOL ] ; then
    EXIT ERROR "Cannot find $DEPLOYTOOL to do ${MYNAME}!"
fi

# Get environment variables
output=${tgt:-"/dev/mmcblk0"}
format=${fmt:-"vfat+ext3"}
lvm=${lvm:-1}
lvm_size=${lvm_size}

# Parse command line options
while getopts "p:hvY" opt ; do
    echo "$OPTARG" | grep -q "^-.*" && {
        EXIT ERROR "option -$opt requires an argument!"
    }

    case $opt in
        p) OPTAPPEND -p $OPTARG;;
        v) OPTAPPEND -v 7 ; DEBUG=1 ;;
        h) show_help && exit 0 ;;
        ?) show_help && exit 1 ;;
    esac
done
[ -n "$OPTIND" -a $OPTIND -le $# ] && {
    eval opt=\$$OPTIND
    EXIT ERROR "Invalid option: $opt"
}

if [ ! -b $output ] ; then
    EXIT ERROR "Device $output does not exist or is not block device!"
fi

if mount | grep -q "/dev/.*1 on $VFAT_AUTOMOUNT" ; then
    SHOW INFO "vfat boot partition is already mounted."
else
    sys_get_current_vfat current_vfat
    [ -z "$current_vfat" ] && EXIT ERROR "Cannot find current vfat boot partition!"
    SHOW INFO "Mount vfat boot partition $current_vfat at $VFAT_AUTOMOUNT"
    mount $current_vfat $VFAT_AUTOMOUNT
    if ! mount | grep -q "/dev/.*1 on $VFAT_AUTOMOUNT" ; then
        EXIT ERROR "Fail to mount vfat boot partition!"
    fi
fi

rootfs="$(ls $VFAT_AUTOMOUNT/SFR/*.tar.bz2 2>/dev/null | awk '{print $1}')"
if [ -n "$rootfs" ] ; then
    OPTAPPEND -f $rootfs
else
    SHOW ERROR "Cannot find the rootfs tarball: $VFAT_AUTOMOUNT/SFR/*.tar.bz2"
    EXIT ERROR "Maybe the current system is not a formal recovery image."
fi

OPTAPPEND -d $output
OPTAPPEND -t $format
OPTAPPEND -l $lvm
OPTAPPEND -u
[ -n "$lvm_size" ] && OPTAPPEND -L $lvm_size

if ! get_user_answer "Restore device $output to its factory defaults" ; then
    EXIT INFO "Canceled!"
fi

if [ -x "$CAPSULEUPDATETOOL" ] ; then
    get_user_answer "Update the onboard flash" && OPTAPPEND -c
fi

OPTAPPEND -Y
SHOW INFO "Call deploytool to do ${MYNAME}:"
SHOW INFO "$DEPLOYTOOL $DIOPTS"
$DEPLOYTOOL $DIOPTS && exit 0 || exit 1
