#!/bin/bash

# Get image name as argument
IMAGE_NAME=$1
PLATFORM="linux/arm64"

# Define paths
IMAGE_DIR="docker/images"
IMAGE_FILE="${IMAGE_DIR}/${IMAGE_NAME//[:\/]/_}.tar"
COMPRESSED_FILE="${IMAGE_FILE}.xz"

# Function to pull the image
pull_image() {
    local image=$1
    if [[ $image == arm64v8/* ]]; then
        echo "Pulling image $image without platform specification..."
        docker pull $image
    else
        echo "Pulling image $image for platform $PLATFORM..."
        docker pull --platform $PLATFORM $image
    fi
    return $?
}

# Pull the image if it is not available locally
if ! docker image inspect ${IMAGE_NAME} &>/dev/null; then
    if pull_image ${IMAGE_NAME}; then
        echo "Image $IMAGE_NAME pulled successfully."
    else
        echo "Error occurred while pulling $IMAGE_NAME for platform $PLATFORM"
        echo "Trying to pull $IMAGE_NAME without platform specification..."

        # Attempt to pull again without platform
        if pull_image ${IMAGE_NAME}; then
            echo "Image $IMAGE_NAME pulled successfully without platform."
        else
            echo "Error occurred while pulling $IMAGE_NAME without platform."
            echo "Trying to pull arm64v8/${IMAGE_NAME} instead..."

            # Construct new image name
            NEW_IMAGE_NAME="arm64v8/${IMAGE_NAME}"
            if pull_image ${NEW_IMAGE_NAME}; then
                echo "Image $NEW_IMAGE_NAME pulled successfully."
                IMAGE_NAME=${NEW_IMAGE_NAME}  # Update IMAGE_NAME to use the new one
            else
                echo "Error occurred while pulling $NEW_IMAGE_NAME"
                exit 1
            fi
        fi
    fi
else
    echo "Image $IMAGE_NAME found locally. Skipping pull."
fi

# Save the Docker image
echo "Saving $IMAGE_NAME Docker image..."
docker save ${IMAGE_NAME} > $IMAGE_FILE

# Compress the Docker image (overwriting if file exists)
echo "Compressing $IMAGE_FILE..."
xz -z --force $IMAGE_FILE

if [ $? -eq 0 ]; then
    echo "$IMAGE_NAME Docker image saved and compressed successfully as $COMPRESSED_FILE"
else
    echo "Error occurred while compressing $IMAGE_NAME Docker image"
    exit 1
fi