#!/usr/bin/env bash
# Copyright 2024 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

#### CATEGORY=Other
### Flash Fuchsia onto a Kola board that has an unmodified bootloader.

## usage: fx flash-kola

# Bash strict mode.
set -eo pipefail

source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/../lib/vars.sh || exit $?
fx-config-read

# Use the in-tree fastboot.
readonly FASTBOOT="${FUCHSIA_DIR}/prebuilt/third_party/fastboot/fastboot"


# TODO(b/373727276): Change to use `ffx target flash` instead.
FLASH_SOURCE="$(fx-command-run get-flash-source)"
PRODUCT_BUNDLE="$(fx get-build-dir)/${FLASH_SOURCE#product-bundle:}"
ARTIFACTS="${PRODUCT_BUNDLE}/product_bundle.json"
DTBO_PATH="${PRODUCT_BUNDLE}/$(fx-command-run jq --raw-output '.system_a[] | select(.name=="dtbo-a") | .path' "${ARTIFACTS}")"
FVM_PATH="${PRODUCT_BUNDLE}/$(fx-command-run jq --raw-output '.system_a[] | select(.name=="fxfs.fastboot") | .path' "${ARTIFACTS}")"
ZBI_PATH="${PRODUCT_BUNDLE}/$(fx-command-run jq --raw-output '.system_a[] | select(.name=="zircon-a") | .path' "${ARTIFACTS}")"

RECOVERY_PATH=""
if [[ -d "${PRODUCT_BUNDLE}/system_r" ]]; then
  RECOVERY_PATH="${PRODUCT_BUNDLE}/$(fx-command-run jq --raw-output '.system_r[] | select(.name=="zircon-a") | .path' "${ARTIFACTS}")"
fi

readonly FVM_PATH ZBI_PATH DTBO_PATH RECOVERY_PATH

if [[ -z $ANDROID_SERIAL ]]; then
  readonly DEVICE_COUNT=$(${FASTBOOT} devices | awk NF | wc -l)
  if [[ ${DEVICE_COUNT} -gt 1 ]]; then
    fx-error "Detected ${DEVICE_COUNT} devices in fastboot mode."
    fx-error "Make sure there is only one device or set ANDROID_SERIAL."
    exit 1
  fi
fi

BOOTLOADER="$(${FASTBOOT} getvar product 2>&1 | grep '^product:' | awk '{print $2}')"
# Only attempt to set bootmode if the bootloader is not 'sorrel' and mode is not 'factory'
if [[ "${BOOTLOADER}" != "sorrel" ]]; then
  # Ensure bootmode is set to factory before flashing partitions.
  fx-info "Checking current bootmode..."
  # Capture output and potential errors. Use awk to extract the value after 'bootmode='.
  # Store the exit code separately.
  if BOOTMODE_OUTPUT=$(${FASTBOOT} oem get_config bootmode 2>&1); then
    BOOTMODE_EXIT_CODE=0
  else
    BOOTMODE_EXIT_CODE=$?
  fi

  INITIAL_BOOTMODE="" # Initialize to empty

  if [[ ${BOOTMODE_EXIT_CODE} -eq 0 ]]; then
    CURRENT_BOOTMODE=$(echo "${BOOTMODE_OUTPUT}" | awk '/bootmode:/ {print $3}')
    INITIAL_BOOTMODE="${CURRENT_BOOTMODE}" # Store the initial bootmode
    if [[ "${CURRENT_BOOTMODE}" != "factory" ]]; then
      fx-info "Current bootmode is '${CURRENT_BOOTMODE}', setting it to 'factory'."
      ${FASTBOOT} oem set_config bootmode factory || {
        fx-warn "Failed to set bootmode to factory. Attempting to flash anyway."
      }
    else
      fx-info "Bootmode: 'factory'."
    fi
  else
    fx-warn "Failed to get current bootmode (Exit code: ${BOOTMODE_EXIT_CODE}). Output: ${BOOTMODE_OUTPUT}"
    fx-warn "Proceeding with flashing, but bootmode might not be 'factory'."
  fi
else
  fx-info "Bootloader is '${BOOTLOADER}'."
fi

${FASTBOOT} flash boot_a "$ZBI_PATH"
${FASTBOOT} flash boot_b "$ZBI_PATH"
${FASTBOOT} flash dtbo_a "$DTBO_PATH"
${FASTBOOT} flash dtbo_b "$DTBO_PATH"
if [[ -n "$RECOVERY_PATH" ]]; then
  ${FASTBOOT} flash vendor_boot_a "$RECOVERY_PATH"
  ${FASTBOOT} flash vendor_boot_b "$RECOVERY_PATH"
fi

# Reboot back into the bootloader.
# Some builds put firmware in the boot image, in which case we reboot here to
# pick up the new firmware we just flashed. For builds that don't put firmware
# in boot, this is still fine and just triggers a reboot now rather than at the
# end of flashing.
fx-info "Rebooting to bootloader"
${FASTBOOT} reboot bootloader
# Pause to make sure fastboot actually drops the device; sometimes without this
# the next fastboot command will execute before the host realizes the device is
# gone which causes the command to fail.
sleep 4
# Attempt to flash SSH keys. This is only supported on devices running custom
# Fuchsia firmware, so if it doesn't work just log a warning.
function push_ssh_keys {
  AUTH_KEYS="$(get-ssh-authkeys)" || {
    fx-warn "Failed to locate host SSH keys"
    return 1
  }

  # Split this into two commands for clarity on which command fails.
  # Staging a file failing is very different from the oem command failing
  ${FASTBOOT} stage "${AUTH_KEYS}" || {
    fx-warn "Failed to upload authorized keys"
    return 1
  }

  ${FASTBOOT} oem add-staged-bootloader-file ssh.authorized_keys || {
    fx-warn "Failed to stage SSH keys to fastboot"
    fx-warn "Device bootloader may not support it"
    return 1
  }
}

if [[ "${BOOTLOADER}" == "sorrel" ]]; then
  if push_ssh_keys; then
    fx-info "Successfully provisioned SSH keys via bootloader"
  else
    fx-warn "Skipping bootloader provisioning of SSH keys"
  fi
else
  fx-info "Skipping bootloader provisioning of SSH keys because bootloader is ${BOOTLOADER}"
fi

# Flash the super partition last; some Kola devices exhibit USB instability
# after flashing large super partition data so we want to do as much as
# possible before this while the USB is more reliable (b/405436515).
${FASTBOOT} flash super "$FVM_PATH"

# Attempt to restore the initial bootmode if it was changed
if [[ -n "${INITIAL_BOOTMODE}" && "${INITIAL_BOOTMODE}" != "factory" ]]; then
  fx-info "Attempting to restore initial bootmode to '${INITIAL_BOOTMODE}'..."
  # Check the exit status of the set_config command
  if ${FASTBOOT} oem set_config bootmode "${INITIAL_BOOTMODE}"; then
    fx-info "Successfully restored bootmode to '${INITIAL_BOOTMODE}'."
  else
    fx-warn "Failed to restore initial bootmode to '${INITIAL_BOOTMODE}'. Device might remain in 'factory' mode."
  fi
fi

# SSH keys require `fastboot continue`, since they are stored in RAM and would
# get lost in a reboot. We rebooted earlier to ensure a complete boot chain
# using the new images.
${FASTBOOT} continue || {
  fx-warn "Fastboot got a communication error, but this may be OK"
  fx-warn "If the device is now booting then flash succeeded (b/405436515)"
}

fx-info "Done flashing Kola device."