#!/usr/bin/env bash
set -euo pipefail

IMAGES_FILE="scripts/bootstrap/harbor-bootstrap-images.txt"
BUNDLE_FILE="artifacts/harbor-bootstrap-v2.14.1-arm64.tar.zst"
DOCKER_CONFIG_PATH=""
PLATFORM="linux/arm64"
ZSTD_LEVEL="${ZSTD_LEVEL:-19}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --images-file)
      IMAGES_FILE="${2:?missing images file}"
      shift 2
      ;;
    --bundle-file)
      BUNDLE_FILE="${2:?missing bundle file}"
      shift 2
      ;;
    --docker-config)
      DOCKER_CONFIG_PATH="${2:?missing docker config path}"
      shift 2
      ;;
    --platform)
      PLATFORM="${2:?missing platform}"
      shift 2
      ;;
    --zstd-level)
      ZSTD_LEVEL="${2:?missing zstd compression level}"
      shift 2
      ;;
    -h|--help)
      cat <<USAGE
Usage: scripts/build_harbor_bootstrap_bundle.sh [--images-file <path>] [--bundle-file <path>] [--docker-config <path>] [--platform <linux/arm64>] [--zstd-level <level>]
USAGE
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      exit 1
      ;;
  esac
done

if [[ -n "${DOCKER_CONFIG_PATH}" ]]; then
  export DOCKER_CONFIG="${DOCKER_CONFIG_PATH}"
fi

mapfile -t IMAGES < <(grep -v '^[[:space:]]*#' "${IMAGES_FILE}" | sed '/^[[:space:]]*$/d')
if [[ ${#IMAGES[@]} -eq 0 ]]; then
  echo "No images found in ${IMAGES_FILE}" >&2
  exit 1
fi

source_image_for_alias() {
  local image="$1"
  local tag="${image##*:}"
  case "${image}" in
    registry.bstein.dev/infra/longhorn-engine:*) echo "docker.io/longhornio/longhorn-engine:${tag}" ;;
    registry.bstein.dev/infra/longhorn-manager:*) echo "docker.io/longhornio/longhorn-manager:${tag}" ;;
    registry.bstein.dev/infra/longhorn-ui:*) echo "docker.io/longhornio/longhorn-ui:${tag}" ;;
    registry.bstein.dev/infra/longhorn-instance-manager:*) echo "docker.io/longhornio/longhorn-instance-manager:${tag}" ;;
    registry.bstein.dev/infra/longhorn-share-manager:*) echo "docker.io/longhornio/longhorn-share-manager:${tag}" ;;
    registry.bstein.dev/infra/longhorn-backing-image-manager:*) echo "docker.io/longhornio/backing-image-manager:${tag}" ;;
    registry.bstein.dev/infra/longhorn-support-bundle-kit:*) echo "docker.io/longhornio/support-bundle-kit:${tag}" ;;
    registry.bstein.dev/infra/longhorn-csi-attacher:*) echo "registry.k8s.io/sig-storage/csi-attacher:${tag}" ;;
    registry.bstein.dev/infra/longhorn-csi-provisioner:*) echo "registry.k8s.io/sig-storage/csi-provisioner:${tag}" ;;
    registry.bstein.dev/infra/longhorn-csi-node-driver-registrar:*) echo "registry.k8s.io/sig-storage/csi-node-driver-registrar:${tag}" ;;
    registry.bstein.dev/infra/longhorn-csi-resizer:*) echo "registry.k8s.io/sig-storage/csi-resizer:${tag}" ;;
    registry.bstein.dev/infra/longhorn-csi-snapshotter:*) echo "registry.k8s.io/sig-storage/csi-snapshotter:${tag}" ;;
    registry.bstein.dev/infra/longhorn-livenessprobe:*) echo "registry.k8s.io/sig-storage/livenessprobe:${tag}" ;;
    *) echo "${image}" ;;
  esac
}

pull_or_tag_image() {
  local image="$1"
  local source_image
  if docker image inspect "${image}" >/dev/null 2>&1; then
    echo "Using cached ${image}" >&2
    return 0
  fi
  echo "Pulling ${image}" >&2
  if docker pull --platform "${PLATFORM}" "${image}" >/dev/null; then
    return 0
  fi
  source_image="$(source_image_for_alias "${image}")"
  if [[ "${source_image}" == "${image}" ]]; then
    return 1
  fi
  echo "Pulling ${source_image} for ${image}" >&2
  docker pull --platform "${PLATFORM}" "${source_image}" >/dev/null
  docker tag "${source_image}" "${image}"
}

mkdir -p "$(dirname "${BUNDLE_FILE}")"
for image in "${IMAGES[@]}"; do
  pull_or_tag_image "${image}"
done

tmp_bundle="${BUNDLE_FILE}.tmp"
rm -f "${tmp_bundle}"
docker save "${IMAGES[@]}" | zstd -T0 -"${ZSTD_LEVEL}" -o "${tmp_bundle}"
mv "${tmp_bundle}" "${BUNDLE_FILE}"
echo "Wrote ${BUNDLE_FILE}" >&2