#!/bin/bash

# Load images from a Wallaroo airgap release into an ECR registry. See usage below.

set -euo pipefail

PROG=load-images-ecr
ECR_USERNAME=AWS
region=''
app_version=''
repl_license=''
repl_username=''
load_image=''
declare -A repos
export AWS_PAGER=""

die() {
    echo "$PROG: $*" >&2
    exit 1
}

usage() {
    echo "$@" >&2
    echo >&2
    echo 'USAGE: load-images-ecr ARGUMENTS...

Required Arguments
  --region REGION                AWS region. Examples: us-east-1, us-west-2, etc.
  --wallaroo-version VERSION     Wallaroo product version (with no v).  Example: 2024.4.0
  --wallaroo-license LICENSE     Wallaroo license ID: a 28-character string.
  --wallaroo-username USERNAME   Wallaroo username: Typically an email address.

Optional Arguments
  --load-image IMAGE_URI         Load only this one image instead of all the Wallaroo ones
' >&2
    exit 1
}

command -v aws &> /dev/null || die "'aws' CLI must be installed and configured via env or ~/.aws/credentials"
command -v curl &> /dev/null || die "'curl' must be installed"
command -v docker &> /dev/null || die "'docker' must be installed"
command -v jq &> /dev/null || die "'jq' must be installed"

while [ $# -gt 0 ]; do
    case "$1" in
	--region) region=$2; shift 2;;	    
	--wallaroo-version) app_version=$2; shift 2;;
	--wallaroo-license) repl_license=$2; shift 2;;
	--wallaroo-username) repl_username=$2; shift 2;;
	--load-image) load_image=$2; shift 2;;
	*) usage "unknown argument $1";;
    esac
done

[[ -n "$repl_license" ]] || usage "missing argument --license"
[[ -n "$repl_username" ]] || usage "missing argument --username"
[[ -n "$region" ]] || usage "missing argument --region"
[[ -n "$app_version" ]] || usage "missing argument --version"

case $app_version in
    v*) usage "Do not include the 'v' in --version string";;
    *) ;;
esac

ecr_token=$(aws ecr get-login-password --region "$region")
ecr_reg_id=$(aws ecr describe-registry --region "$region" | jq -r .registryId)
ecr_hostname=${ecr_reg_id}.dkr.ecr.${region}.amazonaws.com

echo "$ecr_token" | docker login -u $ECR_USERNAME --password-stdin "$ecr_hostname"
echo "$repl_license" | docker login proxy.replicated.com -u "$repl_username" --password-stdin 

# ECR requires us to create repositories before pushing. First, generate a list of existing repositories
# to cache what is there.
echo Reading ecr
for r in $(aws ecr describe-repositories --region "$region" | jq -r '.repositories | .[] | .repositoryName'); do
    repos["$r"]=1
done

image-present() {
    aws ecr list-images --region "$region" --repository-name "$1" |
	jq -r '.imageIds | .[] | .imageTag' |
	grep -Eq "^${2}$"
}

load-image()
{
    local img="$1"
    
    # Bare repo name without host and tag
    local repo=${img##*/}
    repo=wallaroo/${repo%:*}

    # New full repo name and tag
    local ecr_repo=${ecr_hostname}/wallaroo/${img##*/}

    # just the tag
    local tag=${img/#*:/}

    # create if not in the cache
    if [[ -v repos["$repo"] ]]; then 
        echo repo "$repo" exists
    else
        aws ecr create-repository --region "$region" --repository-name "$repo"
        repos["$repo"]=1
        echo repo "$repo" created       
    fi

    if image-present "$repo" "$tag"; then
        echo "image $repo:$tag exists"
        return
    fi
    
    echo "loading $repo:$tag"

    case "$img" in
	ghcr.io/wallaroolabs/*) img="proxy.replicated.com/proxy/wallaroo/$img";;

	# docker.io/foo:123 should be just foo:123 through the proxy
	docker.io/*) img="proxy.replicated.com/proxy/anonymous/$(basename $img)";;

	# Just pull the image directly. This host is unsupported by replicated proxy.
	cr.fluentbit.io/*) ;;
	nvcr.io/*) ;;
	registry.k8s.io/*) ;;

	# Supported by the proxy, use that
	*) img="proxy.replicated.com/anonymous/$img";;
    esac

    docker pull -q "$img"
    docker tag "$img" "$ecr_repo"
    docker push -q "$ecr_repo"
}

if [[ "$load_image" ]]; then
    echo "Loading only $load_image"
    load-image "$load_image"
else
    echo Loading all Wallaroo images

    # Replicated images
    curl -s -H "Authorization: $repl_license" "https://replicated.app/v3/app/wallaroo/images?version=$app_version" |
	jq -r '.[] | .image' |
	while read -r img; do
	    load-image "$img"
	done
fi