#!/bin/bash

#
# File .......... swuk_upload_bin
# Author ........ Steve Haywood
# Website ....... http://www.spacewire.co.uk
# Project ....... Common (SpaceWire UK Tutorial)
# Date .......... 18 Jun 2025
# Version ....... 1.0
# Description ...
#   Create a binary file from a bitstream file & upload it.
#

# Strict
#set -euo pipefail


################################################################################
# Main function.
# Arguments ...... ${@}
# Return ......... None
# Exit code ...... Status (0=success, 1=failure)
# Shared (In) .... None
# Shared (Out) ... None
main()
{
  # Declare terminal colour constants
  local -r  c_red='\033[0;31m'
  local -r  c_green='\033[0;32m'
  local -r  c_normal='\033[0m'

  # Declare argument constants
  local -r  c_opthelp="--help"  # Str: Help option name
  local -Ar c_options=(         # ARR: Options & associated help information
    [${c_opthelp}]="Display this help and exit."
  )
  local -ar c_optorder=(        # Arr: Help options display order
    ${c_opthelp}
  )
  local -ar argv=(${@})         # Arr: Get argument values (space-separated) into array

  # Declare constants
  local -r  c_ipre='^(0*(1?[0-9]{1,2}|2([0-4][0-9]|5[0-5]))\.){3}0*(1?[0-9]{1,2}|2([0-4][0-9]|5[0-5]))$'
                                                       # Str: IP validation regular expression
  local -r  c_pwd=$(pwd)                               # Str: Current direcory
  local -r  c_base=$(basename ${c_pwd})                # Str: Project name
  local -r  c_local="fw/vivado/${c_base}.runs/impl_1"  # Str: Location of project Bitstream
  local -r  c_archive="${swuk_user}/firmware"          # Str: Location of user Bitstreams
  local -r  c_remote="/media/sd-mmcblk0p1/firmware"    # Str: Firmware location
  local -r  c_user="root"                              # Str: PetaLinux username
  local -r  c_password="root"                          # Str: PetaLinux password

  # Declare variables
  local     arg       # Str: Current argument from argv array
  local     option    # Str: Current option from c_optorder array
  local     ip=""     # Str: IP Address (x.x.x.x)
  local -i  selcnt=0  # Int: Selection count
  local -A  table     # ARR: Selection table
  local     sel       # Str: User selection input
  local     fpath     # Str: Bitstream location (path & filename)
  local     fname     # Str: Bitstream filename (excluding extension)

  # Display help information
  if [[ " ${argv[*]} " =~ " ${c_opthelp} " ]]; then
    echo "Usage: $(basename ${0}) IP-ADDRESS... [OPTION]..."
    echo "Upload the binary part of a bitstream to the remote system located at IP-ADDRESS."
    echo
    for option in ${c_optorder[@]}
    do
      echo "      ${option} $(printf ' %.0s' {1..12} | head -c $((12-${#option}))) ${c_options[${option}]}"
    done
    echo
    exit 0
  fi

  # Get & check the arguments
  for arg in ${argv[@]}; do
    if [[ ${arg:0:2} == "--" ]]; then  # Option
      [[ ! -v c_options[${arg}] ]] &&
        echo "Option (${arg}) is not recognised!" >&2 &&
        exit 1
    elif [[ -z ${ip} ]]; then  # IP address
      ip=${arg}
    else
      echo "Unexpected argument (${arg}) found!" >&2
      exit 1
    fi
  done

  # Further check the arguments
  [[ -z ${ip} ]] &&
    echo "No IP address specified!" >&2 &&
    exit 1

  # Check the IP ADDRESS
  [[ ! ${ip} =~ ${c_ipre} ]] &&
    echo "The IP address (${ip}) is incorrectly specified!" >&2 &&
    exit 1

  # Check directory at local location
  echo -e "\nBitstream files @ ${c_local} :-\n"
  if [ -d "${c_local}" ]; then
    # Get file at local location
    file="${c_base}.bit"
    if [ -f "${c_local}/${file}" ]; then
      # Add local file to selection list
      table[${selcnt}]="${c_local}/${file}"
      # Display file details
      echo -e "${selcnt}) - ${c_green}${file}${c_normal}"
      # Increase selection count
      selcnt=$((selcnt+1))
    else
      echo -e "${c_red}Bitstream file does not exist in directory!${c_normal}"
    fi
  else
    echo -e "${c_red}Directory does not exist!${c_normal}"
  fi

  # Check directory at archive location
  echo -e "\nBitstream files @ ${c_archive} :-\n"
  if [ -d "${c_archive}" ]; then
    # Get list of files at archive location
    readarray -t files < <(find "${c_archive}" -mindepth 1 -maxdepth 1 -type f -name "*.bit" -printf '%P\n')
    # Sort list of files
    files=($(IFS=$'\n'; echo "${files[*]}" | sort))
    # Iterate through list of files
    for file in "${files[@]}"; do
      # Add archive file to selection list
      table[${selcnt}]="${c_archive}/${file}"
      # Display file details
      echo -e "${selcnt}) - ${c_green}${file}${c_normal}"
      # Increase selection count
      selcnt=$((selcnt+1))
    done
  else
    echo -e "${c_red}Directory does not exist!${c_normal}"
  fi

  # Obtain user input, if good attempt upload
  echo -e "\nq) Quit\n"
  echo -n "Select file for upload or quit"
  while true; do
    read -p " : " sel
    if [[ ("${sel}" == "q") || ("${sel}" == "Q") ]]; then
      echo "Quitting without selection!"
      exit 0
    elif ! [[ "${sel}" =~ ^[0-9]+$ ]] ; then
      # User input is not a number
      echo -n "Selection is not a number, try again"
    elif [ ${sel} -lt 0 ] || [ ${sel} -ge ${selcnt} ]; then
      # User input is out of range
      echo -n "Selection is out of range, try again"
    else
      break
    fi
  done

  # Extract filename from full path
  fpath=${table[${sel}]}
  fname=$(basename -- "${fpath}")
  fname="${fname%.*}"

  # Establish SSH connection
  swuk_ssh ${ip} > /dev/null
  [ ${?} -ne 0 ] && exit 1

  # Extract binary from bitstream
  cp "${fpath}" /tmp/bitstream.bit
  [ ${?} -ne 0 ] && exit 1
  echo -e "all:\n{\n  /tmp/bitstream.bit\n}" > /tmp/bitstream.bif
  bootgen -arch zynq -image /tmp/bitstream.bif -process_bitstream bin -w > /dev/null
  [ ${?} -ne 0 ] && exit 1

  # Upload binary to
  sshpass -p ${c_password} scp /tmp/bitstream.bit.bin ${c_user}@${ip}:${c_remote}/${fname}.bin
  [ ${?} -ne 0 ] && exit 1
  rm /tmp/{bitstream.bit,bitstream.bif,bitstream.bit.bin}

  # Success message
  echo -e "\nUploaded ${fname}.bin to ${c_remote} @ ${ip}."
}


################################################################################
# Opening gambit.
# Arguments ...... ${@}
# Return ......... None
# Exit code ...... Status (0=success, 1=failure)
# Shared (In) .... None
# Shared (Out) ... None
main ${@}
exit 0