#!/bin/bash -ex
usage () {
  echo "Create a Mixed Build archive with the given system and device archives."
  echo
  echo "Usage: $0 [-v <vendor_version>] [-m <modify_system_image_path>]"
  echo "    [-t <prebuilt_otatools_path>] [-p <override_vbmeta_image_path>]"
  echo "    [-b <override_boot_image_path>]"
  echo "    [-s] [-d] system_build_dir device_build_dir out_dir"
  echo
  echo "Options -v, -m, -t, -p, -b, -s, -d must precede positional arguments."
  echo
  echo "vendor_version is the version of the vendor image when Keymaster v3"
  echo "    related modifications to the system image is necessary. Optional."
  echo "    eg. 8.1.0 for a mixed build of GSI and O-MR1 vendor image."
  echo "modify_system_image_path is the path to the script that modifies the"
  echo "    system image, needed for Keymaster v3. Optional."
  echo "prebuilt_otatools_path is the path to otatools.zip file that has all"
  echo "    required host binaries to modify system image. It also must include"
  echo "    VINTF check tool to verify the compatibility of the given images."
  echo "    Optional."
  echo "override_vbmeta_image_path is the path to a vbmeta.img to use"
  echo "    to override the existing vbmeta.img of device. Optional."
  echo "override_boot_image_path is the path to a boot imgage to use to"
  echo "    override the existing boot.img of device. Optional."
  echo "-s is used to fetch and flash both product.img and system.img from"
  echo "    the system_build_dir for devices with a product partition."
  echo "    product.img will be removed if system_build_dir does not have"
  echo "    product.img when -s option is declared."
  echo "    By default, only system.img is flashed to the target device for"
  echo "    independent system update. No parameter required. Optional"
  echo "-d is used to skip vbmeta.img replacement in its entirety and the"
  echo "    one from the device build (if exists) is left untouched"
  echo "system_build_dir is the path to the system build"
  echo "    eg. aosp_arm64-userdebug."
  echo "device_build_dir is the path to the device build"
  echo "    eg. sailfish-user."
  echo "out_dir is the path to where the new build will be placed."
}

# Print error message and exit.
# Usage: exit_badparam message
#
# message is a string to be displayed before exit.
exit_badparam () {
  echo "ERROR: $1" >&2
  usage
  exit 1
}

cleanup_and_exit () {
  readonly result="$?"
  rm -rf "$TEMP_DIR"
  exit "$result"
}

trap cleanup_and_exit EXIT

while getopts :v:m:p:b:t:sd opt; do
  case "$opt" in
    v)
      readonly VENDOR_VERSION="$OPTARG"
      ;;
    m)
      readonly MODIFY_SYSTEM_SCRIPT="$OPTARG"
      ;;
    p)
      readonly OVERRIDE_VBMETA_IMAGE_PATH="$OPTARG"
      ;;
    b)
      readonly OVERRIDE_BOOT_IMAGE_PATH="$OPTARG"
      ;;
    t)
      readonly OTATOOLS_ZIP="$OPTARG"
      ;;
    s)
      readonly INCLUDE_PRODUCT=true
      ;;
    d)
      readonly SKIP_VBMETA_REPLACE=true
      ;;
    \?)
      exit_badparam "Invalid options: -"$OPTARG""
      ;;
    :)
      exit_badparam "Option -"$OPTARG" requires an argument."
      ;;
  esac
done

if [[ -z "${VENDOR_VERSION+x}" && ! -z "${MODIFY_SYSTEM_SCRIPT+x}" ]] || \
  [[ ! -z "${VENDOR_VERSION+x}" && -z "${MODIFY_SYSTEM_SCRIPT+x}" ]]; then
  exit_badparam "Options -v and -m must be set together."
fi

shift "$((OPTIND-1))"

if [[ $# -lt 3 ]]; then
  exit_badparam "Unexpected number of arguments"
fi

readonly SYSTEM_DIR="$1"
readonly DEVICE_DIR="$2"
readonly DIST_DIR="$3"
readonly TEMP_DIR="$(mktemp -d /tmp/"$(basename $0)"_XXXXXXXX)"

readonly SYSTEM_TARGET_FILES_ARCHIVE="$(find "$SYSTEM_DIR" -name "*-target_files-*.zip" -print)"
if [[ ! -f "$SYSTEM_TARGET_FILES_ARCHIVE" ]]; then
  exit_badparam "Could not find system target files archive in $SYSTEM_DIR."
fi

readonly DEVICE_ARCHIVE="$(find "$DEVICE_DIR" -name "*-img-*.zip" -print)"
if [[ ! -f "$DEVICE_ARCHIVE" ]]; then
  exit_badparam "Could not find device img archive in $DEVICE_DIR."
fi

readonly DEVICE_TARGET_FILES_ARCHIVE="$(find "$DEVICE_DIR" -name "*-target_files-*.zip" -print)"
if [[ ! -f "$DEVICE_TARGET_FILES_ARCHIVE" ]]; then
  exit_badparam "Could not find device target_files archive in $DEVICE_DIR."
fi

if [[ ! -z "${MODIFY_SYSTEM_SCRIPT+x}" && ! -f "$MODIFY_SYSTEM_SCRIPT" ]]; then
  exit_badparam "Script not found: "$MODIFY_SYSTEM_SCRIPT""
fi

if [[ ! -z "${OVERRIDE_VBMETA_IMAGE_PATH+x}" && ! -f "$OVERRIDE_VBMETA_IMAGE_PATH" ]]; then
  exit_badparam "Specified vbmeta.img not found: "$OVERRIDE_VBMETA_IMAGE_PATH""
fi

if [[ ! -z "${OVERRIDE_BOOT_IMAGE_PATH+x}" && ! -f "$OVERRIDE_BOOT_IMAGE_PATH" ]]; then
  exit_badparam "Specified boot image not found: "$OVERRIDE_BOOT_IMAGE_PATH""
fi

readonly DEVICE_ARTIFACTS_DIR="$TEMP_DIR"/device_archive_artifacts
readonly DEVICE_IMAGES_DIR="$DEVICE_ARTIFACTS_DIR"/IMAGES
readonly SYSTEM_ARTIFACTS_DIR="$TEMP_DIR"/system_artifacts
readonly SYSTEM_IMAGES_DIR="$SYSTEM_ARTIFACTS_DIR"/IMAGES
readonly OTATOOLS_DIR="$TEMP_DIR"/otatools

readonly SPL_PROPERTY_NAME="ro.build.version.security_patch"
readonly SYSTEM_BUILD_PROP="SYSTEM/build.prop"

declare -a EXTRACT_SYSTEM_FILE_LIST
EXTRACT_SYSTEM_FILE_LIST=(
  IMAGES/system.img \
  IMAGES/vbmeta.img \
  "$SYSTEM_BUILD_PROP" \
)

declare -a EXTRACT_VINTF_SYSTEM_FILE_LIST
EXTRACT_VINTF_SYSTEM_FILE_LIST=(
  "$SYSTEM_BUILD_PROP" \
)

declare -a EXTRACT_DEVICE_FILE_LIST
EXTRACT_DEVICE_FILE_LIST=(
  */build.prop \
  META/* \
)

declare -A SYSTEM_SEARCH_PATH
SYSTEM_SEARCH_PATH=( \
  [/system]="SYSTEM" \
  [/product]="PRODUCT SYSTEM/product" \
  [/system_ext]="SYSTEM_EXT SYSTEM/system_ext" \
)

declare -A DEVICE_SEARCH_PATH
# Mixed build will not have /vendor to SYSTEM/vendor case
DEVICE_SEARCH_PATH=( \
  [/vendor]="VENDOR" \
  [/odm]="ODM VENDOR/odm" \
)

###
# Uncompress otatools.zip and get vintf file list.
if [[ ! -f "$OTATOOLS_ZIP" ]]; then
  echo "WARNING: otatools.zip is missing. Add \"-t otatools.zip\" to enable checkvintf"
else
  readonly OTATOOLS_AVAILABLE=true
  # Uncompress otatools
  mkdir -p "$OTATOOLS_DIR"
  unzip "$OTATOOLS_ZIP" bin/* lib64/* -d "$OTATOOLS_DIR"
  # Set paths for using prebuilt host binaries.
  export PATH="$OTATOOLS_DIR"/bin:"$PATH"
  export LD_LIBRARY_PATH="$OTATOOLS_DIR"/lib64:"$LD_LIBRARY_PATH"

  # Add vintf file to extract file list
  declare -a VINTF_DUMP_FILE_LIST
  VINTF_DUMP_FILE_LIST=( "$(checkvintf --dump-file-list)" )

  for vintf_file_list in ${VINTF_DUMP_FILE_LIST[*]}; do
    if [[ "$vintf_file_list" == */ ]]; then
      vintf_file_list="$vintf_file_list"\*
      # Create system vintf file list for system target files archive
      for system_dir in "${!SYSTEM_SEARCH_PATH[@]}"; do
        if [[ "$vintf_file_list" == "$system_dir"/* ]]; then
          for search_dir in ${SYSTEM_SEARCH_PATH["$system_dir"]}; do
            search_file=${vintf_file_list/$system_dir/$search_dir}
            unzip -l "$SYSTEM_TARGET_FILES_ARCHIVE" "$search_file" > /dev/null && \
              EXTRACT_VINTF_SYSTEM_FILE_LIST+=( "$search_file" )
          done
          break
        fi
      done
      # Create device vintf file list for device target files archive
      for device_dir in "${!DEVICE_SEARCH_PATH[@]}"; do
        if [[ "$vintf_file_list" == "$device_dir"/* ]]; then
          for search_dir in ${DEVICE_SEARCH_PATH["$device_dir"]}; do
            search_file=${vintf_file_list/$device_dir/$search_dir}
            unzip -l "$DEVICE_TARGET_FILES_ARCHIVE" "$search_file" > /dev/null && \
              EXTRACT_DEVICE_FILE_LIST+=( "$search_file" )
          done
          break
        fi
      done
     fi
  done
fi

###
# Uncompress the system archives.
if [[ "$INCLUDE_PRODUCT" == true ]]; then
  unzip -l "$SYSTEM_TARGET_FILES_ARCHIVE" | grep -q IMAGES/product.img &&
  EXTRACT_SYSTEM_FILE_LIST+=(IMAGES/product.img)
fi

mkdir -p "$SYSTEM_ARTIFACTS_DIR"
# Get system images.
unzip "$SYSTEM_TARGET_FILES_ARCHIVE" "${EXTRACT_SYSTEM_FILE_LIST[@]}" \
  -d "$SYSTEM_ARTIFACTS_DIR"

###
# Uncompress the device archives.
mkdir -p "$DEVICE_IMAGES_DIR"
# Get device images.
unzip "$DEVICE_ARCHIVE" -d "$DEVICE_IMAGES_DIR"
# Get the device meta data.
unzip "$DEVICE_TARGET_FILES_ARCHIVE" "${EXTRACT_DEVICE_FILE_LIST[@]}" \
  -d "$DEVICE_ARTIFACTS_DIR"

###
# Modify system.img if vendor version is provided.
if [[ ! -z "${VENDOR_VERSION+x}" ]]; then
  # Create copy of system target files package that can be modified
  # since the original $SYSTEM_TARGET_FILES_ARCHIVE is a symlink to
  # prebuilt files in cache
  cp "$SYSTEM_TARGET_FILES_ARCHIVE" "$TEMP_DIR"
  readonly COPY_SYSTEM_TARGET_FILES_ARCHIVE="$TEMP_DIR"/"$(basename "$SYSTEM_TARGET_FILES_ARCHIVE")"

  # Check compatibility of security patch level
  readonly SYSTEM_SPL=$(sed -n -r "s/^"$SPL_PROPERTY_NAME"=(.*)$/\1/p" "$SYSTEM_ARTIFACTS_DIR"/"$SYSTEM_BUILD_PROP")
  readonly VENDOR_SPL=$(sed -n -r "s/^"$SPL_PROPERTY_NAME"=(.*)$/\1/p" "$DEVICE_ARTIFACTS_DIR"/"$SYSTEM_BUILD_PROP")
  declare -a args
  args=(-v "$VENDOR_VERSION" "$COPY_SYSTEM_TARGET_FILES_ARCHIVE")
  if [[ "$SYSTEM_SPL" != "$VENDOR_SPL" ]]; then
    echo "Security patch level mismatch detected..."
    echo "  SPL of system: "$SYSTEM_SPL""
    echo "  SPL of vendor: "$VENDOR_SPL""
    args+=("$VENDOR_SPL")
  fi
  "$MODIFY_SYSTEM_SCRIPT" "${args[@]}"
  # Replace system.img with newly modified system.img
  unzip -o "$COPY_SYSTEM_TARGET_FILES_ARCHIVE" IMAGES/system.img -d "$SYSTEM_ARTIFACTS_DIR"
fi

# Check vintf
if [[ "$OTATOOLS_AVAILABLE" == true ]]; then
  # Overwrite VINTF system matrix to device artifacts dir
  unzip -o "$SYSTEM_TARGET_FILES_ARCHIVE" "${EXTRACT_VINTF_SYSTEM_FILE_LIST[@]}" \
    -d "$DEVICE_ARTIFACTS_DIR"
  check_target_files_vintf "$DEVICE_ARTIFACTS_DIR"
fi

###
# Overwrite artifacts in the device archive to create the Mixed Build artifacts.
cp "$SYSTEM_IMAGES_DIR"/system.img "$DEVICE_IMAGES_DIR"/
if [[ "$INCLUDE_PRODUCT" == true ]]; then
  if [[ -f "$SYSTEM_IMAGES_DIR"/product.img ]]; then
    cp "$SYSTEM_IMAGES_DIR"/product.img "$DEVICE_IMAGES_DIR"/
  else
    rm -f "$DEVICE_IMAGES_DIR"/product.img
    # Removed product partition from required partition list
    sed -i "/partition-exists=product$/d" "$DEVICE_IMAGES_DIR"/android-info.txt
  fi
fi

if [[ "$SKIP_VBMETA_REPLACE" == true ]]; then
    # Totally skip the vbmeta.img replacement
    echo "Skip vbmeta.img replacement."
else
    # Only override vbmeta if it is already present since fastboot update will try
    # to flash whatever is in the archive.
    if [[ -f "$DEVICE_IMAGES_DIR"/vbmeta.img ]]; then
      readonly VBMETA_IMAGE_PATH="${OVERRIDE_VBMETA_IMAGE_PATH:-"$SYSTEM_IMAGES_DIR"/vbmeta.img}"
      cp "$VBMETA_IMAGE_PATH" "$DEVICE_IMAGES_DIR"/
    fi
fi

# Override boot.img with the provided boot image file since fastboot update cmd
# will try to flash boot.img in the archive.
if [[ ! -z "${OVERRIDE_BOOT_IMAGE_PATH+x}" && -f "$DEVICE_IMAGES_DIR"/boot.img ]]; then
  cp "$OVERRIDE_BOOT_IMAGE_PATH" "$DEVICE_IMAGES_DIR"/boot.img
fi

###
# Create the Mixed Build archive.
(
  cd "$DEVICE_IMAGES_DIR"
  zip -r mixed.zip ./*
)

###
# Archive the artifacts.
if [ -n "$DIST_DIR" ]; then
  mkdir -p "$DIST_DIR" || true
fi
# Archive all the device artifacts.
rsync --archive --verbose --copy-links --exclude='logs' \
  "$DEVICE_DIR"/* "$DIST_DIR"
# Overwrite the image archive with the Mixed Build archive.
OUT_ARCHIVE="$DIST_DIR"/"$(basename $DEVICE_ARCHIVE)"
cp "$DEVICE_IMAGES_DIR"/mixed.zip "$OUT_ARCHIVE"
# Overwrite android-info.txt with the updated one.
cp "$DEVICE_IMAGES_DIR"/android-info.txt "$DIST_DIR"/