#!/bin/bash
################################################################################
# Copyright 2025 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-esc"
USER=$(whoami)
TOOLS=/usr/share/modalai/voxl-esc-tools/
WAS_PX4_ENABLED=false
WAS_PX4_ACTIVE=false
FORCE=false

STARLING_SILVER_MOTORS_FILE="../voxl-esc-params/mavic_mini_2/mavic_mini_2.xml"
STARLING_BLACK_MOTORS_FILE="../voxl-esc-params/D0005_starling_1504_3000kv_dji_props/D0005_starling_1504_3000kv_dji_props.xml"
SENTINEL_V1_FILE="../voxl-esc-params/Sentinel_V1/Sentinel_V1.xml"
FPV_REVB_V3_FILE="../voxl-esc-params/FPV_RevB/FPV_RevB.xml"
FPV_REVB_V4_FILE="../voxl-esc-params/FPV_RevB/m0138_xing2_1855_5x4x3.xml"
D0012_FILE="../voxl-esc-params/D0012_Starling_MAX/esc_params_m0129_f2203_5_1500kv_7x3.7x3.xml"
D0012_V2_FILE="../voxl-esc-params/D0012_Starling_MAX/esc_params_m0129_custom_2204_1500kv_7x3.7x3.xml"
D0013_FILE="../voxl-esc-params/D0013/lumenier_2305_2400kv_GF_D90_4S.xml"
D0014_STARLING_2_FILE="../voxl-esc-params/D0014_starling_1504_3000kv_ms_props/D0014_starling_1504_3000kv_ms_props.xml"
D0019_FILE="../voxl-esc-params/D0019/m0138_xing2_1855kv_gemfan_hurricane_7035_tri.xml"
D0020_FILE="../voxl-esc-params/D0020/m0138_3110_900kv_gemfan_10x4.5_tri.xml"

M0049_FW_FILE="modalai_esc_firmware_m0049"
M0117_1_FW_FILE="modalai_esc_firmware_m0117_1"
M0117_3_FW_FILE="modalai_esc_firmware_m0117_3"
M0129_3_FW_FILE="modalai_esc_firmware_m0129_3"
M0129_6_FW_FILE="modalai_esc_firmware_m0129_6"
M0134_1_FW_FILE="modalai_esc_firmware_m0134_1"
M0134_3_FW_FILE="modalai_esc_firmware_m0134_3"
M0134_6_FW_FILE="modalai_esc_firmware_m0134_6"
M0138_1_FW_FILE="modalai_esc_firmware_m0138_1"
TMOTOR_F55A_FW_FILE="modalai_esc_firmware_tmotor_f55a_revc"

SPIN_POWER="10"
SPIN_TIMEOUT="5"


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


print_usage()
{
	echo ""
	echo "Start wizard with prompts:"
	echo "voxl-esc"
	echo ""
	echo "Shortcut configuration arguments for scripted setup."
	echo ""
	echo "voxl-esc scan"
	echo "voxl-esc detect"
	echo "voxl-esc spin"
	echo "voxl-esc setup_starling_silver_motors"
	echo "voxl-esc setup_starling_black_motors"
	echo "voxl-esc setup_sentinel_v1"
	echo "voxl-esc setup_fpv_revB_v3"
	echo "voxl-esc setup_fpv_revB_v4"
	echo "voxl-esc setup_D0012"
	echo "voxl-esc setup_D0013"
	echo "voxl-esc upgrade_firmware"
	echo ""
	echo "optionally add the 'force' argument to any of the setup_* args above"
	echo "to force flash the ESCs even if the firmware matches and doesn't need"
	echo "updating, this is good for testing."
	echo ""
	echo ""
	echo "show this help message:"
	echo "voxl-esc help"
	echo ""
	exit 0
}

reset_slpi()
{
	if command -v voxl-reset-slpi &> /dev/null
	then
		voxl-reset-slpi -t 8000
	fi
}


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
	reset_slpi
	echo "bridge enabled"
}


disable_bridge()
{
	echo "disabling bridge"
	reset_slpi

	## 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-ESC WITH ERROR${RESET_ALL}"
	exit 1
}


scan()
{
	set +e ## don't exit on error
	cd $TOOLS
	python3 voxl-esc-scan.py
	RET=$?
	if [ $RET -eq 0 ]; then
		echo "successfully pinged ESCs"
	else
		echo "FAILED to ping ESCs when calling scan script"
		return 1
	fi
	return
}

detect()
{
	set +e ## don't exit on error
	cd $TOOLS
	OUTPUT=$(python3 voxl-esc-board-detect.py | grep --text "ESC detected" | head -1)
	
	if [[ $OUTPUT != *"detected"* ]]; then
		echo -e "${RED}[ERROR] No ESCs detected${RESET_ALL}"
		disable_bridge_and_exit_error
	fi
	echo ""
	echo -e "${GRN}${SET_BOLD}$OUTPUT${RESET_ALL}"
	echo ""
	return
}

spin()
{
	set +e ## don't exit on error
	cd $TOOLS
	python3 voxl-esc-scan.py
	RET=$?
	if [ $RET -eq 0 ]; then
		for mot in {0..3}
		do
			python3 voxl-esc-spin.py --id $mot --rpm 1000 --timeout $SPIN_TIMEOUT --skip-prompt True
		done
		RET=$?
	else
		echo -e "${RED}FAILED to ping ESCs when calling spin script${RESET_ALL}"
		disable_bridge_and_exit_error
	fi
	return
}

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

	# set VOXL to performance mode to reduce latency doing ghe firmware upload
	echo "[INFO] Setting CPU to performance mode"
	voxl-set-cpu-mode performance

	# reset_slpi ## just to be safe, sometimes this scan script will fail during SDK flash

	echo "[INFO] Scanning for ESC..."
	# bash gets mad if we don't strip out null bytes
	OUTPUT=$(python3 voxl-esc-scan.py  | tr '\0' ' ' )

	RET=$?
	if [ $RET -eq 0 ]; then
		echo "successfully pinged ESCs"
	else
		echo "FAILED to ping ESCs when calling scan script during upgrade firmware process"
		disable_bridge_and_exit_error
	fi

	BOARD=$(echo "$OUTPUT" | grep --text "Board" | head -1)

	ESC_FW="NONE"
	HASH="NONE"

	if [[ "$BOARD" == *"M0049"* ]]; then
		echo -e "[INFO] M0049 detected"
		ESC_FW=$(ls firmware/ | grep $M0049_FW_FILE)
	elif [[ "$BOARD" == *"M0117-1"* ]]; then
		echo -e "[INFO] M0117-1 detected"
		ESC_FW=$(ls firmware/ | grep $M0117_1_FW_FILE)
	elif [[ "$BOARD" == *"M0117-3"* ]]; then
		echo -e "[INFO] M0117-3 detected"
		ESC_FW=$(ls firmware/ | grep $M0117_3_FW_FILE)
	elif [[ "$BOARD" == *"M0129-3"* ]]; then
		echo -e "[INFO] M0129-3 detected"
		ESC_FW=$(ls firmware/ | grep $M0129_3_FW_FILE)
	elif [[ "$BOARD" == *"M0129-6"* ]]; then
		echo -e "[INFO] M0129-6 detected"
		ESC_FW=$(ls firmware/ | grep $M0129_6_FW_FILE)
	elif [[ "$BOARD" == *"M0134-1"* ]]; then
		echo -e "[INFO] M0134-1 detected"
		ESC_FW=$(ls firmware/ | grep $M0134_1_FW_FILE)
	elif [[ "$BOARD" == *"M0134-3"* ]]; then
		echo -e "[INFO] M0134-3 detected"
		ESC_FW=$(ls firmware/ | grep $M0134_3_FW_FILE)
	elif [[ "$BOARD" == *"M0134-6"* ]]; then
		echo -e "[INFO] M0134-6 detected"
		ESC_FW=$(ls firmware/ | grep $M0134_6_FW_FILE)
	elif [[ "$BOARD" == *"M0138-1"* ]]; then
		echo -e "[INFO] M0138-1 detected"
		ESC_FW=$(ls firmware/ | grep $M0138_1_FW_FILE)
	elif [[ "$BOARD" == *"G071"* ]]; then
		echo -e "[INFO] Tmotor F55A PRO G071 detected"
		ESC_FW=$(ls firmware/ | grep $TMOTOR_F55A_FW_FILE)
	elif [[ "$BOARD" == *"F051"* ]]; then
		echo -e "[INFO] Tmotor F55A PRO F051 detected"
	elif [[ "$BOARD" == *"M0065"* ]]; then
		echo -e "[INFO] ModalAi M0065 PX4IO (M0065) detected"
		echo -e "Currently no firmware for M0065 in the voxl-esc-tool yet"
		disable_bridge
		exit 0
	else
		echo -e "${RED}[ERROR] Unknown board detected: ${BOARD}${RESET_ALL}"
		disable_bridge_and_exit_error
	fi

	## Now extract the has from the file name
	HASH=${ESC_FW%.*} ## trim everything after the .
	HASH=${HASH##*_}
	echo "[INFO] has 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 ESC firmware info:"
	echo -e "\tFirmware: $ESC_FW"
	echo -e "\tFirmware hash: $HASH\n"

	N_CORRECT=$(echo "$OUTPUT" | grep --text "$HASH" | wc -l)
	echo "[INFO] Number of ESCs with correct FW: $N_CORRECT"


	if [[ "$N_CORRECT" == "4" ]]; then
		echo -e "[INFO] All of the ESCs have the correct firmware loaded"
		if ! $FORCE; then
			return
		fi
	fi

	echo -e "\n[INFO] $N_CORRECT of the ESCs do not have the current firmware"
	echo "[INFO] Uploading current ESC firmware..."
	for ESC_ID in {0..3}; do
		python3 voxl-esc-upload-firmware.py --firmware-file firmware/$ESC_FW --id $ESC_ID
		if [[ $? -ne 0 ]]; then
			echo -e "${YLW}[WARNING] failed to upload firmware to ESC${ESC_ID}, trying again${RESET_ALL}"
			reset_slpi
			python3 voxl-esc-upload-firmware.py --firmware-file firmware/$ESC_FW --id $ESC_ID
			if [[ $? -ne 0 ]]; then
				echo -e "${RED}[ERROR] failed to upload firmware to ESC${ESC_ID}${RESET_ALL}"
				disable_bridge_and_exit_error
			fi
		fi
	done

	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  | tr '\0' ' ' )
	RET=$?
	if [ $RET -eq 0 ]; then
		echo "successfully pinged ESCs"
	else
		echo "FAILED to ping ESCs when calling scan script after firmware update"
		disable_bridge_and_exit_error
	fi

	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, only $N_CORRECT of 4 ESCs are now current"
		echo "rescanning:"
		OUTPUT=$(python3 voxl-esc-scan.py  | tr '\0' ' ' )
		RET=$?
		if [ $RET -eq 0 ]; then
			echo "successfully pinged ESCs"
		else
			echo "FAILED to ping ESCs when calling scan script after firmware update during retry"
			disable_bridge_and_exit_error
		fi
		N_CORRECT=$(echo "$OUTPUT" | grep --text "$HASH" | wc -l)
	fi
	if [[ "$N_CORRECT" != "4" ]]; then
		disable_bridge_and_exit_error
	fi

	echo "successfully flashed firmware to all 4 escs"
	return
}


upload_params()
{
	echo "uploading params file $1"
	set +e ## don't exit on error
	cd $TOOLS
	# reset_slpi ## just to be safe

	## scan to make sure we can ping ESCs firs, perhaps not necessary? but it's fast
	python3 voxl-esc-scan.py
	RET=$?
	if [ $RET -eq 0 ]; then
		echo "successfully pinged ESCs"
	else
		echo "ERROR scan script failed before attempting to upload params."
		disable_bridge_and_exit_error
	fi

	## now upload params
	python3 voxl-esc-upload-params.py --params-file "$1"
	RET=$?
	if [ $RET -eq 0 ]; then
		echo "successfully uploaded params"
	else
		echo "ERROR upload params failed"
		disable_bridge_and_exit_error
	fi
	return
}

param_wizard()
{
	echo " "
	echo "Which ?"
	select opt in "d0014_starling_2" "starling_silver_motors" "starling_black_motors" "sentinel_v1" "fpv_revB" "d0013"; do
	case $opt in
	d0014_starling_2)
		upload_params $D0014_STARLING_2_FILE
		break;;
	starling_silver_motors )
		upload_params $STARLING_SILVER_MOTORS_FILE
		break;;
	starling_black_motors )
		upload_params $STARLING_BLACK_MOTORS_FILE
		break;;
	sentinel_v1 )
		upload_params $SENTINEL_V1_FILE
		break;;
	fpv_revB )
		upload_params $FPV_REVB_FILE
		break;;
	d0013 )
		upload_params $D0013_FILE
		break;;
	*)
		echo "invalid option"
		esac
	done
}

## on drones with m0138 esc and a m0185 vtx board, vtx will reset during esc
## update. if static ip is enabled, restart it to get our old ip back
reset_static_ip()
{
	SERVICE_NAME="voxl-static-ip"
	VTX_SERVICE="voxl-vtx"

	if systemctl is-enabled --quiet "$SERVICE_NAME" && ! systemctl is-enabled --quiet "$VTX_SERVICE"; then

		echo "$SERVICE_NAME is enabled AND $VTX_SERVICE is disabled. Resetting $STATIC_IP_SERVICE..."

		# Stop the service if it's running
		systemctl restart $SERVICE_NAME
		echo "$SERVICE_NAME has been reset."
	else
		echo "$SERVICE_NAME is not enabled. No action taken."
	fi

}


################################################################################
## 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:]')


## look for a force flag in case it comes after the other argument
for arg in "$@"
do
	case ${arg} in
		"-f"|"--force"|"f"|force)
			FORCE=true;
		;;
	esac
done


## parse arguments
for arg in "$@"
do
	case ${arg} in
	"")
		echo "Starting Wizard"
		;;
	"-f"|"--force"|"f"|force)
		FORCE=true;
		;;
	"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
		;;
	"setup_starling_v2" | "setup_starling_silver_motors")
		enable_bridge
		upgrade_firmware
		upload_params $STARLING_SILVER_MOTORS_FILE
		disable_bridge
		exit 0
		;;
	"setup_starling_v3" | "setup_starling_black_motors")
		enable_bridge
		upgrade_firmware
		upload_params $STARLING_BLACK_MOTORS_FILE
		disable_bridge
		exit 0
		;;
	"setup_sentinel_v1")
		enable_bridge
		upgrade_firmware
		upload_params $SENTINEL_V1_FILE
		disable_bridge
		exit 0
		;;
	"setup_fpv_revb_v3")
		enable_bridge
		upgrade_firmware
		upload_params $FPV_REVB_V3_FILE
		disable_bridge
		exit 0
		;;
	"setup_fpv_revb_v4")
		enable_bridge
		upgrade_firmware
		upload_params $FPV_REVB_V4_FILE
		disable_bridge
		reset_static_ip
		exit 0
		;;
	"setup_d0012" | "setup_d0012_tmotor_2203.5")
		enable_bridge
		upgrade_firmware
		upload_params $D0012_FILE
		disable_bridge
		exit 0
		;;
	"setup_d0012_custom_2204")
		enable_bridge
		upgrade_firmware
		upload_params $D0012_V2_FILE
		disable_bridge
		exit 0
		;;
	"setup_d0013")
		enable_bridge
		upgrade_firmware
		upload_params $D0013_FILE
		disable_bridge
		exit 0
		;;

	"setup_d0014_starling_2")
		enable_bridge
		upgrade_firmware
		upload_params $D0014_STARLING_2_FILE
		disable_bridge
		exit 0
		;;

	"setup_d0019")
		enable_bridge
		upgrade_firmware
		upload_params $D0019_FILE
		disable_bridge
		reset_static_ip
		exit 0
		;;

	"setup_d0020")
		enable_bridge
		upgrade_firmware
		upload_params $D0020_FILE
		disable_bridge
		reset_static_ip
		exit 0
		;;

	"upgrade_firmware"|"validate_firmware")
		enable_bridge
		upgrade_firmware
		disable_bridge
		exit 0
		;;
	*)
		echo "invalid option"
		print_usage
		exit 1
	esac
done



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

echo " "
echo "What do you want to do?"
select opt in "scan" "detect" "spin" "upgrade_firmware" "upload_params"; 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;;
upload_params )
	enable_bridge
	param_wizard
	disable_bridge
	break;;
*)
	echo "invalid option"
	esac
done

echo "DONE"
exit 0
