#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

usage() {
    cat <<EOF
Usage: $(basename "$0") [--help]

Pull Azaion AI Training Docker images from the container registry.

Environment:
  DOCKER_REGISTRY     Registry URL (required)
  DOCKER_IMAGE_TAG    Image tag to pull (default: latest)
EOF
    exit 0
}

[[ "${1:-}" == "--help" ]] && usage

if [[ -f "$PROJECT_ROOT/.env" ]]; then
    set -a
    source "$PROJECT_ROOT/.env"
    set +a
fi

DOCKER_REGISTRY="${DOCKER_REGISTRY:?DOCKER_REGISTRY is required}"
DOCKER_IMAGE_TAG="${DOCKER_IMAGE_TAG:-latest}"

IMAGES=(
    "${DOCKER_REGISTRY}/azaion/training:${DOCKER_IMAGE_TAG}"
    "${DOCKER_REGISTRY}/azaion/annotation-queue:${DOCKER_IMAGE_TAG}"
)

echo "Pulling images (tag: ${DOCKER_IMAGE_TAG})..."

for image in "${IMAGES[@]}"; do
    echo "  Pulling $image ..."
    if docker pull "$image"; then
        echo "  OK: $image"
    else
        echo "  FAILED: $image"
        exit 1
    fi
done

echo "All images pulled successfully."
