#!/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_amd64_image() {
	local image=$1

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

	# Check if amd64 layers should be fetched.
	if docker manifest inspect "${REMOTE_IMAGE}-amd64" &> /dev/null; then
		echo -e "\t* Layers for amd64 architecture for this image exist."

		return 2
	fi

	echo -ne "\t* Getting amd64 layers... "

	if ! docker pull --platform=linux/amd64 "${image}" &> /dev/null; then
		echo "FAIL!"

		return 1
	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}-amd64" &> /dev/null; then
		echo "but tagging failed!"
		exit 1
	fi

	echo "tagged."
}

function get_arm64_image() {
	local image=$1

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

	# Check if arm64 layers should be fetched.
	if docker manifest inspect "${REMOTE_IMAGE}-arm64" &> /dev/null; then
		echo -e "\t* Layers for arm64 architecture for this image exist."

		return 2
	fi

	echo -ne "\t* Getting arm64 layers... "

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

		return 1
	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}-arm64" &> /dev/null; then
		echo "but tagging failed!"
		exit 1
	fi

	echo "tagged."

	MULTIARCH=1
}

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_amd64_image "${image}"; then
		echo "! Image mirroring failed! Cannot obtain amd64 image!"

		return 1
	fi

	get_arm64_image "${image}"
	push_multiarch_image "${image}"
	cleanup "${image}"
}

function push_multiarch_image() {
	local image=$1

	if ! docker push -f "${REMOTE_IMAGE}"-amd64 &> /dev/null; then
		echo -e "\t! amd64 image push failed!"
	else
		echo -e "\t* amd64 image pushed"
	fi

	if ! docker push -f ="${REMOTE_IMAGE}"-arm64 &> /dev/null; then
		echo -e "\t! arm64 image push failed!"
	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
		echo "! Failed to mirror package ${package}!"

		exit 1
	fi
done