#!/usr/bin/bash

usage() {
  echo "Usage: $0 --tifs <TIFS image> --spl <SPL image> --uboot <U-Boot image>"
  exit 1
}

shutdown_board() {
  echo "Shutting down board..."
  am62pxx-power --power 1 --serial 2 --xds 3 off
  sleep 1
}

# Initialize variables
ARG_TIFS=""
ARG_SPL=""
ARG_UBOOT=""

# Parse long options
while [[ $# -gt 0 ]]; do
  case "$1" in
    --tifs)
      ARG_TIFS="$2"
      shift 2
      ;;
    --spl)
      ARG_SPL="$2"
      shift 2
      ;;
    --uboot)
      ARG_UBOOT="$2"
      shift 2
      ;;
    *)
      echo "Unknown option: $1"
      usage
      ;;
  esac
done

if [[ -z $ARG_TIFS || -z $ARG_SPL || -z $ARG_UBOOT ]]; then
  echo "Error: Missing required options."
  usage
fi

if [ ! -f "$ARG_TIFS" ]; then
  echo "File $ARG_TIFS does not exist"
  exit 1
fi
if [ ! -f "$ARG_SPL" ]; then
  echo "File $ARG_SPL does not exist"
  exit 1
fi
if [ ! -f "$ARG_UBOOT" ]; then
  echo "File $ARG_UBOOT does not exist"
  exit 1
fi

# Function to install required software
install_package() {
  local package=$1
  if ! command -v "$package" &> /dev/null; then
    echo "$package is not installed. Installing $package..."
    apt update
    apt install -y -t bookworm-backports "$package"
  else
    echo "$package is already installed."
  fi
}

# Install required packages
install_package tio
install_package bc

# Check tio version is at least 3.8
TIO_VERSION=$(tio --version | grep -oP 'tio \K[0-9]+\.[0-9]+' || echo "0.0")
REQUIRED_VERSION="3.8"

echo "Installed tio version: $TIO_VERSION"
echo "Required tio version: $REQUIRED_VERSION"

# shellcheck disable=SC2046
if [ $(echo "$TIO_VERSION < $REQUIRED_VERSION" | bc -l) -eq 1 ]; then
  echo "Error: tio version must be at least $REQUIRED_VERSION"
  exit 1
fi

shutdown_board

# Make sure to start board in UART boot mode
# Supported boot modes for [am62pxx-sk]:
#   mmc - 0243
#   uart - 003B
#   emmc - 004B
#   ospi - 0273
#   qspi - 0273
#   usbmsc - 0153
#   spldfu - 0053
echo "Booting on UART mode"
am62pxx-power --power 1 --serial 2 --xds 3 --sysboot 003B on
sleep 1

TFTP_DIR=${PWD#/var/lib/lava/dispatcher/tmp/}

# In accordance with board TRM, the first operation must upload the tiboot3.bin
# file via the X-MODEM protocol
#
# Next, SPL (tispl) must be uploaded via Y-MODEM
#
# Last, u-boot.img same settings as SPL
send_files="\
local timeout = 300000 -- 5 mins
send('${ARG_TIFS}',XMODEM_1K) \
expect('CCC', timeout) \
print('\n') \
\
send('${ARG_SPL}',YMODEM) \
expect('CCC', timeout) \
print('\n') \
\
send('${ARG_UBOOT}',YMODEM) \
expect('U-Boot', timeout) \
print('\n') \
\
print('\nUART Transfer complete\n') \
\
expect('Hit any key to stop autoboot:', timeout) \
write('\n') \
expect('=> ', timeout) \
write('setenv serverip 198.18.0.1; dhcp\n') \
expect('=> ', timeout) \
write('mmc dev 0 1\n') \
expect('=> ', timeout) \
write('tftp 0x82000000 ${TFTP_DIR}/${ARG_TIFS}; mmc write 0x82000000 0x0 0x400\n') \
expect('=> ', timeout) \
write('tftp 0x82000000 ${TFTP_DIR}/${ARG_SPL}; mmc write 0x82000000 0x400 0x1000\n') \
expect('=> ', timeout) \
write('tftp 0x82000000 ${TFTP_DIR}/${ARG_UBOOT}; mmc write 0x82000000 0x1400 0x3000\n') \
expect('=> ', timeout) \
write('mmc partconf 0 1 1 1; mmc bootbus 0 2 0 0\n') \
expect('=> ', timeout) \
\
print('\neMMC Transfer complete\n') \
exit(0)"

# Make sure tio is not running to avoid serial port is locked
pkill tio

# Based on the board Technical Reference Manual (TRM), files must be
# uploaded at 115200 baud, 1 stop bits, and no parity
# tio uses same default required by board
# script tio does not return != 0 when an error occurs, so we need to check the text output
# We want also to see the output real time
while IFS= read -r line; do
  echo "$line"
  if echo "$line" | grep -q "Error: Device file is locked by another process"; then
    echo "tio failed — exiting. Please check that there are no active telnet sessions occupying the serial port"
    shutdown_board
    exit 1
  elif echo "$line" | grep -q "Error: "; then
    echo "tio failed — exiting."
    shutdown_board
    exit 1
  fi
done < <(
  script -q -c "tio --color none --script \"${send_files}\" /dev/ttyUSB0" 2>&1 | grep --line-buffered -E '.*'
)

sleep 1

# Start board on eMMC boot mode
echo "Booting on eMMC mode"
am62pxx-power --power 1 --serial 2 --xds 3 --sysboot 004B on
sleep 1