#!/bin/bash
#
# Copyright 2020 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
EMMC_SIZE=8G

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

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

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

  if [[ ! -z "${overwrite_home}" ]]; then
    try ${FASTBOOT_CMD} erase mmc0
    try ${FASTBOOT_CMD} erase mmc0boot0
    try ${FASTBOOT_CMD} erase mmc0boot1
  fi
  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"
    "fip.bin"
    "partition-table.img"
    "rootfs_arm64.img"
    "u-boot-env.bin"
  )

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

  # BL2 and partition tables 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}/partition-table_${EMMC_SIZE}.img" ]]; then
    ln -sf "${files_dir}/partition-table_${EMMC_SIZE}.img" "${link_dir}/partition-table.img"
  else
    echo "partition-table_${EMMC_SIZE}.img 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]
  -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.
EOF
)
  local detect_retries=1          # -r <retry_count>
  local serial_number             # -s <serial>
  local files_dir                 # -d <files_dir>
  local lk_fastboot               # -l
  local overwrite_home            # -H
  local args=$(getopt hHlous: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
        EMMC_SIZE=4G
        shift 1
        ;;

      -H)  # overwrite home
        overwrite_home=true
        shift 1
        ;;

      -u)  # Flash with u-boot (deprecated)
        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 "${serial_number}" ]]; then
    detect_device_or_die "${detect_retries}" "${serial_number}"
  fi

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

  echo "Flash completed."
}

main "$@"