#!/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 I2C module and 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)
i2cDriverName=i2c_gpio_param                                    # I2C bit-banged driver module name    
i2cDriverPath=i2c-gpio-param.ko                                 # I2C bit-banged driver module path (relative path)
i2cDeviceNr=5                                                   # the device number that should be assigned to the GPIO bit-banged bus
i2cSdaPinNumber=2                                               # The GPIO pin number that should be used as SDA. If using the default I2C bus on the
                                                                # Raspberry Pi, use pin 2.
i2cSclPinNumber=3                                               # The GPIO pin number that should be used as SCL. If using the default I2C bus on the
                                                                # Raspberry Pi, use pin 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\"`"

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

if [ ! -e $i2cDriver ]
then
    error "The I2C driver \"$i2cDriver\" is missing."
    exit 1
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
    
    lsm=$(lsmod | grep $i2cDriverName)
    if [ $? -eq 0 ]
    then
        echo "WARNING: Module \"$i2cDriverName\" is loaded, we will unload and reload afterwards."
        unload $i2cDriverName
    fi

    # keep track about current i2c devices
    i2cDevicesBeforeLoad=$(ls /sys/class/i2c-dev)

    echo "INFO: Loading driver \"${dir}/$i2cDriverPath\"."
    insmod "${dir}/$i2cDriverPath" busid=${i2cDeviceNr} sda=${i2cSdaPinNumber} scl=${i2cSclPinNumber}
    if [ $? -ne 0 ]
    then
        error "Could not load \"${dir}/$i2cDriverPath\"."
        #exit 1
        errorOccured=true
    fi

    #### we've loaded the bit-banged driver, so the i2c device list might have been changed
    ###i2cDevicesAfterLoad=$(ls /sys/class/i2c-dev)
    ###
    #### convert the strings holding the i2c devices to an array, so we can process the individual devices
    ###i2cDevicesBeforeLoadList=($(echo $i2cDevicesBeforeLoad | tr " " "\n"))
    ###i2cDevicesAfterLoadList=($(echo $i2cDevicesAfterLoad | tr " " "\n"))
    ###
    #### create an array, that holds the differences --> will hold the i2c devices that were added due to bit-banged driver loading
    ###i2cDevicesToDeleteList=()
    ###for i in "${i2cDevicesAfterLoadList[@]}"; do
    ###     skip=
    ###     for j in "${i2cDevicesBeforeLoadList[@]}"; do
    ###         [[ $i == $j ]] && { skip=1; break; }
    ###     done
    ###     [[ -n $skip ]] || i2cDevicesToDeleteList+=("$i")
    ###done
    ###declare -p i2cDevicesToDeleteList
    ###
    #### remove the devices, so we can cleanly add them afterwards
    ###for i2cDev in "${i2cDevicesToDeleteList[@]}"
    ###do
    ###    echo "INFO: Deleting I2C device ${i2cDev}."
    ###    echo ${i2cDev##*-} > /sys/class/i2c-gpio/remove_bus
    ###    if [ $? -ne 0 ]
    ###    then
    ###        error "Could not delete I2C device ${i2cDev}."
    ###        exit 1
    ###    fi
    ###done        
    ###
    ####if [ -e "/sys/class/i2c-dev/i2c-${i2cDeviceNr}" ]
    ####then
    ####    echo "INFO: Deleting I2C device."
    ####    echo ${i2cDeviceNr} > /sys/class/i2c-gpio/remove_bus
    ####fi
    ###
    ###echo "INFO: Creating I2C bus."
    ###echo ${i2cDeviceNr} ${i2cSdaPinNumber} ${i2cSclPinNumber} 2 100 > /sys/class/i2c-gpio/add_bus
    ###if [ $? -ne 0 ]
    ###then
    ###    error "Could not create I2C bus."
    ###    unload $i2cDriverName
    ###    exit 1
    ###fi  

    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."
            unload $i2cDriverName
            #exit 1
            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
        unload $i2cDriverName
        insmod "${dir}/$moitessierDriverPath"
        if [ $? -ne 0 ]
        then
            error "Could not load \"${dir}/$moitessierDriverPath\"."
            exit 1
        fi   
    fi
else
    echo 
    echo "**************************************************************************************"
    echo "*                                                                                    *"
    echo "* FIRMWARE UPDATE FAILED!!!                                                          *"
    echo "*                                                                                    *"
    echo "**************************************************************************************"
    echo
    
    exit 1
fi

exit 0