#!/bin/bash
set -euo pipefail

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

OS=linux

# IMPORTANT:
# This must match what the target system will use at runtime.
# If your target CRI-O uses overlay, leave this as overlay.
# If overlay seeding fails in your build environment, fix the build env;
# do not silently switch to vfs unless the target will also use vfs.
STORAGE_DRIVER="${STORAGE_DRIVER:-overlay}"

# Optional. Needed by many overlay-in-container setups.
FUSE_OVERLAYFS="${FUSE_OVERLAYFS:-/usr/bin/fuse-overlayfs}"

# Keep archive cache version/arch scoped so downloads do not get mixed.
ARCHIVE_DIR="${KUBE_IMG_CACHE}/archives/${ARCH}/${KUBE_VERSION}"

# Seed directly into the target rootfs image store.
TARGET_GRAPHROOT="${ROOTFS}/var/lib/containers/storage"

# Build-host temporary runtime state for containers/storage.
TMPDIR="$(mktemp -d)"
RUNROOT="${TMPDIR}/runroot"
STORAGE_CONF="${TMPDIR}/storage.conf"

cleanup() {
    rm -rf "${TMPDIR}"
}
trap cleanup EXIT

mkdir -p "${ARCHIVE_DIR}"
mkdir -p "${TARGET_GRAPHROOT}"
mkdir -p "${RUNROOT}"

echo "============================================================"
echo "Preparing Kubernetes CRI-O image store"
echo "  KUBE_VERSION      = ${KUBE_VERSION}"
echo "  ARCH              = ${ARCH}"
echo "  ARCHIVE_DIR       = ${ARCHIVE_DIR}"
echo "  TARGET_GRAPHROOT  = ${TARGET_GRAPHROOT}"
echo "  STORAGE_DRIVER    = ${STORAGE_DRIVER}"
echo "============================================================"

# Build a temporary containers/storage config that writes directly into
# the target rootfs image store.
if [ "${STORAGE_DRIVER}" = "overlay" ]; then
    cat > "${STORAGE_CONF}" <<EOF
[storage]
driver = "overlay"
runroot = "${RUNROOT}"
graphroot = "${TARGET_GRAPHROOT}"

[storage.options]
mount_program = "${FUSE_OVERLAYFS}"
EOF
else
    cat > "${STORAGE_CONF}" <<EOF
[storage]
driver = "${STORAGE_DRIVER}"
runroot = "${RUNROOT}"
graphroot = "${TARGET_GRAPHROOT}"
EOF
fi

export CONTAINERS_STORAGE_CONF="${STORAGE_CONF}"

# Sanity check: list required images using the target kubeadm binary.
mapfile -t kube_images < <(
    qemu-aarch64-static \
        "${ROOTFS}/usr/local/bin/kubeadm" \
        config images list \
        --kubernetes-version "${KUBE_VERSION}"
)

if [ "${#kube_images[@]}" -eq 0 ]; then
    echo "No kubeadm images returned"
    exit 1
fi

for img in "${kube_images[@]}"; do
    [ -n "${img}" ] || continue

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

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

    # If already present in the target containers/storage, skip.
    if skopeo inspect "containers-storage:${img}" >/dev/null 2>&1; then
        echo "Store hit: ${img}"
        continue
    fi

    echo "Seeding CRI-O store: ${img}"
    skopeo copy \
        --override-os "${OS}" \
        --override-arch "${ARCH}" \
        "oci-archive:${archive}" \
        "containers-storage:${img}"
done

echo
echo "Done."
echo "Archives:       ${ARCHIVE_DIR}"
echo "CRI-O storage:  ${TARGET_GRAPHROOT}"