#!/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=''
create_kots=false
create_wallaroo=false
registry_namespace=wallaroo/
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.
  --registry-namespace NAMESPACE   Desired namespace in private registry. Often "wallaroo".

Optional Arguments
  --load-image IMAGE_URI           Load only this one image instead of all the Wallaroo ones
  --create-kots                    Create KOTS repositories for use before a KOTS install. Nothing is pushed.
  --create-wallaroo                Create Wallaroo repositories. Nothing is pushed.
' >&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;;
	--registry-namespace) registry_namespace=$2; shift 2;;
	--load-image) load_image=$2; shift 2;;
	--create-kots) create_kots=true; shift;;
	--create-wallaroo) create_wallaroo=true; shift;;
	*) 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"
[[ -n "$registry_namespace" ]]	|| usage "missing argument --registry-namespace"

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

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

if $create_kots || $create_wallaroo; then
    echo skipping docker logins
else
    ecr_token=$(aws ecr get-login-password --region "$region")
    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 
fi

create-kots() {
    for img in local-volume-provider minio kotsadm kotsadm-migrations dex rqlite; do
	create-repository "${registry_namespace}${img}"
    done
}

create-wallaroo() {
    local repo
    curl -s -H "Authorization: $repl_license" "https://replicated.app/v3/app/wallaroo/images?version=$app_version" |
	jq -r '.[] | .image' |
	while read -r img; do
	    repo=${img##*/}
	    repo=${registry_namespace}${repo%:*}
	    create-repository "$repo"
	done
}    

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

# Only needed for ECR. Attempt to create repo and ignore dup errors.
create-repository() {
    local repo=$1
    local temp=/tmp/load-images-error
    set +e
    aws ecr create-repository --region "$region" --repository-name "$repo" >& $temp
    case $? in
	0) echo repo "$repo" created;;
	254) echo repo "$repo" exists;;
	*) die "Error creating registry: $(cat $temp)"
    esac
    set -e
}

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

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

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

    create-repository "$repo"	# may grumble but can be ignored

    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 $create_kots; then
    echo Creating the KOTS repositories
    create-kots
fi

if $create_wallaroo; then
    echo Creating the Wallaroo repositories
    create-wallaroo
fi

if $create_wallaroo || $create_kots; then
    exit 0
fi

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