#!/bin/sh

backlights_full()
{
    # Turn backlights full on
    echo 255 > /sys/class/leds/lcdled-pwm/brightness 
    echo 255 > /sys/class/leds/keyled-pwm/brightness
}

setup_usb_eject_script()
{
    # Create USB stick unplug script to flash backlights until it is removed 
    echo "#!/bin/sh" > /media/usbremovalwait.sh
    echo "exec 1>/dev/ttyAMA0" >> /media/usbremovalwait.sh
    echo "exec 2>/dev/ttyAMA0" >> /media/usbremovalwait.sh                           
    echo "while [ -d \"${USBROOTPATH}\" ]; do" >> /media/usbremovalwait.sh
    echo "  echo 255 > /sys/class/leds/lcdled-pwm/brightness" >> /media/usbremovalwait.sh  
    echo "  echo 255 > /sys/class/leds/keyled-pwm/brightness" >> /media/usbremovalwait.sh  
    echo "  sleep 1" >> /media/usbremovalwait.sh
    echo "  echo 50 > /sys/class/leds/lcdled-pwm/brightness" >> /media/usbremovalwait.sh  
    echo "  echo 50 > /sys/class/leds/keyled-pwm/brightness" >> /media/usbremovalwait.sh  
    echo "  sleep 1" >> /media/usbremovalwait.sh
    echo "done" >> /media/usbremovalwait.sh
    echo "echo \"Rebooting...\"" >> /media/usbremovalwait.sh 
    echo "reboot" >> /media/usbremovalwait.sh
    # Make it executable and run it
    chmod 775 /media/usbremovalwait.sh
    /media/usbremovalwait.sh &
}

# Successful update
exit_success()
{
    fbv -s 1 ${USBROOTPATH}/images/${PRODUCTID}/recoveryupdatesucceeded.png
    echo SUCCESSFUL UPDATE!
    
    setup_usb_eject_script
    
    exit 0
}

# Error during update
exit_error()
{
    fbv -s 1 ${USBROOTPATH}/images/${PRODUCTID}/recoveryupdateerror.png
    echo ERROR $1 DURING UPDATE!

    setup_usb_eject_script
    
    exit $1
}

display_updating()
{
    fbv -s 1 ${USBROOTPATH}/images/${PRODUCTID}/recoveryupdating.png
}

display_rebooting_to_update()
{
    fbv -s 1 ${USBROOTPATH}/images/${PRODUCTID}/recoveryrebooting.png
}

read_settings()
{
    # Set default settings
    APPPARTSIZEMB=${APPPARTSIZEDEFAULTMB}
    KEEPOLDDATA="N"
    RESTOREOLDDATA="N"

    echo "Extracting settings.sh to /tmp..."
    cd /tmp
    tar xf ${USBROOTPATH}/app/app.tar.gz bin/settings.sh
    chmod 755 /tmp/bin/settings.sh
    # and import the settings (the above defaults may get overriden)
    source /tmp/bin/settings.sh
    rm -rf /tmp/bin
}

# Display message and reboot to continue update
exit_and_reboot_to_continue_update()
{
    # Tell user we are rebooting to updpate
    display_rebooting_to_update
    
    # Wait for a couple of seconds 
    sleep 2
    
    # and reboot
    reboot
    
    # Get out of script        
    exit 0
}

# Check we have all the firmware we require
firmware_validate()
{
    echo Firmware: Validating...

    # Does the directory for this product exist?
    if [ ! -d ${USBROOTPATH}/firmware/${PRODUCTID} ]; then
        exit_error 100
    fi

    # Check we have all the required files
    FIRMWAREFILES="UBOOTVERSION KERNELVERSION ROOTFSVERSION u-boot.sb rootfs.ubi rootfs.tar.gz ${PRODUCTID}/uImage"
    for FIRMWAREFILE in ${FIRMWAREFILES}
    do
        if [ ! -f ${USBROOTPATH}/firmware/${FIRMWAREFILE} ]; then
            echo Firmware: ${USBROOTPATH}/firmware/${FIRMWAREFILE} not found!
            exit_error 101
        fi
    done
}

# Update u-boot if required
# Returns 1 if u-boot was updated, 0 if it wasnt
uboot_update_if_required()
{
    # Get version of U-Boot and update if we need to
    UBOOTVER=$(fw_printenv -n version)
    NEWUBOOTVER=$(cat ${USBROOTPATH}/firmware/UBOOTVERSION)
    if [ "${UBOOTVER}" != "${NEWUBOOTVER}" ]; then
        echo "Bootloader: Updating to ${NEWUBOOTVER}..."
        flash_eraseall /dev/mtd0
        flash_eraseall /dev/mtd1
        kobs-ng init ${USBROOTPATH}/firmware/u-boot.sb
        # The default bootloader environment will check the USB so don't set the usb_recovery flag
        return 1
    fi

    echo "Bootloader: Already running ${NEWUBOOTVER}"
    return 0
}

# Returns 1 if firmware (Kernel or rootfs) need updating, 0 if they are up to date
# NOTE: This is only used in --arsmode
does_firmware_need_update()
{
    # Check running kernel version is up to date
    KERNELVER=$(uname -r)
    NEWKERNELVER=$(cat ${USBROOTPATH}/firmware/KERNELVERSION)
    if [ "$KERNELVER" != "$NEWKERNELVER" ]; then
        echo "Kernel: Needs updating from ${KERNELVER} to ${NEWKERNELVER}"
        return 1
    else
        echo "Kernel: Already running ${NEWKERNELVER}"
    fi
    
    # Check rootfs is up to date
    ROOTFSVER=$(cat /etc/VERSION)
    NEWROOTFSVER=$(cat $USBROOTPATH/firmware/ROOTFSVERSION)
    if [ "$ROOTFSVER" != "$NEWROOTFSVER" ]; then
        echo "RootFs: Needs updating from ${ROOTFSVER} to ${NEWROOTFSVER}"
        return 1
    else
        echo "RootFs: Already running ${NEWROOTFSVER}"
    fi

    # Update not required    
    return 0
}

# Update the firmware
firmware_update()
{
    echo Firmware: Updating...

    uboot_update_if_required
    
    # Update the kernel
    echo "Kernel: Updating..."
    flash_eraseall /dev/mtd2
    nandwrite -p /dev/mtd2 ${USBROOTPATH}/firmware/${PRODUCTID}/uImage
    
    # Format the UBI partition
    ubiformat /dev/mtd4 --yes --flash-image=${USBROOTPATH}/firmware/rootfs.ubi
    ubiattach /dev/ubi_ctrl -m 4 -d 0
    
    mkdir -p /tmp/rootfs
    mount -t ubifs ubi0:rootfs /tmp/rootfs

    echo "RootFs: Updating..."    
    cd /tmp/rootfs
    tar pxf ${USBROOTPATH}/firmware/rootfs.tar.gz
    cd /tmp
    sync
    umount /tmp/rootfs
}

# Check we have all the app files we require
app_validate()
{
    echo App: Validating...

    # Does the directory for the app exist?
    if [ ! -d ${USBROOTPATH}/app ]; then
        exit_error 200
    fi

    # Check we have all the required files
    APPFILES="app.tar.gz"
    for APPFILE in ${APPFILES}
    do
        if [ ! -f ${USBROOTPATH}/app/${APPFILE} ]; then
            echo App: ${USBROOTPATH}/app/${APPFILE} not found!
            exit_error 201
        fi
    done
}

# Update the app
app_update()
{
    echo "App: Removing partitions..."
    ubirmvol /dev/ubi0 --name=data
    ubirmvol /dev/ubi0 --name=app
    
    echo "App: Creating partitions..."
    ubimkvol /dev/ubi0 --name=app --size=${APPPARTSIZEMB}MiB
    ubimkvol /dev/ubi0 --name=data --maxavsize
    
    echo "App: Mounting partitions..."
    mkdir -p /tmp/app
    mount -t ubifs ubi0:app /tmp/app
    mkdir -p /tmp/app/data
    mount -t ubifs ubi0:data /tmp/app/data

    restore_data

    echo "App: Updating..."
    cd /tmp/app
    tar xf ${USBROOTPATH}/app/app.tar.gz
    
    chown -R 0:0 /tmp/app/*
    chmod -R 755 /tmp/app/*
    sync

    echo "App: Updating splashscreen..."
    flash_eraseall /dev/mtd3
    if [ -f "/tmp/app/bin/splashscreen${PRODUCTID}.bmp" ]; then
        # Platform dependent splashscreen
        nandwrite -p /dev/mtd3 /tmp/app/bin/splashscreen${PRODUCTID}.bmp
    else
        # Generic splashscreen
        nandwrite -p /dev/mtd3 /tmp/app/bin/splashscreen.bmp
    fi

    # Dismount and detach UBI
    cd /tmp
    umount /tmp/app/data
    umount /tmp/app
    rm -rf /tmp/app
}

# Backup the data partition to RAMDISK if KEEPOLDDATA is set to "Y"
# $1 is optionally "ATTACHUBI" if this is required (only from USB kernel)
backup_data()
{
    # Only backup data if we need to
    if [ "${KEEPOLDDATA}" == "Y" ]; then
        # Do we need to attach to UBI?
        if [ "$1" == "ATTACHUBI" ]; then
            # Attach to UBI
            ubiattach /dev/ubi_ctrl -m4 -d0
            if [ $? -ne 0 ]; then
                return
            fi
        fi
            
        # Mount the data partition
        mkdir -p /tmp/datatmp
        mount -t ubifs ubi0:data /tmp/datatmp
        if [ $? -eq 0 ]; then
            echo "Backing up data partition..."
            
            # Create the RAMDISK
            mkdir -p ${RAMDISKPATH} 
            mount -n -t tmpfs rwfs ${RAMDISKPATH} -o size=50m

            # Backup data to tar.bz2            
            cd /tmp/datatmp
            tar cjf ${RAMDISKPATH}/databackup.tar.bz2 .
            
            # Set flag to restore the data later on
            RESTOREOLDDATA="Y"

            cd /tmp
            umount /tmp/datatmp
        fi
        
        rm -rf /tmp/datatmp

        # Do we need to dettach to UBI?
        if [ "$1" == "ATTACHUBI" ]; then
            # Detach to UBI
            ubidetach -m4
        fi
    fi
}

# Restore the data we previously saved in backup_data
restore_data()
{
    # Only restore data if we need to
    if [ "${RESTOREOLDDATA}" == "Y" ]; then
        echo "Restoring data partition..."
        
        cd /tmp/app/data
        tar pxf ${RAMDISKPATH}/databackup.tar.bz2
        cd /tmp   
        sync
    fi        
}

###############################################################################

# Redirect output to debug serial
exec 1>/dev/ttyAMA0
exec 2>/dev/ttyAMA0

echo "Running recovery.sh..."

# Get the root pach of the USB
USBROOTPATH=$(cd "${0%/*}/../" && echo $PWD)

# Set the ramdisk path
RAMDISKPATH="/tmp/ramdisk"

# Get the product id and serial from OTP
PRODUCTID=$(printf "%d" $(cat /sys/uid/otp/HW_OCOTP_CUST1))
SERIAL=$(printf "%d" $(cat /sys/uid/otp/HW_OCOTP_CUST0))
echo "Recovery: Running on Product ${PRODUCTID} with serial ${SERIAL}"

# The default APP partition size
APPPARTSIZEDEFAULTMB=32

backlights_full

display_updating

firmware_validate
app_validate

# Has this script been run from the .ars file (from a unit running from NAND rather than USB)?
if [ "$1" == "--arsmode" ]; then
    echo "Running in --arsmode"
    
    # Update u-boot if required 
    uboot_update_if_required
    if [ $? -eq 1 ]; then
        # u-boot was updated. 
        # NOTE: No need to set usb_recovery flag as default environment already has this set! 
        exit_and_reboot_to_continue_update
    fi
    
    # Does the firmware need updating?
    does_firmware_need_update
    if [ $? -eq 1 ]; then
        # Yes. Make bootloader check USB stick on next reboot by setting the flag
        fw_setenv usb_recovery 1

        exit_and_reboot_to_continue_update
    fi
    
    # Its only the app that needs updating        
    read_settings

    # Dismount UBI partitions and detach
    umount /app/data
    umount /app/event
    umount /app
    
    backup_data
    
    app_update
    
    exit_success
else
    read_settings

    # Backup data if required but we need to attach the UBI to do this!
    backup_data ATTACHUBI
    
    firmware_update
    
    app_update
    
    exit_success
fi
