#!/bin/bash
#
# This script is used to update the firmware of the Moitessier HAT microcontroller.
# The latest firmware will be loaded automatically form an online repository.
# 
# Ensure that you have the flashing tool compiled (see 
# https://github.com/mr-rooney/moitessier for further information). You might
# need to change the paths below.
#
# Usage: ./fw_update_moitessier -h

moitessierDriverName=moitessier                                 # the name of the Moitessier HAT driver module
moitessierDriverPath=moitessier.ko                              # the path of the Moitessier HAT driver module (relative path)
flasher=stm32flash                                              # STM32 flasher (path/file name) (relative path)
configFile=/boot/config.txt                                     # the boot configuration file used by the Raspberry Pi (absolute path)
i2cDeviceNr=3                                                   # The bit-banged I2C interface. This parameter must equal with the settings in the 
                                                                # device tree overlay. SDA and SCL for this interface are specified in the overlay.
                                                                # FIXME: currently the kernel 4.19 only supports I2C bus number 3.
repo=https://github.com/mr-rooney/moitessier-firmware.git       # repository holding the latest HAT firmware
fwPath=fw                                                       # the path to store the firmware
fwFileRemote=$fwPath/application.binary                         # path + file name of the firmware, if loaded from the remote repository
blackList=fw_update_moitessier_black_list.txt                   # file that holds the processes/programms that need to be killed before the update can be executed (relative path)
retries=5                                                       # number of retries before giving up the firmware update

# calling path
dir="`dirname \"$0\"`"


dtoverlayI2CgpioValid=0

# https://misc.flogisoft.com/bash/tip_colors_and_formatting
FORMAT_RED="\e[1;31m"
FORMAT_YELLOW="\e[1;93m"
FORMAT_GREEN="\e[1;32m"
FORMAT_DEFAULT="\e[0m"

unload()
{
    echo "INFO: Unloading \"$1\"."
	rmmod $1
}

help()
{
    echo
    echo "******************************************************************************************"
    echo "This script is used to flash the STM32 microcontroller on the Moitessier HAT."
    echo       
    echo "One of the following options can be applied:"
    echo "    -r : Will load the latest firmware release from the remote repository $repo."
    echo "         The firmware will be located in \"${pwd}/${fwFileRemote}\"."
    echo "    -s : Specifies the firmware that should be used."
    echo "         This option requires the path + filename as parameter"
    echo "    -f : Forces the Raspberry Pi to shutdown after firmware update. A shutdown + power cycle"
    echo "         is required only for virgin Moitessier HATs without any firmware."
    echo
    echo "Note: Options -s and -r must not be applied at the same time!!"
    echo "******************************************************************************************"
    echo
}

error()
{
    echo
    echo "******************************************************************************************"
    echo "ERROR occured!!"
    echo
    echo $1
    echo "******************************************************************************************"
    echo
}

pwd=$(pwd)

rOption=0
sOption=0
fOption=0
fwFile=""

help

while getopts 'rfs:h' option
do
    case $option in
        r) rOption=1
            ;;
        f) fOption=1
            ;;
        s) sOption=1
           fwFile="$OPTARG"
            ;;
        h) exit 0
            ;;
        \?) echo
            error "Invalid option."
            exit 1
            ;;
    esac
done

if [ "$EUID" -ne 0 ]
then 
    error "You must run this script with root privileges."
    exit
fi

if [ $rOption == "1" -a $sOption == "1" ]
then
    error "Option -r and -s must not be applied at the same time."
    exit 1
fi

if [ $rOption == "0" -a $sOption == "0" ]
then
    error "Option -r or -s must be applied."
    exit 1
fi

if [ ! -e "${dir}/$flasher" ]
then
    error "The flasher \"${dir}/$flasher\" is missing."
    exit 1
fi

# check if the relevant parameters are set in the configuration file to load the device tree
while IFS='' read -r line || [[ -n "$line" ]]; do
    if [ "$line" = "dtoverlay=i2c-gpio,i2c_gpio_sda=2,i2c_gpio_scl=3,bus=${i2cDeviceNr}" ]
    then
        dtoverlayI2CgpioValid=1
    fi
done < "$configFile"

if [ $dtoverlayI2CgpioValid == "0" ]
then    
    sh -c "echo 'dtoverlay=i2c-gpio,i2c_gpio_sda=2,i2c_gpio_scl=3,bus=${i2cDeviceNr}' >> $configFile"
fi

# we need to reboot the Raspberry Pi to apply the changes in the configuration
if [ $dtoverlayI2CgpioValid == "0" ] 
then
    echo -e "${FORMAT_RED}${configFile} has been changed. System will reboot. You need to call this script afterwards again.${FORMAT_DEFAULT}"
    sleep 10
    reboot
fi

if [ $rOption -eq 1 ]
then
    fwFile=$fwFileRemote

    echo "INFO: Loading latest firmware binary form \"$repo\"."
    if [ -e $fwFile ]
    then
        echo "WARNING: The firmware file \"$fwFile\" already exists."
        
        while true
	    do
		    read -r -p "Overwrite? [y/n]" answer
		    if ([ "$answer" != "y" ] && [ "$answer" != "n" ])
		    then
			    echo "Wrong entry... Try again."
		    else
			    if [ "$answer" = "n" ]
			    then
				    exit 1
			    elif [ "$answer" = "y" ]
			    then
			        rm -r $fwFile
			        rm -r $fwPath
				    break
			    fi
		    fi
	    done
    fi

    git clone $repo fw
    if [ $? -ne 0 ]
    then
        error "Getting firmware from remote respository failed."
        exit 1
    fi   
fi

if [ ! -e $fwFile ]
then
    error "The firmware file \"$fwFile\" is missing."
    exit 1
fi

echo 
echo "**************************************************************************************"
echo "The firmware file \"${fwFile}\" is used for the update"
echo "**************************************************************************************"
echo 

for((ry = 0; ry < $retries; ry++))
do
    errorOccured=false

    if [ "$errorOccured" = false ]
    then
        # Check if OpenPlotter, and especially kplex, is running. If running, we need to stop it, otherwise loading the
        # driver for the Moitessier HAT fails as resource is busy.
        # The processes/programms to kill are specified in the black list file
        echo "INFO: Terminating running processes..."
        while IFS='' read -r line || [[ -n "$line" ]]; do
            toKill="$line"
            # we need to replace carriage return and line feed
            toKill=$(echo ${toKill} | tr -d '\n','\r')
            
            pid=$(pidof ${toKill})
            if [ $? -eq 0 ]
            then
                echo "INFO: Killing process ${toKill}."
                pkill ${toKill}
            else
                echo "INFO: ${toKill} not running."
            fi
        done < "${dir}/${blackList}"

        lsm=$(lsmod | grep $moitessierDriverName)
        if [ $? -eq 0 ]
        then
            echo "WARNING: Module \"$moitessierDriverName\" is loaded, we will unload to gain access to the required GPIO pins."
            unload $moitessierDriverName
            if [ $? -ne 0 ]
            then
                error "Could not unload \"$moitessierDriverName\""
                #exit 1
                errorOccured=true
            fi  
        fi 
    fi

    if [ "$errorOccured" = false ]
    then
        # delay
        secondsElapsed=0
        startTimestamp=$(date +%s)
        while [ $secondsElapsed -lt 8 ]
        do
            printf "."
            sleep 0.2
            currentTimestamp=$(date +%s)
            secondsElapsed=$((currentTimestamp-startTimestamp))
        done
        echo
        echo "INFO: Flashing firmware."
        # GPIO18...nRESET pin
        # GPIO17...BOOT0 pin
        "${dir}/$flasher" -R -i 18,,,17,,,-18,,,,,18,,:,,-17,,,18,,,-18,,,18 -w $fwFile -v -g 0x00 -a 0x41 /dev/i2c-${i2cDeviceNr}
        if [ $? -ne 0 ]
        then
            error "Flashing firmware failed."
            errorOccured=true
        fi   
    fi
    
    if [ "$errorOccured" = false ]
    then
        break;
    else
        error "An error occured, we will retry."
    fi
done

if [ "$errorOccured" = false ]
then
    echo 
    echo "**************************************************************************************"
    echo "*                                                                                    *"
    echo "* FIRMWARE UPDATE SUCCESSFULL!!!                                                     *"
    echo "*                                                                                    *"
    if [ $fOption == "1" ]
    then
    echo "* Shutting down system. YOU NEED TO UNPLUG AND REPLUG THE POWER SUPPLY AFTERWARDS!!! *"
    echo "*                                                                                    *"
    fi
    echo "**************************************************************************************"
    echo

    if [ $fOption == "1" ]
    then
        sleep 4
        shutdown now
    else
        # get kernel release
        kernelRelease=$(uname -r)
        # get machine
        machine=$(uname -m)
        moduleToLoad=${moitessierDriverName}_${kernelRelease}.ko
        insmod "${dir}/${moduleToLoad}"
        if [ $? -ne 0 ]
        then
            error "Could not load \"${dir}/${moduleToLoad}\"."
            exit 1
        fi   
    fi
else
    echo 
    echo "**************************************************************************************"
    echo "*                                                                                    *"
    echo "* FIRMWARE UPDATE FAILED!!!                                                          *"
    echo "*                                                                                    *"
    echo "**************************************************************************************"
    echo
    
    exit 1
fi

exit 0