#!/bin/bash
set -euo pipefail

: "${KUBE_IMG_CACHE:?KUBE_IMG_CACHE is required}"
: "${KUBE_VERSION:?KUBE_VERSION is required}"
: "${ARCH:?ARCH is required}"

OS=linux
REGISTRY_PORT=5000

# Keep everything version/arch scoped so caches do not get mixed.
ARCHIVE_DIR="${KUBE_IMG_CACHE}/archives/${ARCH}/${KUBE_VERSION}"
REGISTRY_DATA_DIR="${KUBE_IMG_CACHE}/registry-data/${ARCH}/${KUBE_VERSION}"
REGISTRY_TAR="${KUBE_IMG_CACHE}/registry.tar"
REGISTRY_IMAGE="docker.io/library/registry:2"

if podman image exists "${REGISTRY_IMAGE}"; then
    echo "Registry image already present: ${REGISTRY_IMAGE}"

elif [ -f "${REGISTRY_TAR}" ]; then
    echo "Loading registry image from cache: ${REGISTRY_TAR}"
    podman load -i "${REGISTRY_TAR}"

else
    echo "Cache miss → pulling and saving registry image"

    podman pull "${REGISTRY_IMAGE}"
    podman save -o "${REGISTRY_TAR}" "${REGISTRY_IMAGE}"
fi

mkdir -p "${ARCHIVE_DIR}"
rm -rf "${REGISTRY_DATA_DIR}"
mkdir -p "${REGISTRY_DATA_DIR}"

echo "============================================================"
echo "Preparing Kubernetes image cache"
echo "  KUBE_VERSION      = ${KUBE_VERSION}"
echo "  ARCH              = ${ARCH}"
echo "  ARCHIVE_DIR       = ${ARCHIVE_DIR}"
echo "  REGISTRY_DATA_DIR = ${REGISTRY_DATA_DIR}"
echo "============================================================"

# Start a temporary local registry backed by REGISTRY_DATA_DIR.
# We use host network here to avoid messing with nested networking.
podman rm -f temp-registry >/dev/null 2>&1 || true

podman run -d \
    --name temp-registry \
    --network host \
    --cgroups=disabled \
    -v "${REGISTRY_DATA_DIR}:/var/lib/registry" \
    docker.io/library/registry:2

cleanup() {
    podman rm -f temp-registry >/dev/null 2>&1 || true
}
trap cleanup EXIT

# Wait for registry to answer.
for _ in $(seq 1 30); do
    if curl -fsS "http://127.0.0.1:${REGISTRY_PORT}/v2/" >/dev/null; then
        break
    fi
    sleep 1
done

if ! curl -fsS "http://127.0.0.1:${REGISTRY_PORT}/v2/" >/dev/null; then
    echo "Temporary registry did not start"
    exit 1
fi

# Cache and seed all kubeadm-required images.
while read -r img; do
    [ -n "$img" ] || continue

    safe_name=$(printf '%s' "$img" | sed 's#/#_#g; s#:#__#g')
    archive="${ARCHIVE_DIR}/${safe_name}.tar"
    ref="${img#registry.k8s.io/}"

    if [ ! -f "$archive" ]; then
        echo "Caching: $img"
        skopeo copy \
            --override-os "${OS}" \
            --override-arch "${ARCH}" \
            "docker://$img" \
            "oci-archive:$archive"
    else
        echo "Cache hit: $img"
    fi

    if skopeo inspect --tls-verify=false "docker://127.0.0.1:${REGISTRY_PORT}/${ref}" >/dev/null 2>&1; then
        echo "Registry hit: 127.0.0.1:${REGISTRY_PORT}/${ref}"
        continue
    fi

    echo "Seeding registry: $ref"
    skopeo copy \
        --dest-tls-verify=false \
        "oci-archive:$archive" \
        "docker://127.0.0.1:${REGISTRY_PORT}/${ref}"
done < <(qemu-aarch64-static "$ROOTFS/usr/local/bin/kubeadm" config images list --kubernetes-version "${KUBE_VERSION}")

echo
echo "Done."
echo "Archives:      ${ARCHIVE_DIR}"
echo "Registry data: ${REGISTRY_DATA_DIR}"