#!/bin/bash
set -e

# Usage function
usage() {
    echo "Usage: $0 <source-image> [target-image]"
    echo ""
    echo "Examples:"
    echo "  $0 ghcr.io/evanjarrett/myapp:latest"
    echo "  $0 ghcr.io/evanjarrett/myapp:latest atcr.io/evan.jarrett.net/myapp:latest"
    echo ""
    echo "If target-image is not specified, it will use atcr.io/<username>/<repo>:<tag>"
    exit 1
}

# Check arguments
if [ $# -lt 1 ]; then
    usage
fi

SOURCE_IMAGE="$1"
TARGET_IMAGE="${2:-}"

# Parse source image to extract components
# Format: [registry/]repository[:tag|@digest]
parse_image_ref() {
    local ref="$1"
    local registry=""
    local repository=""
    local tag="latest"

    # Remove digest if present (we'll fetch the manifest-list)
    ref="${ref%@*}"

    # Extract tag
    if [[ "$ref" == *:* ]]; then
        tag="${ref##*:}"
        ref="${ref%:*}"
    fi

    # Extract registry and repository
    if [[ "$ref" == */*/* ]]; then
        # Has registry
        registry="${ref%%/*}"
        repository="${ref#*/}"
    else
        # No registry, assume Docker Hub
        registry="docker.io"
        repository="$ref"
    fi

    echo "$registry" "$repository" "$tag"
}

# Parse source image
read -r SOURCE_REGISTRY SOURCE_REPO SOURCE_TAG <<< "$(parse_image_ref "$SOURCE_IMAGE")"

# If no target specified, auto-generate it
if [ -z "$TARGET_IMAGE" ]; then
    # Extract just the repo name (last component)
    REPO_NAME="${SOURCE_REPO##*/}"
    # Try to extract username from source
    if [[ "$SOURCE_REPO" == */* ]]; then
        USERNAME="${SOURCE_REPO%/*}"
        USERNAME="${USERNAME##*/}"
    else
        USERNAME="default"
    fi
    TARGET_IMAGE="atcr.io/${USERNAME}/${REPO_NAME}:${SOURCE_TAG}"
fi

# Parse target image
read -r TARGET_REGISTRY TARGET_REPO TARGET_TAG <<< "$(parse_image_ref "$TARGET_IMAGE")"

echo "=== Migrating multi-arch image ==="
echo "Source: ${SOURCE_REGISTRY}/${SOURCE_REPO}:${SOURCE_TAG}"
echo "Target: ${TARGET_REGISTRY}/${TARGET_REPO}:${TARGET_TAG}"
echo ""

# Full source reference
SOURCE_REF="${SOURCE_REGISTRY}/${SOURCE_REPO}:${SOURCE_TAG}"
TARGET_REF="${TARGET_REGISTRY}/${TARGET_REPO}:${TARGET_TAG}"

# Fetch the manifest list
echo ">>> Fetching manifest list from source..."
MANIFEST_JSON=$(docker manifest inspect "$SOURCE_REF" 2>/dev/null || {
    echo "Error: Failed to fetch manifest list. This may not be a multi-arch image."
    echo "Trying as single-arch image..."

    # Try pulling as single image
    docker pull "$SOURCE_REF"
    docker tag "$SOURCE_REF" "$TARGET_REF"
    docker push "$TARGET_REF"
    echo "=== Migration complete (single-arch) ==="
    exit 0
})

# Check if this is a manifest list
MEDIA_TYPE=$(echo "$MANIFEST_JSON" | jq -r '.mediaType // .schemaVersion')
if [[ ! "$MEDIA_TYPE" =~ "manifest.list" ]] && [[ ! "$MEDIA_TYPE" =~ "index" ]]; then
    echo "Warning: Source appears to be a single-arch image, not a manifest list."
    docker pull "$SOURCE_REF"
    docker tag "$SOURCE_REF" "$TARGET_REF"
    docker push "$TARGET_REF"
    echo "=== Migration complete (single-arch) ==="
    exit 0
fi

echo "Found multi-arch manifest list"
echo ""

# Extract platform information and digests (skip unknown/unknown for cosign artifacts)
PLATFORMS=$(echo "$MANIFEST_JSON" | jq -r '.manifests[] | select(.platform.os != "unknown" and .platform.architecture != "unknown") | "\(.platform.os)|\(.platform.architecture)|\(.platform.variant // "")|\(.digest)"')

# Arrays to store pushed images for manifest creation
declare -a PUSHED_IMAGES
declare -a PLATFORM_INFO

# Process each platform
while IFS='|' read -r os arch variant digest; do
    # Create platform tag (e.g., "linux-amd64" or "linux-arm-v7")
    PLATFORM_TAG="${os}-${arch}"
    if [ -n "$variant" ]; then
        PLATFORM_TAG="${PLATFORM_TAG}-${variant}"
    fi

    echo ">>> Processing ${os}/${arch}${variant:+/$variant}..."
    echo "    Digest: $digest"

    # Pull by digest
    echo "    Pulling image..."
    docker pull "${SOURCE_REGISTRY}/${SOURCE_REPO}@${digest}"

    # Tag for target
    TARGET_PLATFORM_REF="${TARGET_REGISTRY}/${TARGET_REPO}:${TARGET_TAG}-${PLATFORM_TAG}"
    echo "    Tagging as: ${TARGET_PLATFORM_REF}"
    docker tag "${SOURCE_REGISTRY}/${SOURCE_REPO}@${digest}" "${TARGET_PLATFORM_REF}"

    # Push platform-specific image
    echo "    Pushing..."
    docker push "${TARGET_PLATFORM_REF}"

    # Store for manifest creation
    PUSHED_IMAGES+=("${TARGET_PLATFORM_REF}")
    PLATFORM_INFO+=("${os}|${arch}|${variant}")

    echo ""
done <<< "$PLATFORMS"

# Create multi-arch manifest
echo ">>> Creating multi-arch manifest..."
MANIFEST_CREATE_CMD="docker manifest create ${TARGET_REF}"
for img in "${PUSHED_IMAGES[@]}"; do
    MANIFEST_CREATE_CMD="${MANIFEST_CREATE_CMD} --amend ${img}"
done

eval "$MANIFEST_CREATE_CMD"
echo ""

# Annotate each platform
echo ">>> Annotating manifest with platform information..."
for i in "${!PUSHED_IMAGES[@]}"; do
    IFS='|' read -r os arch variant <<< "${PLATFORM_INFO[$i]}"

    ANNOTATE_CMD="docker manifest annotate ${TARGET_REF} ${PUSHED_IMAGES[$i]} --os ${os} --arch ${arch}"
    if [ -n "$variant" ]; then
        ANNOTATE_CMD="${ANNOTATE_CMD} --variant ${variant}"
    fi

    echo "  Annotating ${os}/${arch}${variant:+/$variant}..."
    eval "$ANNOTATE_CMD"
done
echo ""

# Push the manifest list
echo ">>> Pushing multi-arch manifest..."
docker manifest push "${TARGET_REF}"
echo ""

echo "=== Migration complete! ==="
echo "You can now pull: docker pull ${TARGET_REF}"
echo ""
echo "Migrated platforms:"
for i in "${!PLATFORM_INFO[@]}"; do
    IFS='|' read -r os arch variant <<< "${PLATFORM_INFO[$i]}"
    echo "  - ${os}/${arch}${variant:+/$variant}"
done