#!/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.

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FASTBOOT_CMD="$(which fastboot)"
DRAM_SIZE=2G

function die {
  echo "$@" >/dev/stderr
  exit 1
}

function try {
  $@ || die "Failed to execute '$@'"
}

function load_fastboot {
  local files_dir="$1"
  python3 "${files_dir}/flashtools/debug_board_reset.py" --rom &
  sleep 1
  pushd "${files_dir}"
  python3 "${files_dir}/flashtools/fbtool.py" -f "${files_dir}/dl_addr.ini"
  sleep 1
  popd
}

function flash_partitions {
  local files_dir="$1"
  local flash_with_uboot="$2"

  # Load LK if we aren't using u-boot
  # Also, only LK can handle the mmc0* partitions.
  if [[ ${flash_with_uboot} != "true" ]]; then
    load_fastboot "${files_dir}"

    try ${FASTBOOT_CMD} erase mmc0
  fi
  try ${FASTBOOT_CMD} erase mmc0boot0
  try ${FASTBOOT_CMD} erase mmc0boot1
  try ${FASTBOOT_CMD} flash mmc0 "${files_dir}/partition-table.img"
  try ${FASTBOOT_CMD} flash mmc0boot0 "${files_dir}/bl2.img"
  try ${FASTBOOT_CMD} flash mmc0boot1 "${files_dir}/u-boot-env.bin"
  try ${FASTBOOT_CMD} flash bootloaders "${files_dir}/fip.bin"
  try ${FASTBOOT_CMD} flash boot "${files_dir}/boot_arm64.img"
  try ${FASTBOOT_CMD} flash rootfs "${files_dir}/rootfs_arm64.img"
}

function ensure_flash_files_present {
  local found_all_files=true
  local files_dir="$1"; shift
  local link_dir="$1"; shift
  local files=(
    "bl2.img"
    "boot_arm64.img"
    "dl_addr.ini"
    "fip.bin"
    "flashtools/debug_board_reset.py"
    "flashtools/dl_addr.ini"
    "flashtools/fbtool.py"
    "lk.bin"
    "partition-table.img"
    "rootfs_arm64.img"
    "u-boot-env.bin"
  )

  if [[ "${PRODUCT_OUT}" == "${files_dir}" ]]; then
    ln -sf "${ROOTDIR}/board/flashtools/dl_addr.ini" "${files_dir}/dl_addr.ini"
    rm -f "${files_dir}/flashtools"
    ln -sf "${ROOTDIR}/board/flashtools" "${files_dir}/flashtools"
  fi

  for file in "${files_dir}/*"; do
    ln -sf ${file} "${link_dir}"
  done

  # BL2 and LK need to be handled separatedly for now.
  if [[ -f "${files_dir}/bl2_${DRAM_SIZE}.img" ]]; then
    ln -sf "${files_dir}/bl2_${DRAM_SIZE}.img" "${link_dir}/bl2.img"
  else
    echo "bl2_${DRAM_SIZE}.img is missing!"
    found_all_files=""
  fi
  if [[ -f "${files_dir}/lk_${DRAM_SIZE}.bin" ]]; then
    ln -sf "${files_dir}/lk_${DRAM_SIZE}.bin" "${link_dir}/lk.bin"
  else
    echo "lk_${DRAM_SIZE}.bin is missing!"
    found_all_files=""
  fi

  for file in "${files[@]}"; do
    if [[ ! -f "${link_dir}/${file}" ]]; then
      echo "${file} is missing!"
      found_all_files=""
    fi
  done

  if [[ -z "${found_all_files}" ]]; then
    die "Required files are missing. Can not continue."
  fi
}

function detect_device_or_die {
  local retries="$1"; shift
  local serial="$1"; shift
  local count=0
  local output
  local found

  echo -n "Waiting up to ${retries} seconds for device ${serial}"
  for ((count=0; count<"${retries}"; count++)); do
    output=$(${FASTBOOT_CMD} devices)
    if [[ "$?" != "0" ]]; then
      die "Unable to communicate with fastboot (nonzero exit, output was [${output}])"
    fi

    if [[ "${output}" =~ "${serial}" ]]; then
      found=true
      break
    fi

    echo -n '.'
    sleep 1
  done
  echo

  if [[ -z "${found}" ]]; then
    die "Couldn't find device with serial ${serial} in fastboot output"
  else
    echo "Found device ${serial}"
  fi
}

function main {
    local usage=$(cat <<EOF
Usage: flash.sh [-d <files_dir>] [-s <serial>] [-r <detect_retries>] [-o] [-l]
  -l                  load LK into ram for fastboot mode
  -d <files_dir>      flashes files from <files_dir> (defaults to current dir)
  -s <serial>         only flashes the board with the given fastboot serial number
  -r <detect_retries> number of times to retry waiting for a device (defaults to 0)
  -o                  flash the old boards.
  -u                  flash with u-boot fastboot.
EOF
)
  local detect_retries=1          # -r <retry_count>
  local serial_number             # -s <serial>
  local files_dir                 # -d <files_dir>
  local lk_fastboot               # -l
  local flash_with_uboot=false    # -u
  local args=$(getopt hlous:d:r: $*)
  set -- $args

  for i; do
    case "$1" in
      -l)  # lk
        local lk_fastboot=true
        shift 1
        ;;

      -s)  # serial
        serial_number="$2"
        shift 2
        ;;

      -d)  # files dir
        files_dir="$2"
        shift 2
        ;;

      -r)  # retry count
        detect_retries="$2"
        if ! [[ "${detect_retries}" -gt "0" ]]; then
          die "Retry count must be a positive integer!"
        fi
        shift 2
        ;;

      -o)  # old board
        DRAM_SIZE=1G
        shift 1
        ;;

      -u)  # Flash with u-boot
        flash_with_uboot=true
        shift 1
        ;;

      --)  # end of args
        shift
        break
        ;;

      *)  # help
        die "${usage}"
        ;;

    esac
  done

  if [[ ! -x "${FASTBOOT_CMD}" ]]; then
    die "Couldn't find fastboot on your PATH -- did you install it?"
  fi

  if [[ -n "${serial_number}" ]]; then
    FASTBOOT_CMD="${FASTBOOT_CMD} -s ${serial_number}"
  fi

  if [[ -z "${files_dir}" ]]; then
    if [[ -n "${PRODUCT_OUT}" ]]; then
      files_dir="${PRODUCT_OUT}"
    else
      files_dir="${SCRIPT_DIR}"
    fi
  fi

  link_dir=$(mktemp -d)
  trap 'rm -rf ${link_dir}' EXIT
  ensure_flash_files_present "${files_dir}" "${link_dir}"

  if [[ -n "${lk_fastboot}" ]]; then
    load_fastboot "${link_dir}"
    exit 0
  fi

  if [[ -n "${serial_number}" ]]; then
    detect_device_or_die "${detect_retries}" "${serial_number}"
  fi

  flash_partitions "${link_dir}" "${flash_with_uboot}"
  try ${FASTBOOT_CMD} reboot

  echo "Flash completed."
}

main "$@"