#!/bin/bash

# This is a small helper script which might be used as a post update service
# to handle "migration"-like tasks after a successful mender update.
# Upgrade scripts (if any) should be placed in /etc/update-hook.d
# or a different directory which then might be specified by using the --scripts-dir parameter.
#
# Upgrade scripts should be named in a way that they are executed in the correct order e.g. by using a prefix
# like 01_script.sh, 05_script.sh, 08_script.sh, ...
# They should be carefully designed to be idempotent (that is subsequent calls should have the same result and
# no undesired side effects) and they should be able to rollback changes if necessary.
# If an upgrade-script fails, this handler script will attempt to rollback changes by calling the previously
# "successfully" run upgrade-scripts in reverse order with a "--rollback" parameter.

# Check if fw_printenv, fw_setenv and mender-update are available
if ! command -v fw_printenv &> /dev/null
then
   echo "ERROR: fw_printenv could not be found"
   exit 0 # exit with 0 since we do not want systemd to reboot in this case
fi

if ! command -v fw_setenv &> /dev/null
then
   echo "ERROR: fw_setenv could not be found"
   exit 0 # exit with 0 since we do not want systemd to reboot in this case
fi

if ! command -v mender-update &> /dev/null
then
   echo "ERROR: mender-update could not be found"
   exit 0 # exit with 0 since we do not want systemd to reboot in this case
fi

# Set static variables
UPGRADE_SCRIPTS_DIR="/etc/update-hook.d"

# Get version info before commiting the update (also still the old version after deployment)
# Use show-provides instead of show-artifact because show-artifact can be altered by module update
OS_VERSION=$(mender-update show-provides 2>/dev/null | grep ^osVersion | cut -d'=' -f2)
APP_VERSION=$(mender-update show-provides 2>/dev/null | grep ^version | cut -d'=' -f2)

# Readout the necessary variables from the U-Boot environment
UPGRADE_AVAILABLE=$(fw_printenv --no-header upgrade_available)
BOOTCOUNT=$(fw_printenv --no-header bootcount)
BOOTLIMIT=$(fw_printenv --no-header bootlimit)
CEL_UPDATE_HOOK_FAILED=$(fw_printenv --no-header cel_update_hook_failed)

# Handle the case where the variable is not set yet
if [ -z "${CEL_UPDATE_HOOK_FAILED}" ]; then
    CEL_UPDATE_HOOK_FAILED=0
fi

echo "System is running with the following parameters:"
echo "upgrade_available: ${UPGRADE_AVAILABLE}"
echo "bootcount: ${BOOTCOUNT}"
echo "bootlimit: ${BOOTLIMIT}"
echo "cel_update_hook_failed: ${CEL_UPDATE_HOOK_FAILED}"

# Check which of the four states we are in
if [ "${UPGRADE_AVAILABLE}" -eq 1 ] && [ "${BOOTCOUNT}" -le "${BOOTLIMIT}" ]; then
    # There was an update and we could boot successfully
    echo "==> Update successful"
    STATE="update_successful"
elif [ "${BOOTCOUNT}" -gt "${BOOTLIMIT}" ]; then
    # There was an update but we failed to boot previously
    echo "==> Update failed"
    STATE="update_failed"
elif [ "${CEL_UPDATE_HOOK_FAILED}" -eq 1 ]; then
    # There was a previous update which failed due to an upgrade-script failure
    echo "==> Previous update failed due to an upgrade-script failure"
    STATE="update_failed_previously"
else
    # There is no update
    echo "==> No Update"
    STATE="no_update"
fi

# Each state need to be handled differently
case $STATE in
    "update_failed")
        echo "Still OS version ${OS_VERSION} with App version ${APP_VERSION} is running"

        # Although we are running from the old partition and upgrade_available was already resetted 
        # by U-Boot, we should still call mender-update rollback to be again in a defined state
        mender-update rollback

        # Reset bootcount to 0 here because mender will not reset this variable until the next 
        # mender-update install command.
        # Do it to avoid triggering mender-update rollback each time the device is rebooted.
        fw_setenv bootcount 0

        # Setting the two custom flag variables to 1 to signalize external services that the update failed
        # This is used by services like the settings service
        fw_setenv cel_update_requested 1
        fw_setenv cel_update_failed 1

        exit 0 # exit with 0 since we do not want systemd to reboot in this case
        ;;

    "update_successful")
        echo "Update in progress"

        # Create a list of files contained within the directory ${UPGRADE_SCRIPTS_DIR}
        scriptsToRun=()

        # Check if the directory with the hook scripts exists and add each script to a list
        if [ ! -d "$UPGRADE_SCRIPTS_DIR" ]; then
            echo "Directory ${UPGRADE_SCRIPTS_DIR} does not exist!"
        else
            # Collect scripts and sort them using version sort (handles numeric prefixes correctly)
            # Version sort treats numeric prefixes as numbers, so 2_ comes before 10_
            # Scripts without numeric prefixes will be sorted alphabetically
            while IFS= read -r -d '' file; do
                scriptsToRun+=("$file")
            done < <(find "${UPGRADE_SCRIPTS_DIR}" -maxdepth 1 -type f -name "*.sh" -print0 | sort -z -V)

            echo "Scripts found in ${UPGRADE_SCRIPTS_DIR}: ${scriptsToRun[*]}"
        fi

        # Commit the update directly if no scripts are found
        if [ ${#scriptsToRun[@]} -eq 0 ]; then
            echo "No scripts found in ${UPGRADE_SCRIPTS_DIR}"
            echo "Committing update..."
            mender-update commit

            # Get the new version info after the commit
            OS_VERSION=$(mender-update show-provides 2>/dev/null | grep ^osVersion | cut -d'=' -f2)
            APP_VERSION=$(mender-update show-provides 2>/dev/null | grep ^version | cut -d'=' -f2)
            echo "Update to OS version ${OS_VERSION} with App version ${APP_VERSION} successful"

            # Setting the custom flag variable for the request to 1 and for the result to 0
            # to signalize external services that the update succeeded
            fw_setenv cel_update_requested 1
            fw_setenv cel_update_failed 0

            # Do not trigger a reboot
            exit 0
        fi

        echo "Scripts to run: ${scriptsToRun[*]}"

        # Create a list of files which are successfully called
        scriptsCalled=()
        upgradeScriptsSuccessful=true
        for script in "${scriptsToRun[@]}"; do
            echo "Executing script: ${script}"

            # Execute the script with the --upgrade parameter and catch any error
            if ! bash "${script}" --upgrade; then
                # Handle the error
                upgradeScriptsSuccessful=false
                echo "Script ${script} failed."
                echo "Attempting rollback on scripts called so far: ${scriptsCalled[*]}"

                # Create a list with scipts which should be rolled back
                reversedScriptsToRun=()
                for ((i=${#scriptsCalled[@]}-1; i>=0; i--)); do
                    reversedScriptsToRun+=("${scriptsCalled[i]}")
                done

                # Execute the rollback scripts in reverse order
                for rollbackScript in "${reversedScriptsToRun[@]}"; do
                    echo "Rolling back script: ${rollbackScript}"

                    # Execute the script with the --rollback parameter
                    # at this point we don't care if the rollback script fails
                    # because it couln't be handled anyways
                    bash "${rollbackScript}" --rollback
                done

                # At this point it's important to break the loop, because we don't
                # want to execute the remaining scripts after one has failed
                break
            fi

            # Add the script to the list of scripts that have been called
            scriptsCalled+=("${script}")
        done

        if [ ${upgradeScriptsSuccessful} = true ]; then
            echo "All upgrade-scripts executed successfully"
            echo "Committing update..."
            mender-update commit

            # Get the new version info after the commit
            OS_VERSION=$(mender-update show-provides 2>/dev/null | grep ^osVersion | cut -d'=' -f2)
            APP_VERSION=$(mender-update show-provides 2>/dev/null | grep ^version | cut -d'=' -f2)
            echo "Update to OS version ${OS_VERSION} with App version ${APP_VERSION} successful"

            # Setting the custom flag variable for the request to 1 and for the result to 0
            # to signalize external services that the update succeeded
            fw_setenv cel_update_requested 1
            fw_setenv cel_update_failed 0

            # Do not trigger a reboot
            exit 0
        else
            echo "Update failed due to an upgrade-script failure"
            echo "Rolling back the update ..."
            mender-update rollback

            # Setting the two custom flag variables to 1 to signalize external services that the update failed
            fw_setenv cel_update_requested 1
            fw_setenv cel_update_failed 1

            # Set another flag variable only used by this service to handle the single case where
            # the update failed and we need a reboot for the rollback
            fw_setenv cel_update_hook_failed 1

            # Trigger a reboot using the systemd service functionality
            exit 1
        fi
        ;;

    "update_failed_previously")
        echo "Still OS version ${OS_VERSION} with App version ${APP_VERSION} is running"

        # Reset the flag variable which was used to handle the single case where
        # the update failed and we needed a reboot for the rollback
        fw_setenv cel_update_hook_failed 0

        # Skip resetting the two variables cel_update_requested and cel_update_failed
        # to signalize external services that the previous update failed

        # Do not trigger a reboot
        exit 0
        ;;

    "no_update")        
        echo "No action needed."

        # Set / reset the custom flag variables to 0 to avoid triggering
        # some mechanism of external services
        fw_setenv cel_update_requested 0
        fw_setenv cel_update_failed 0

        # Do not trigger a reboot
        exit 0
        ;;
esac
