#!/bin/bash
################################################################################
# Copyright 2024 ModalAI Inc.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# 4. The Software is used solely in conjunction with devices provided by
#    ModalAI Inc.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
################################################################################

NAME="voxl-2-io"
USER=$(whoami)
TOOLS=/usr/share/modalai/voxl-esc-tools/
FIRMWARE=/usr/share/modalai/voxl2-io-tools/firmware/
WAS_PX4_ENABLED=false
WAS_PX4_ACTIVE=false

M0065_FW_FILE="modalai_m0065_firmware"

SPIN_POWER="10"
SPIN_TIMEOUT="5"

RESET_ALL="\e[0m"
RED="\e[91m"
YLW="\e[33m"
GRN="\e[32m"
SET_BOLD="\e[1m"

PORT="2"

print_usage()
{
	echo ""
	echo "Start wizard with prompts:"
	echo "voxl-2-io"
	echo ""
	echo "Shortcut configuration arguments for scripted setup."
	echo ""
	echo "voxl-2-io scan"
	echo "voxl-2-io detect"
	echo "voxl-2-io spin"
	echo "voxl-2-io upgrade_firmware"
	echo ""
	echo "show this help message:"
	echo "voxl-2-io help"
	echo ""
	exit 0
}


enable_bridge()
{
	echo "enabling bridge"
	if [[ "$(systemctl is-enabled voxl-px4)" == "enabled" ]]; then
		WAS_PX4_ENABLED=true
		echo "detected voxl-px4 is enabled"
		systemctl disable voxl-px4
	fi
	if [[ "$(systemctl is-active voxl-px4)" == "active" ]]; then
		WAS_PX4_ACTIVE=true
		echo "detected voxl-px4 is running, stopping it now"
		systemctl stop voxl-px4
		sleep 5
	fi
	echo "bridge enabled"
}


disable_bridge()
{
	echo "disabling bridge"

	## put back the px4 service to how it was
	if $WAS_PX4_ENABLED; then
		echo "re-enabling voxl-px4"
		systemctl enable voxl-px4
	fi
	if $WAS_PX4_ACTIVE; then
		echo "restarting voxl-px4"
		systemctl start voxl-px4
	fi
	echo "bridge disabled"
}


disable_bridge_and_exit_error()
{
	disable_bridge
	echo -e "${RED}${SET_BOLD}EXITING VOXL-2-IO WITH ERROR${RESET_ALL}"
	exit 1
}

display_device_info(){

	# Pass in output to be parsed
	OUTPUT=$1
	
	# Firmware version and hash
	firmware_hash=$(echo "$OUTPUT" | grep -m 1 "Firmware" | awk '{print $6}' | cut -c 1- )
	firmware_version=$(echo "$OUTPUT" | grep -m 1 "Firmware" | awk '{print $4}' | cut -c 1)
	
	# Bootloader version and hash
	bootloader_hash=$(echo "$OUTPUT" | grep -m 1 "Bootloader" | awk '{print $6}' | cut -c 1- )
	bootloader_version=$(echo "$OUTPUT" | grep -m 1 "Bootloader" | awk '{print $4}' | cut -c 1)

	# Board Info
	board_name=$(echo "$OUTPUT" | grep -m 1 "Board" | cut -d ':' -f 3 | cut -c 2- )
	board_version=$(echo "$OUTPUT" | grep -m 1 "Board" | cut -d ':' -f 2 | cut -c 2- )

	# UID 
	uid=$(echo "$OUTPUT" | grep -m 1 "UID" | cut -d ':' -f 2  | cut -c 2- )

	echo -e "INFO: VOXL2 IO Information:"
	echo -e "---------------------"	
	echo -e "	Board      : $board_version : $board_name"
	echo -e "	UID        : $uid"
	echo -e "	Firmware   : version $firmware_version, hash $firmware_hash"
	echo -e "	Bootloader : version $bootloader_version, hash $bootloader_hash"
	echo -e "---------------------"
}


scan()
{
	set +e ## don't exit on error
	cd $TOOLS

	OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 2>/dev/null | tr '\0' ' ')
	BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)
	if [[ "$BOARD" == *"M0065"* ]]; then
		echo "Successfully pinged VOXL2 IO on default port (${PORT})."
		display_device_info "${OUTPUT}"
		return
	else
		echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on default port (${PORT}).${RESET_ALL}"
		PORT="7"
	fi

	OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 --device /dev/slpi-uart-${PORT} 2>/dev/null | tr '\0' ' ' )
	BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)
	if [[ "$BOARD" == *"M0065"* ]]; then
		echo "Successfully pinged VOXL2 IO on alternate port (${PORT})."
		display_device_info "${OUTPUT}"
	else
		echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on alternate port (${PORT}).${RESET_ALL}"
		return 1
	fi

	return
}

detect()
{
	set +e ## don't exit on error
	cd $TOOLS

	OUTPUT=$(python3 voxl-esc-board-detect.py 2>/dev/null | grep --text "ESC detected" | head -1)
	
	if [[ $OUTPUT != *"detected"* ]]; then
		echo -e "${RED}[ERROR] No VOXL2 IO detected on default port (${PORT}).${RESET_ALL}"
		
		# Check alternate port
		PORT="7"
		OUTPUT=$(python3 voxl-esc-board-detect.py --device /dev/slpi-uart-${PORT} 2>/dev/null | grep --text "ESC detected" | head -1)
		if [[ $OUTPUT != *"detected"* ]]; then
			echo -e "${RED}[ERROR] No VOXL2 IO detected on alternate port (${PORT}).${RESET_ALL}"
			disable_bridge_and_exit_error
		fi
		echo -e "[INFO] Detected VOXL2 IO on alternate port (${PORT}).${RESET_ALL}"
	else
		echo -e "[INFO] Detected VOXL2 IO on default port (${PORT}).${RESET_ALL}"
	fi

	echo ""
	echo -e "${GRN}${SET_BOLD}$OUTPUT${RESET_ALL}"
	echo ""
	return
}

spin()
{
	set +e ## don't exit on error
	cd $TOOLS

	OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 2>/dev/null | tr '\0' ' ' )
	BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)

	# Check alternate port if no VOXL2 IO response
	if [[ "$BOARD" != *"M0065"* ]]; then
		echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on default port (${PORT}).${RESET_ALL}"
		PORT="7"
		OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 --device /dev/slpi-uart-${PORT} 2>/dev/null | tr '\0' ' ' )
		BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)
		if [[ "$BOARD" != *"M0065"* ]]; then
			echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on alternate port (${PORT}).${RESET_ALL}"
			disable_bridge_and_exit_error
		else 
			echo "[INFO] Detected VOXL2 IO on alternate port (${PORT})."
		fi
	else 
		echo "[INFO] Detected VOXL2 IO on default port (${PORT})."
	fi

	if [[ "$BOARD" == *"M0065"* ]]; then
		for mot in {0..3}
		do
			python3 voxl-esc-spin.py --device /dev/slpi-uart-${PORT} --id $mot --power $SPIN_POWER --init-cmd 0 --init-time 1.5 --cmd-rate 500 --timeout $SPIN_TIMEOUT --skip-prompt True
		done
		RET=$?
	else
		echo -e "${RED}[ERROR] FAILED to ping VOXL2 IO.${RESET_ALL}"
		disable_bridge_and_exit_error
	fi
	return
}

upgrade_firmware()
{
	set +e ## don't exit on error

	cd $TOOLS
	
	echo "[INFO] Scanning for VOXL2 IO..."

	# bash gets mad if we don't strip out null bytes
	OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 2>/dev/null | tr '\0' ' ' )
	BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)
	BOOTLOADER_MODE=$(echo "$OUTPUT" | grep --text "protocol: bootloader" | head -1)
	
	# Check alternate port if no VOXL2 IO response
	if [[ ! -z "$BOOTLOADER_MODE" ]]; then
		echo -e "[INFO] Bootloader mode detected, assuming M0065 hardware"
		echo -e "Please hit enter to continue updating M0065/VOXL2 IO hardware"
		read
		BOARD="M0065"
	# Check alternate port if no VOXL2 IO response
	elif [[ "$BOARD" != *"M0065"* ]]; then
		echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on default port (${PORT}).${RESET_ALL}"
		PORT="7"
		OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 --device /dev/slpi-uart-$PORT 2>/dev/null | tr '\0' ' ' )
		BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)
		if [[ "$BOARD" != *"M0065"* ]]; then
			echo -e "${RED}[ERROR] Failed to ping VOXL2 IO on alternate port (${PORT}).${RESET_ALL}"
			disable_bridge_and_exit_error
		else 
			echo "[INFO] Detected VOXL2 IO on alternate port (${PORT})."
		fi
	else
		echo "[INFO] Detected VOXL2 IO on default port (${PORT})."
	fi

	ESC_FW="NONE"
	HASH="NONE"

	if [[ "$BOARD" == *"M0049"* ]]; then
		echo -e "[INFO] M0049 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0049.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0117-1"* ]]; then
		echo -e "[INFO] M0117-1 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0117-1.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0117-3"* ]]; then
		echo -e "[INFO] M0117-3 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0117-3.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0129-3"* ]]; then
		echo -e "[INFO] M0129-3 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0129-3.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0134-1"* ]]; then
		echo -e "[INFO] M0134-1 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0134-1.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0134-3"* ]]; then
		echo -e "[INFO] M0134-3 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your M0134-3.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"G071"* ]]; then
		echo -e "[INFO] Tmotor F55A PRO G071 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your G071.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"F051"* ]]; then
		echo -e "[INFO] Tmotor F55A PRO F051 detected"
		echo -e "${YLW}Please use voxl-esc tool to upgrade the firmware on your F051.${RESET_ALL}"
		disable_bridge
		exit 0
	elif [[ "$BOARD" == *"M0065"* ]]; then
		echo "[INFO] ModalAI VOXL2 IO (M0065) detected."
		cd $FIRMWARE
		ESC_FW=$(ls | grep $M0065_FW_FILE)
		cd $TOOLS
	else
		echo -e "${RED}[ERROR] Unknown board detected: ${BOARD}${RESET_ALL}"
		disable_bridge_and_exit_error
	fi

	## Now extract the hash from the file name
	HASH=${ESC_FW%.*} ## trim everything after the .
	HASH=${HASH##*_}
	echo "[INFO] Git hash for most recent firmware: $HASH"

	## should never get here unless a mistake is made adding a new board to the
	## above switch case, keep as a santiy check
	if [[ "$ESC_FW" == "NONE" ]]; then
		echo -e "${RED}[ERROR] No matching FW detected for ESC${RESET_ALL}"
		disable_bridge_and_exit_error
	fi

	echo "[INFO] Expected VOXL2 IO firmware info:"
	echo -e "\tFirmware: $ESC_FW"
	echo -e "\tFirmware hash: $HASH\n"

	firmware_hash=$(echo "$OUTPUT" | grep -m 1 "Firmware" | awk '{print $6}' | cut -c 1- )
	N_CORRECT=$(echo "$OUTPUT" | grep --text "$HASH" | wc -l)
	# echo "[INFO] Number of ESCs with correct FW: $N_CORRECT"


	if [[ "$N_CORRECT" == "4" ]]; then
		if [[ "$firmware_hash" == *"*"* ]]; then
			echo -e "${YLW}[WARNING] Tainted Git hash detected: ${firmware_hash}. Would you like to upgrade firmware to the most recent: ${HASH}? ${RESET_ALL}"
			select opt in "yes" "no"; do
			case $opt in
			yes )
				break;;
			no )
				echo -e "${GRN}[INFO] VOXL2 IO has the correct firmware loaded.${RESET_ALL}"
				return
				break;;
			*)
				echo "invalid option"
				esac
			done
		else 
			echo -e "${GRN}[INFO] VOXL2 IO has the correct firmware loaded.${RESET_ALL}"
			return		
		fi
	fi

	echo -e "\n[INFO] VOXL2 IO does not have the most current firmware"
	echo "[INFO] Uploading most current VOXL2 IO firmware..."
	python3 voxl-esc-upload-firmware.py --device /dev/slpi-uart-$PORT --firmware-file $FIRMWARE/$ESC_FW --id 0 2>/dev/null
	if [[ $? -ne 0 ]]; then
		echo -e "${RED}[ERROR] Failed to upload firmware to VOXL2 IO.${RESET_ALL}"
		disable_bridge_and_exit_error
	fi

	echo -e "\n[INFO] Verifying firmware upload was succecssful..."
	sleep 3 ## sleep so the last ESC we flashed has time to start up
	OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 --device /dev/slpi-uart-$PORT 2>/dev/null| tr '\0' ' ' )
	N_CORRECT=$(echo "$OUTPUT" | grep --text "$HASH" | wc -l)

	## sometimes if an esc doesn't restart in time we get 3 or 4 are correct
	## even though all 4 are good, in this case recheck
	if [[ "$N_CORRECT" != "4" ]]; then
		echo -e "WARNING, VOXL2 IO firmware may not be current"
		echo "rescanning:"
		OUTPUT=$(python3 voxl-esc-scan.py --verify-params 0 --device /dev/slpi-uart-$PORT 2>/dev/null | tr '\0' ' ' )
		N_CORRECT=$(echo "$OUTPUT" | grep --text "$HASH" | wc -l)
	fi
	if [[ "$N_CORRECT" != "4" ]]; then
		disable_bridge_and_exit_error
	fi
	echo -e "${GRN}[INFO] Successfully flashed firmware to VOXL2 IO.${RESET_ALL}"
	return
}


################################################################################
## actual start of execution, handle optional arguments first
################################################################################

## sanity checks
if [ "${USER}" != "root" ]; then
	echo "Please run this script as root"
	exit 1
fi


## convert argument to lower case for robustness
arg=$(echo "$1" | tr '[:upper:]' '[:lower:]')

## parse arguments
case ${arg} in
	"")
		echo "Starting Wizard"
		;;
	"h"|"-h"|"help"|"--help")
		print_usage
		exit 0
		;;
	"enable_bridge")
		enable_bridge
		exit 0
		;;
	"disable_bridge")
		disable_bridge
		exit 0
		;;
	"scan")
		enable_bridge
		scan
		disable_bridge
		exit 0
		;;
	"detect")
		enable_bridge
		detect
		disable_bridge
		exit 0
		;;
	"spin")
		enable_bridge
		spin
		disable_bridge
		exit 0
		;;
	"upgrade_firmware"|"validate_firmware")
		enable_bridge
		upgrade_firmware
		disable_bridge
		exit 0
		;;
	*)
		echo "invalid option"
		print_usage
		exit 1
esac



################################################################################
## no optional arguments, start config wizard prompts
################################################################################

echo " "
echo "What do you want to do?"
select opt in "scan" "detect" "spin" "upgrade_firmware"; do
case $opt in
scan )
	enable_bridge
	scan
	disable_bridge
	break;;
detect )
	enable_bridge
	detect
	disable_bridge
	break;;
spin )
	enable_bridge
	spin
	disable_bridge
	break;;
upgrade_firmware )
	enable_bridge
	upgrade_firmware
	disable_bridge
	break;;
*)
	echo "invalid option"
	esac
done

echo "DONE"
exit 0
