#!/bin/bash -e
  
  help() {
  
  bn=`basename $0`
  cat << EOF
  
  Version: 1.5
  Last change: generate super.img when flash images with dynamic partition feature
  
  eg: sudo ./fastboot_imx_flashall.sh -f imx8mm -a -D ~/android10/evk_8mm/
  eg: sudo ./fastboot_imx_flashall.sh -f imx7ulp -D ~/android10/evk_7ulp/
  
  Usage: $bn <option>
  
  options:
    -h                displays this help message
    -f soc_name       flash android image file with soc_name
    -a                only flash image to slot_a
    -b                only flash image to slot_b
    -c card_size      optional setting: 14 / 28
                          If not set, use partition-table.img/partition-table-dual.img (default)
                          If set to 14, use partition-table-14GB.img for 16GB SD card
                          If set to 28, use partition-table-28GB.img/partition-table-28GB-dual.img for 32GB SD card
                      Make sure the corresponding file exist for your platform
    -m                flash mcu image
    -u uboot_feature  flash uboot or spl&bootloader image with "uboot_feature" in their names
                          For Standard Android:
                              If the parameter after "-u" option contains the string of "dual", then spl&bootloader image will be flashed,
                              otherwise uboot image will be flashed
                          For Android Automative:
                              only dual bootloader feature is supported, by default spl&bootloader image will be flashed
                          Below table lists the legal value supported now based on the soc_name provided:
                             ┌────────────────┬──────────────────────────────────────────────────────────────────────────────────────────────────────┐
                             │   soc_name     │  legal parameter after "-u"                                                                          │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mm       │  dual trusty-dual 4g-evk-uuu 4g ddr4-evk-uuu ddr4 evk-uuu trusty-4g trusty-secure-unlock trusty      │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mn       │  dual trusty-dual evk-uuu trusty-secure-unlock trusty ddr4-evk-uuu ddr4                              │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mp       │  dual trusty-dual evk-uuu trusty-secure-unlock trusty                                                │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤

                             │   imx8mq       │  smarc_2g trusty-smarc_2g trusty-dual-smarc_2g trusty-secure-unlock-smarc_2g smarc_4g trusty-smarc_4g trusty-dual-smarc_4g trusty-secure-unlock-smarc_4g                             │

                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
  ───────────────────────────────────────────────────────────────────┤
                             │   imx8mq       │  smarc_4g-uuu smarc_4g-trusty-secure-unlock smarc_2g-uuu smarc_2g-trusty-secure-unlock               │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8qxp      │  mek-uuu trusty-secure-unlock trusty secure-unlock c0 trusty-c0 mek-c0-uuu                           │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8qm       │  mek-uuu trusty-secure-unlock trusty secure-unlock md hdmi                                           │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx7ulp      │  evk-uuu                                                                                             │
                             └────────────────┴──────────────────────────────────────────────────────────────────────────────────────────────────────┘
  
    -d dtb_feature    flash dtbo, vbmeta and recovery image file with "dtb_feature" in their names
                          If not set, default dtbo, vbmeta and recovery image will be flashed
                          Below table lists the legal value supported now based on the soc_name provided:
                             ┌────────────────┬──────────────────────────────────────────────────────────────────────────────────────────────────────┐
                             │   soc_name     │  legal parameter after "-d"                                                                          │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mm       │  ddr4 m4 mipi-panel                                                                                  │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mn       │  mipi-panel rpmsg ddr4 ddr4-mipi-panel ddr4-rpmsg                                                    │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤

                             │   imx8mq       │  hdmi dp dcss-lvds lcdif-lvds dual-display                                                           │

                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8mp       │  rpmsg hdmi lvds-panel lvds mipi-panel basler ov5640                                                 │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8qxp      │                                                                                                      │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx8qm       │  hdmi mipi-panel md xen                                                                              │
                             ├────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────┤
                             │   imx7ulp      │  evk-mipi evk mipi                                                                                   │
                             └────────────────┴──────────────────────────────────────────────────────────────────────────────────────────────────────┘
  
    -e                erase user data after all image files being flashed
    -l                lock the device after all image files being flashed
    -D directory      the directory of images
                          No need to use this option if images are in current working directory
    -s ser_num        the serial number of board
                          If only one board connected to computer, no need to use this option
    -super            do not generate super.img when flash the images with dynamic partition feature enabled.
                          Under the condition that dynamic partition feature are enabled:
                            if this option is not used, super.img will be generated under temporary directory and flashed to the board.
                            if this option is used, make sure super.img already exists together with other images.
    -tmp temp_dir     specify the temporary directory. Default temporary directory is "/tmp".
  EOF
  
  }
  
  # this function checks whether the value of first parameter is in the array value of second parameter
  # pass the name of the (array)variable to this function. the first is potential element, the second one is array.
  # make sure the first parameter is not empty
  function whether_in_array
  {
      local potential_element=`eval echo \$\{${1}\}`
      local array=(`eval echo \$\{${2}\[\*\]\}`)
      local array_length=${#array[*]}
      local last_element=${array[${array_length}-1]}
      for arg in ${array[*]}
      do
          if [ "${arg}" = "${potential_element}" ]; then
              result_value=0
              return 0
          fi
          if [ "${arg}" = "${last_element}" ]; then
              result_value=1
              return 0
          fi
      done
  }
  
  function flash_partition
  {
      if [ ${support_dual_bootloader} -eq 1 ] && [ "$(echo ${1} | grep "bootloader_")" != "" ]; then
          img_name=${uboot_proper_to_be_flashed}
      elif [ "$(echo ${1} | grep "system")" != "" ]; then
          img_name=${systemimage_file}
      elif [ "$(echo ${1} | grep "vendor")" != "" ]; then
          img_name=${vendor_file}
      elif [ "$(echo ${1} | grep "product")" != "" ]; then
          img_name=${product_file}
      elif [ "$(echo ${1} | grep "bootloader")" != "" ]; then
           img_name=${bootloader_flashed_to_board}
      elif [ ${support_dtbo} -eq 1 ] && [ "$(echo ${1} | grep "boot")" != "" ]; then
          img_name="boot.img"
      elif [ "$(echo ${1} | grep `echo ${mcu_os_partition}`)" != "" ]; then
          if [ "${soc_name}" = "imx7ulp" ]; then
              img_name="${soc_name}_m4_demo.img"
          else
              img_name="${soc_name}_mcu_demo.img"
          fi
      elif [ "$(echo ${1} | grep -E "dtbo|vbmeta|recovery")" != "" -a "${dtb_feature}" != "" ]; then

          img_name="${1%_*}-${soc_name}-smarc-${dtb_feature}.img"

      elif [ "$(echo ${1} | grep "gpt")" != "" ]; then
          img_name=${partition_file}
      elif [ "$(echo ${1} | grep "super")" != "" ]; then
          if [ ${dont_generate_super} -eq 0 ]; then
              make_super_image
          fi
          img_name=${super_file}
      else
          img_name="${1%_*}-${soc_name}.img"
      fi
  
      echo -e flash the file of ${GREEN}${img_name}${STD} to the partition of ${GREEN}${1}${STD}
      if [ "${img_name}" = "${super_file}" ] && [ ${dont_generate_super} -eq 0 ]; then
          ${fastboot_tool} flash ${1} "${temp_dir}${img_name}"
      else
          ${fastboot_tool} flash ${1} "${image_directory}${img_name}"
      fi
  }
  
  function flash_userpartitions
  {
      if [ ${support_dtbo} -eq 1 ]; then
          flash_partition ${dtbo_partition}
      fi
  
      flash_partition ${boot_partition}
  
      if [ ${support_recovery} -eq 1 ]; then
          flash_partition ${recovery_partition}
      fi
  
      if [ ${support_dynamic_partition} -eq 0 ]; then
          flash_partition ${system_partition}
          flash_partition ${vendor_partition}
          flash_partition ${product_partition}
      fi
      flash_partition ${vbmeta_partition}
  }
  
  function flash_partition_name
  {
      boot_partition="boot"${1}
      recovery_partition="recovery"${1}
      system_partition="system"${1}
      vendor_partition="vendor"${1}
      product_partition="product"${1}
      vbmeta_partition="vbmeta"${1}
      dtbo_partition="dtbo"${1}
  }
  
  function flash_android
  {
      # a precondition: the location of gpt partition and the partition for uboot or spl(in dual bootloader condition)
      # should be the same for the u-boot just boot up the board and the on to be flashed to the board
      flash_partition "gpt"
  
      ${fastboot_tool} getvar all 2>${temp_dir}fastboot_var.log
      grep -q "bootloader_a" ${temp_dir}fastboot_var.log && support_dual_bootloader=1
      grep -q "dtbo" ${temp_dir}fastboot_var.log && support_dtbo=1
      grep -q "recovery" ${temp_dir}fastboot_var.log && support_recovery=1
      # use boot_b to check whether current gpt support a/b slot
      grep -q "boot_b" ${temp_dir}fastboot_var.log && support_dualslot=1
      grep -q "super" ${temp_dir}fastboot_var.log && support_dynamic_partition=1
  
      # some partitions are hard-coded in uboot, flash the uboot first and then reboot to check these partitions
  
      # uboot or spl&bootloader
      if [ ${support_dual_bootloader} -eq 1 ]; then

          bootloader_flashed_to_board="spl-${soc_name}${uboot_feature}.bin"
          uboot_proper_to_be_flashed="bootloader-${soc_name}${uboot_feature}.img"

      else

          bootloader_flashed_to_board="u-boot-${soc_name}${uboot_feature}.imx"

      fi
  
      # in the source code, if AB slot feature is supported, uboot partition name is bootloader0, otherwise it's bootloader
      if [ ${support_dualslot} -eq 1 ]; then
           flash_partition "bootloader0"
      else
           flash_partition "bootloader"
      fi
  
      # if a platform doesn't support dual slot but a slot is selected, ignore it.
      if [ ${support_dualslot} -eq 0 ] && [ "${slot}" != "" ]; then
          slot=""
      fi
  
  
      #if dual-bootloader feature is supported, we need to flash the u-boot proper then reboot to get hard-coded partition info
      if [ ${support_dual_bootloader} -eq 1 ]; then
          if [ "${slot}" != "" ]; then
              dual_bootloader_partition="bootloader"${slot}
              flash_partition ${dual_bootloader_partition}
              ${fastboot_tool} set_active ${slot#_}
          else
              dual_bootloader_partition="bootloader_a"
              flash_partition ${dual_bootloader_partition}
              dual_bootloader_partition="bootloader_b"
              flash_partition ${dual_bootloader_partition}
              ${fastboot_tool} set_active a
          fi
      fi
  
      # full uboot is flashed to the board and active slot is set, reboot to u-boot fastboot boot command
      ${fastboot_tool} reboot bootloader
      sleep 5
  
      ${fastboot_tool} getvar all 2>${temp_dir}fastboot_var.log
      grep -q `echo ${mcu_os_partition}` ${temp_dir}fastboot_var.log && support_mcu_os=1
  
      if [ ${flash_mcu} -eq 1 -a ${support_mcu_os} -eq 1 ]; then
          flash_partition ${mcu_os_partition}
      fi
  
      if [ "${slot}" = "" ] && [ ${support_dualslot} -eq 1 ]; then
          #flash image to a and b slot
          flash_partition_name "_a"
          flash_userpartitions
  
          flash_partition_name "_b"
          flash_userpartitions
      else
          flash_partition_name ${slot}
          flash_userpartitions
      fi
      # super partition does not have a/b slot, handle it individually
      if [ ${support_dynamic_partition} -eq 1 ]; then
          flash_partition ${super_partition}
      fi
  }
  
  # this function will invoke lpmake to create super.img, the super.img will
  # be created in temporary directory, make sure that there is enouth space
  function make_super_image
  {
      rm -rf ${temp_dir}${super_file}
      # now dynamic partition is only enabled in dual slot condition
      if [ ${support_dualslot} -eq 1 ]; then
          if [ "${slot}" == "_a" ]; then
              lpmake_system_image_a="--image system_a=${image_directory}${systemimage_file}"
              lpmake_vendor_image_a="--image vendor_a=${image_directory}${vendor_file}"
              lpmake_product_image_a="--image product_a=${image_directory}${product_file}"
          elif [ "${slot}" == "_b" ]; then
              lpmake_system_image_b="--image system_b=${image_directory}${systemimage_file}"
              lpmake_vendor_image_b="--image vendor_b=${image_directory}${vendor_file}"
              lpmake_product_image_b="--image product_b=${image_directory}${product_file}"
          else
              lpmake_system_image_a="--image system_a=${image_directory}${systemimage_file}"
              lpmake_vendor_image_a="--image vendor_a=${image_directory}${vendor_file}"
              lpmake_product_image_a="--image product_a=${image_directory}${product_file}"
              lpmake_system_image_b="--image system_b=${image_directory}${systemimage_file}"
              lpmake_vendor_image_b="--image vendor_b=${image_directory}${vendor_file}"
              lpmake_product_image_b="--image product_b=${image_directory}${product_file}"
          fi
      fi
  
      ${image_directory}lpmake --metadata-size 65536 --super-name super --metadata-slots 3 --device super:7516192768 \
          --group nxp_dynamic_partitions_a:3747610624 --group nxp_dynamic_partitions_b:3747610624 \
          --partition system_a:readonly:0:nxp_dynamic_partitions_a ${lpmake_system_image_a} \
          --partition system_b:readonly:0:nxp_dynamic_partitions_b ${lpmake_system_image_b} \
          --partition vendor_a:readonly:0:nxp_dynamic_partitions_a ${lpmake_vendor_image_a} \
          --partition vendor_b:readonly:0:nxp_dynamic_partitions_b ${lpmake_vendor_image_b} \
          --partition product_a:readonly:0:nxp_dynamic_partitions_a ${lpmake_product_image_a} \
          --partition product_b:readonly:0:nxp_dynamic_partitions_b ${lpmake_product_image_b} \
          --sparse --output ${temp_dir}${super_file}
  }
  
  
  # parse command line
  soc_name=""
  uboot_feature=""
  dtb_feature=""
  card_size=0
  slot=""
  systemimage_file="system.img"
  vendor_file="vendor.img"
  product_file="product.img"
  partition_file="partition-table.img"
  super_file="super.img"
  support_dtbo=0
  support_recovery=0
  support_dualslot=0
  support_mcu_os=0
  support_dual_bootloader=0
  support_dynamic_partition=0
  dual_bootloader_partition=""
  bootloader_flashed_to_board=""
  uboot_proper_to_be_flashed=""
  boot_partition="boot"
  recovery_partition="recovery"
  system_partition="system"
  vendor_partition="vendor"
  product_partition="product"
  vbmeta_partition="vbmeta"
  dtbo_partition="dtbo"
  mcu_os_partition="mcu_os"
  super_partition="super"
  flash_mcu=0
  lock=0
  erase=0
  image_directory=""
  ser_num=""
  fastboot_tool="fastboot"
  RED='\033[0;31m'
  STD='\033[0;0m'
  GREEN='\033[0;32m'
  lpmake_system_image_a=""
  lpmake_system_image_b=""
  lpmake_vendor_image_a=""
  lpmake_vendor_image_b=""
  lpmake_product_image_a=""
  lpmake_product_image_b=""
  result_value=0
  dont_generate_super=0
  temp_dir="/tmp/"
  
  
  # We want to detect illegal feature input to some extent. Here it's based on SoC names. Since an SoC may be on a
  # board running different set of images(android and automative for a example), so misuse the features of one set of
  # images when flash another set of images can not be detect early with this scenario.
  imx8mm_uboot_feature=(dual trusty-dual 4g-evk-uuu 4g ddr4-evk-uuu ddr4 evk-uuu trusty-4g trusty-secure-unlock trusty)
  imx8mn_uboot_feature=(dual trusty-dual evk-uuu trusty-secure-unlock trusty ddr4-evk-uuu ddr4)

  imx8mq_uboot_feature=(smarc_2g smarc_2g-dp trusty-smarc_2g trustu-dual-smarc_2g trusty-secure-unlock-smarc_2g smarc_4g smarc_4g-dp trusty-smarc_4g trustu-dual-smarc_4g trusty-secure-unlock-smarc_4g)

  imx8mp_uboot_feature=(dual trusty-dual evk-uuu trusty-secure-unlock trusty)
  imx8qxp_uboot_feature=(mek-uuu trusty-secure-unlock trusty secure-unlock c0 trusty-c0 mek-c0-uuu)
  imx8qm_uboot_feature=(mek-uuu trusty-secure-unlock trusty secure-unlock md hdmi)
  imx7ulp_uboot_feature=(evk-uuu)
  
  imx8mm_dtb_feature=(ddr4 m4 mipi-panel)
  imx8mn_dtb_feature=(mipi-panel rpmsg ddr4 ddr4-mipi-panel ddr4-rpmsg)

  imx8mq_dtb_feature=(hdmi dcss-lvds lcdif-lvds dp dual-display)

  imx8mp_dtb_feature=(rpmsg hdmi lvds-panel lvds mipi-panel basler ov5640)
  imx8qxp_dtb_feature=()
  imx8qm_dtb_feature=(hdmi mipi-panel md xen)
  imx7ulp_dtb_feature=(evk-mipi evk mipi)
  
  # an array to collect the supported soc_names

  supported_soc_names=(imx8qm imx8qxp imx8mq imx8mm imx8mn imx8mp imx7ulp)

  
  if [ $# -eq 0 ]; then
      echo -e ${RED}no parameter specified, will directly exit after displaying help message${STD}
      help; exit 1;
  fi
  while [ $# -gt 0 ]; do
      case $1 in
          -h) help; exit ;;
          -f) soc_name=$2; shift;;
          -c) card_size=$2; shift;;
          -u) uboot_feature=-$2; shift;;
          -d) dtb_feature=$2; shift;;
          -a) slot="_a" ;;
          -b) slot="_b" ;;
          -m) flash_mcu=1 ;;
          -e) erase=1 ;;
          -l) lock=1 ;;
          -D) image_directory=$2; shift;;
          -s) ser_num=$2; shift;;
          -super) dont_generate_super=1 ;;
          -tmp) temp_dir=$2; shift;;
          *)  echo -e ${RED}$1${STD} is not an illegal option
              help; exit;;
      esac
      shift
  done
  
  # check whether the soc_name is legal or not
  if [ -n "${soc_name}" ]; then
      whether_in_array soc_name supported_soc_names
      if [ ${result_value} != 0 ]; then
          echo -e >&2 ${RED}illegal soc_name \"${soc_name}\"${STD}
          help; exit 1;
      fi
  else
      echo -e >&2 ${RED}use \"-f\" option to specify the soc name${STD}
  fi
  
  
  # Process of the uboot_feature parameter
  if [[ "${uboot_feature}" = *"dual"* ]]; then
      support_dual_bootloader=1;
  fi
  
  # if card_size is not correctly set, exit.
  if [ ${card_size} -ne 0 ] && [ ${card_size} -ne 7 ] && [ ${card_size} -ne 14 ] && [ ${card_size} -ne 28 ]; then
      help; exit 1;
  fi
  
  # Android Automative by default support dual bootloader, no "dual" in its partition table name
  if [ ${support_dual_bootloader} -eq 1 ]; then
      if [ ${card_size} -gt 0 ]; then
          partition_file="partition-table-${card_size}GB-dual.img";
  
      else
          partition_file="partition-table-dual.img";
      fi
  else
      if [ ${card_size} -gt 0 ]; then
          partition_file="partition-table-${card_size}GB.img";
      else
          partition_file="partition-table.img";
      fi
  fi
  
  if [[ "${temp_dir}" != "" ]]; then
      temp_dir="${temp_dir%/}/";
      mkdir -p ${temp_dir}
  fi
  
  # if directory is specified, make sure there is a slash at the end
  if [[ "${image_directory}" = "" ]]; then
      image_directory=`pwd`
  fi
  image_directory="${image_directory%/}/"
  
  if [[ "${ser_num}" != "" ]]; then
      fastboot_tool="fastboot -s ${ser_num}"
  fi
  
  # check whether provided spl/bootloader/uboot feature is legal
  if [ -n "${uboot_feature}" ]; then
      uboot_feature_no_pre_hyphen=${uboot_feature#-}
      whether_in_array uboot_feature_no_pre_hyphen ${soc_name}_uboot_feature
      if [ ${result_value} != 0 ]; then
          echo -e >&2 ${RED}illegal parameter \"${uboot_feature_no_pre_hyphen}\" for \"-u\" option${STD}
          help; exit 1;
      fi
  fi
  
  # check whether provided dtb feature is legal
  if [ -n "${dtb_feature}" ]; then
      dtb_feature_no_pre_hyphen=${dtb_feature#-}
      whether_in_array dtb_feature_no_pre_hyphen ${soc_name}_dtb_feature
      if [ ${result_value} != 0 ]; then
          echo -e >&2 ${RED}illegal parameter \"${dtb_feature}\" for \"-d\" option${STD}
          help; exit 1;
      fi
  fi
  
  flash_android
  
  if [ ${erase} -eq 1 ]; then
      if [ ${support_dualslot} -eq 0 ] ; then
          ${fastboot_tool} erase cache
      fi
      ${fastboot_tool} erase misc
      ${fastboot_tool} erase metadata
      ${fastboot_tool} erase userdata
  fi
  
  if [ ${lock} -eq 1 ]; then
      ${fastboot_tool} oem lock
  fi
  
  if [ ${support_dualslot} -eq 1 ] && [ "${slot}" != "" ]; then
      ${fastboot_tool} set_active ${slot#_}
  fi
  
  echo
  echo ">>>>>>>>>>>>>> Flashing successfully completed <<<<<<<<<<<<<<"
  
  exit 0