#!/bin/bash

set -e -o pipefail
shopt -qs failglob

SKIP_MANIFEST="${SKIP_MANIFEST:-'false'}"
ONLY_MANIFEST="${ONLY_MANIFEST:-'false'}"

for opt in "$@"; do
   optarg="$(expr "${opt}" : '[^=]*=\(.*\)')"
   case "${opt}" in
      --version=*) VERSION="${optarg}" ;;
      --registry=*) REGISTRY="${optarg}" ;;
      --sdk-name=*) SDK_NAME="${optarg}" ;;
      --short-sha=*) SHORT_SHA="${optarg,,}" ;;
      --skip-manifest=*) SKIP_MANIFEST="${optarg,,}" ;;
      --only-manifest=*) ONLY_MANIFEST="${optarg,,}" ;;
   esac
done

for required_arg in VERSION REGISTRY SDK_NAME SHORT_SHA; do
  [ -z "${!required_arg}" ] && echo "${required_arg} not set!" >&2 && exit 1
done

ECR_OPERATIONS=0
MIN_ECR_OPERATIONS=2
if [[ "${SKIP_MANIFEST}" == "true" ]] && [[ "${ONLY_MANIFEST}" == "true" ]] ; then
  echo "nothing to do" >&2 && exit 2
elif [[ "${SKIP_MANIFEST}" == "true" ]] || [[ "${ONLY_MANIFEST}" == "true" ]] ; then
  ((--MIN_ECR_OPERATIONS))
fi

update_image() {
  set -u

  local local_image="${1:?}"
  local remote_image="${2:?}"
  local arch="${3:?}"

  local host_arch
  host_arch="$(uname -m)"

  local docker_arch
  case "${host_arch}" in
    x86_64) docker_arch="amd64" ;;
    aarch64) docker_arch="arm64" ;;
    *) echo "unknown host arch ${host_arch}" >&2 && exit 1 ;;
  esac

  local alt_arch
  case "${host_arch}" in
    x86_64) alt_arch="aarch64" ;;
    aarch64) alt_arch="x86_64" ;;
    *) echo "unknown host arch ${host_arch}" >&2 && exit 1 ;;
  esac

  local alt_docker_arch
  case "${docker_arch}" in
    amd64) alt_docker_arch="arm64" ;;
    arm64) alt_docker_arch="amd64" ;;
  esac

  local local_host_arch_image="${local_image}-${SHORT_SHA}-${host_arch}"
  local remote_host_arch_image="${remote_image}-${SHORT_SHA}-${host_arch}"
  local remote_alt_arch_image="${remote_image}-${SHORT_SHA}-${alt_arch}"

  if [[ "${ONLY_MANIFEST}" != "true" ]]; then
    if ! docker image inspect "${local_host_arch_image}" >/dev/null 2>&1 ; then
      echo "did not find ${local_host_arch_image}, skipping it ..." >&2
      return
    fi

    if ! docker tag "${local_host_arch_image}" "${remote_host_arch_image}" ; then
      echo "failed to tag ${remote_host_arch_image}" >&2 && exit 1
    fi

    # push the image, if it exists locally
    if ! docker push "${remote_host_arch_image}" ; then
      echo "failed to push ${remote_host_arch_image}" >&2 && exit 1
    fi
    ((++ECR_OPERATIONS))
  fi

  # return early to skip manifest generation steps
  [[ "${SKIP_MANIFEST}" == "true" ]] && return

  # clean up any cached local copy of the manifest
  manifest_dir="${HOME}/.docker/manifests"
  mkdir -p "${manifest_dir}"
  find "${manifest_dir}" -type f -delete
  find "${manifest_dir}" -mindepth 1 -type d -delete

  if ! docker manifest inspect "${remote_alt_arch_image}" >/dev/null 2>&1 ; then
    echo "could not find ${remote_alt_arch_image}, skipping manifest creation" >&2
    return
  fi

  local image="${remote_image##*/}"
  local repo="${image%:*}"

  # clean up any remote (ECR) copy of the manifest
  if [[ "${REGISTRY}" =~ ^public\.ecr\.aws\/ ]]; then
    aws ecr-public batch-delete-image \
      --repository-name "${repo}" \
      --image-ids imageTag="${VERSION}" \
      --region "us-east-1"
  elif [[ "${REGISTRY}" =~ ^[0-9]+\.dkr\.ecr\.([a-z0-9-]+)\.amazonaws\.com$ ]]; then
    aws ecr batch-delete-image \
      --repository-name "${repo}" \
      --image-ids imageTag="${VERSION}" \
      --region "${BASH_REMATCH[1]}"
  fi

  if ! docker manifest create "${remote_image}" \
      "${remote_host_arch_image}" "${remote_alt_arch_image}" ; then
    echo "failed to create manifest ${remote_image}" >&2 && exit 1
  fi

  if ! docker manifest annotate "${remote_image}" \
      "${remote_host_arch_image}" --arch "${docker_arch}" ; then
    echo "failed to annotate manifest ${remote_image} with ${remote_host_arch_image}" >&2 && exit 1
  fi

  if ! docker manifest annotate "${remote_image}" \
      "${remote_alt_arch_image}" --arch "${alt_docker_arch}" ; then
    echo "failed to annotate manifest ${remote_image} with ${remote_alt_arch_image}" >&2 && exit 1
  fi

  # push the manifest and remove the local cache
  if ! docker manifest push --purge "${remote_image}" ; then
    echo "failed to push ${remote_image}" >&2 && exit 1
  fi
    ((++ECR_OPERATIONS))
}

for arch in x86_64 aarch64 ; do
  local_sdk="bottlerocket/sdk-${arch}:${VERSION}"
  remote_sdk="${REGISTRY}/${SDK_NAME}-sdk-${arch}:${VERSION}"
  update_image "${local_sdk}" "${remote_sdk}" "${arch}"

  local_toolchain="bottlerocket/toolchain-${arch}:${VERSION}"
  remote_toolchain="${REGISTRY}/${SDK_NAME}-toolchain-${arch}:${VERSION}"
  update_image "${local_toolchain}" "${remote_toolchain}" "${arch}"
done

if [[ ${ECR_OPERATIONS} -lt ${MIN_ECR_OPERATIONS} ]]; then
  echo "no containers or manifests were pushed" >&2 && exit 1
else
  exit 0
fi