#!/usr/bin/env bash

# The Docker Image Mirror script.
# Designed to mirror amd64 and arm64 images somewhere you want.
# Required environment variables:
#
#   * REGISTRY: to which registry images should be sent. Should not contain protocols,
#     should be the very same as you use 'docker login' with.
#   * REGISTRY_USERNAME: user name to use for destination registry.
#   * REGISTRY_PASSWORD: password to use for destination registry's user.
#   * REGISTRY_PROJECT: path within registry (or it's part) where images will be push'd.
#     Should start with '/'!

# What images we will mirror. Collected by executing scripts in "images" directory, see
# "collect_images" function.
WHAT_TO_MIRROR=()

# New tag for currently processing image as it will be used with 'docker push'.
# Updates for every image. Contains only path, image name and base tag from source.
REMOTE_IMAGE=""

# Is image multiarch? E.g. is we're successfully fetched not only amd64 image,
# but also an arm64? Updates for every image.
MULTIARCH=0

function cleanup() {
	local image=$1

	docker image rm "${image}" &> /dev/null
	docker image rm "${REMOTE_IMAGE}" &> /dev/null
	docker image rm "${REMOTE_IMAGE}"-amd64 &> /dev/null
	docker image rm "${REMOTE_IMAGE}"-arm64 &> /dev/null
}

function collect_images() {
	# Load shell files and execute them to get list of mirrorred images.
	MIRROR_CONFIGS=$(ls ./images/*.sh)
	# shellcheck disable=SC2068
	for file in ${MIRROR_CONFIGS[@]}; do
		echo "Importing ${file}..."
		# shellcheck disable=SC2086,SC2207,SC2206
		WHAT_TO_MIRROR=( ${WHAT_TO_MIRROR[@]} $(bash ${file}) )
	done

	echo "Images to mirror: ${WHAT_TO_MIRROR[*]}"
}

function get_image() {
	local image=$1
	local arch=$2

	image_name=$(echo "${image}" | cut -d":" -f 1)
	image_version=$(echo "${image}" | cut -d":" -f 2)


	# Check if layers for requested architecture should be fetched.
	if docker manifest inspect "${REMOTE_IMAGE}-${arch}" &> /dev/null; then
		echo -e "\t* Layers for ${arch} architecture for this image exist at registry we mirror to!"

		# This means that no mirrorring for arch-specific image is required.
		return 2
	fi

	echo -ne "\t* Getting ${arch} layers... "

	# arm64 layers might be missing. So we just put "FAIL!" here and do nothing else.
	if ! docker pull --platform=linux/"${arch}" "${image}" &> "/tmp/pull-${arch}"; then
		echo "FAIL!"

		echo -e "\nLogs:\n\n" && cat "/tmp/pull-${arch}" && echo -e "\n"

		# We presume that amd64 layers are always present. Returning an error here if they're absent.
		if [ "${arch}" == "amd64" ]; then
			return 1
		fi
	fi

	echo -n "Downloaded, "

	# shellcheck disable=SC1083
	image_hash=$(docker images -a | grep "^${image_name}" | grep "${image_version}" | awk {' print $3 '})

	if ! docker tag "${image_hash}" "${REMOTE_IMAGE}-${arch}" &> /dev/null; then
		echo "but tagging failed!"
		exit 1
	fi

	echo "tagged."

	# Set multiarch flag for any other arch than amd64.
	if [ "${arch}" != "amd64" ]; then
		MULTIARCH=1
	fi
}

function login_to_registry() {
	# Login to registry.
	echo "Logging into '${REGISTRY}' as '${REGISTRY_USERNAME}'..."
	docker login -u "${REGISTRY_USERNAME}" -p "${REGISTRY_PASSWORD}" "${REGISTRY}"
}

function mirror() {
	local image=$1
	image_name=$(echo "${image}" | cut -d":" -f 1)
	image_version=$(echo "${image}" | cut -d":" -f 2)

	echo "* Mirroring ${image}"

	REMOTE_IMAGE="${REGISTRY}${REGISTRY_PROJECT}/${image}"
	MULTIARCH=0

	# We presumes that amd64 image should always be available.
	if ! get_image "${image}" "amd64"; then
		exitcode=$?

		# Shit happened.
		if [ ${exitcode} -eq 1 ]; then
			echo "! Image mirroring failed! Cannot obtain amd64 image!"
		fi

		return ${exitcode}
	fi

	get_image "${image}" "arm64"
	push_multiarch_image "${image}"
	cleanup "${image}"
}

function push_multiarch_image() {
	local image=$1

	if ! docker push "${REMOTE_IMAGE}"-amd64 &> /tmp/push-amd64; then
		echo -e "\t! amd64 image push failed!"

		echo -e "\nLogs:\n\n" && cat /tmp/push-amd64 && echo -e "\n"
	else
		echo -e "\t* amd64 image pushed"
	fi

	if ! docker push "${REMOTE_IMAGE}"-arm64 &> /tmp/push-arm64; then
		echo -e "\t! arm64 image push failed!"

		echo -e "\nLogs:\n\n" && cat /tmp/push-arm64 && echo -e "\n"
	else
		echo -e "\t* arm64 image pushed"
	fi

	if [ ${MULTIARCH} -eq 1 ]; then
		echo -e "\t* Image is multi-arch, creating and pushing a manifest..."
		docker manifest create "${REMOTE_IMAGE}" \
			--amend "${REMOTE_IMAGE}"-amd64 \
			--amend "${REMOTE_IMAGE}"-arm64 \
			&> /dev/null

		docker manifest push "${REMOTE_IMAGE}" &> /dev/null
	fi
}

function start_docker_daemon() {
	# Starting Docker daemon.
	/usr/local/bin/dockerd --data-root=/var/lib/docker --max-concurrent-uploads 1 &

	# Wait for it.
	echo "Waiting for Docker daemon to start..."
	while true; do
		if docker ps &> /dev/null; then
			break
		fi
	done
}

start_docker_daemon
login_to_registry
collect_images

for package in "${WHAT_TO_MIRROR[@]}"; do
	if ! mirror "${package}"; then
		exitcode=$?

		echo "! Failed to mirror package ${package}! Exit code: ${exitcode} (1 is fatal)."

		# Shit happens only for exitcode 1.
		if [ ${exitcode} -eq 1 ]; then
			exit 1
		fi
	fi
done