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

DRIVER_REPO="$1"
SHARDS_COUNT="$2"

# If nothing was built, we don't need to unify anything
if [[ ! -s ~/workspace/build-tasks ]]; then
	if [[ "${COLLECTOR_DRIVERS_TAG}" == "${COLLECTOR_DRIVERS_CACHE}" ]]; then
		exit 0
	fi

	# If the cache and the desired tag are different, we can re-tag and be done
	docker tag \
		"${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_CACHE}" \
		"${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_TAG}"

	docker push "${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_TAG}"

	exit 0
fi

if [[ "${COLLECTOR_DRIVERS_TAG}" != "${COLLECTOR_DRIVERS_CACHE}" ]]; then
	# Create an empty base image if we are not building the cache
	docker pull registry.access.redhat.com/ubi8/ubi:8.5
	docker tag registry.access.redhat.com/ubi8/ubi:8.5 \
		"${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_TAG}"
fi

for ((i=0; i<SHARDS_COUNT; i++)); do
	PARTIAL_DRIVERS_TAG="${COLLECTOR_DRIVERS_TAG}-${i}"

	if ! docker pull "${DRIVER_REPO}/collector-drivers:${PARTIAL_DRIVERS_TAG}"; then
		if [[ -s "${WORKSPACE_ROOT}/task-shard-${i}" ]]; then
			echo >&2 "Failed to pull ${DRIVER_REPO}/collector-drivers:${PARTIAL_DRIVERS_TAG}"
			exit 1
		fi
		# Failed to pull a partial image because it didn't need to build anything.
		# We keep going.
		continue
	fi

	DOCKER_BUILDKIT=1 \
	BUILDKIT_PROGRESS=plain \
	docker build \
		--tag "${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_CACHE}" \
		--build-arg BASE_TAG="${COLLECTOR_DRIVERS_CACHE}" \
		--build-arg DRIVER_TAG="${PARTIAL_DRIVERS_TAG}" \
		-f "${SOURCE_ROOT}/kernel-modules/dockerized/Dockerfile.merge" \
		"${SOURCE_ROOT}/kernel-modules/dockerized"
done

# Once we merged everything together, we squash the image and push it.
DOCKER_BUILDKIT=1 \
BUILDKIT_PROGRESS=plain \
docker build \
	--tag "${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_TAG}" \
	--build-arg DRIVER_REPO="${DRIVER_REPO}" \
	--build-arg DRIVER_TAG="${COLLECTOR_DRIVERS_CACHE}" \
	-f "${SOURCE_ROOT}/kernel-modules/dockerized/Dockerfile.squash" \
	"${SOURCE_ROOT}/kernel-modules/dockerized"

docker push "${DRIVER_REPO}/collector-drivers:${COLLECTOR_DRIVERS_TAG}"
