#!/bin/bash
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -e

SCRIPT_DIR=$(realpath $(dirname $0))
PYTHON3=$(which python3)
FLASHTOOLS_DIR=${SCRIPT_DIR}/flashtools
FBTOOL="${PYTHON3} ${FLASHTOOLS_DIR}/fbtool.py"
EXCELSIOR_DEBUG_RESET_TOOL="${PYTHON3} ${FLASHTOOLS_DIR}/debug_board_reset.py"
DOWNLOAD_AGENT_CONFIG=${FLASHTOOLS_DIR}/dl_addr.ini
DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT=${PRODUCT_OUT}/dl_addr.ini
FASTBOOT_TOOL=$(which fastboot)
USERSPACE_ARCH=armhf

# Check for fastboot tool
if [[ -z "${FASTBOOT_TOOL}" ]]; then
    echo "Fastboot is not found on this machine. Install fastboot before running this script"; exit 1
fi

# Check if PRODUCT_OUT is set
if [[ -z "${PRODUCT_OUT}" ]]; then
    echo "PRODUCT_OUT not set, attempting to use script path ${SCRIPT_DIR}.";
    PRODUCT_OUT=${SCRIPT_DIR}
    DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT=${PRODUCT_OUT}/dl_addr.ini
    # Check that download agent config is available.
    if [[ ! -f "${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}" ]]; then
        echo "Download agent config not found in ${PRODUCT_OUT}."; exit 1
    fi
else
    # Create symlink to download agent config.
    ln -sf ${ROOTDIR}/board/flashtools/dl_addr.ini ${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}
fi


show_help () {
	echo "\
Script for flashing excelsior
Example: $0
	Options:
		-b	option to flash only boot partition."
}


reset_excelsior () {
  ${EXCELSIOR_DEBUG_RESET_TOOL} &
}

reset_excelsior_to_rom() {
  echo "Reset Excelsior to ROM"
  LANG=en_US.UTF-8 LC_ALL=en_US.UTF-8 ${EXCELSIOR_DEBUG_RESET_TOOL} --rom &
}

load_lk_to_ram() {
  pushd ${PRODUCT_OUT}
  echo "Load LK to rom and boot LK"
  ${FBTOOL} -f ${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}
  popd
}

load_fastboot() {
	reset_excelsior_to_rom
	sleep 1
	load_lk_to_ram
	sleep 1
}

erase_mmc() {
  echo "Erase EMMC"
  ${FASTBOOT_TOOL} erase mmc0
  ${FASTBOOT_TOOL} erase mmc0boot0
  ${FASTBOOT_TOOL} erase mmc0boot1
}


flash_boot() {
    pushd ${PRODUCT_OUT}
    ${FASTBOOT_TOOL} flash BOOTIMG1 boot.img
    popd
}

flash_fastboot() {
    pushd ${PRODUCT_OUT}
    ${FASTBOOT_TOOL} flash mmc0 GPT_EMMC
    ${FASTBOOT_TOOL} flash mmc0boot0 lk.img
    ${FASTBOOT_TOOL} flash TEE1 tz.img
    popd
}

flash_all_images() {
    pushd ${PRODUCT_OUT}
    erase_mmc
    ${FASTBOOT_TOOL} flash mmc0 GPT_EMMC
    ${FASTBOOT_TOOL} flash mmc0boot0 lk.img
    ${FASTBOOT_TOOL} flash TEE1 tz.img
    ${FASTBOOT_TOOL} flash BOOTIMG1 boot.img
    ${FASTBOOT_TOOL} flash BOOT boot_${USERSPACE_ARCH}.img
    ${FASTBOOT_TOOL} flash ROOTFS rootfs_${USERSPACE_ARCH}.img
    popd
}


debug=0

FLASHALL=0
BOOTONLY=1
LKONLY=2
mode=$FLASHALL

while getopts "h?vbf" opt; do
	case $opt in
		h|\?)
			show_help
			exit 0
			;;
		v)
			set -x
			debug=1
			;;
		b)
			mode=$BOOTONLY
			;;
		f)
			mode=$LKONLY
			;;
	esac
done

load_fastboot

case $mode in
	$FLASHALL)
		flash_all_images
		sleep 1
		${FASTBOOT_TOOL} reboot
		exit 0
		;;
	$BOOTONLY)
		flash_boot
		sleep 1
		${FASTBOOT_TOOL} reboot
		exit 0
		;;
	$LKONLY)
		flash_fastboot
		sleep 1
		exit 0
		;;
	*) echo "Unknown flash mode."
esac