#!/bin/bash
set -e

usage() {
cat << EOF
Usage: $0 [OPTIONS] [LABEL]
Push a newly-built image with the given LABEL to gcr.io and DockerHub.

Options:
    -g, --gpu                   Push the image with GPU support.
    -t, --tpu                   Push the image with GPU support.
    -s, --source-image IMAGE    Tag for the source image. 
EOF
}

SOURCE_IMAGE_TAG='kaggle/python-build:latest'
SOURCE_IMAGE_TAG_OVERRIDE=''
TARGET_IMAGE='gcr.io/kaggle-images/python'

while :; do
    case "$1" in 
        -h|--help)
            usage
            exit
            ;;
        -g|--gpu)
            SOURCE_IMAGE_TAG='kaggle/python-gpu-build:latest'
            TARGET_IMAGE='gcr.io/kaggle-private-byod/python'
            ;;
        -t|--tpu)
            SOURCE_IMAGE_TAG='kaggle/python-tpuvm-build:latest'
            TARGET_IMAGE='gcr.io/kaggle-private-byod/python-tpuvm'
            ;;
        -s|--source-image)
            if [[ -z $2 ]]; then
                usage
                printf 'ERROR: No IMAGE specified after the %s flag.\n' "$1" >&2
                exit
            fi
            SOURCE_IMAGE_TAG_OVERRIDE=$2
            shift # skip the flag value
            ;;
        -?*)
            usage
            printf 'ERROR: Unknown option: %s\n' "$1" >&2
            exit
            ;;
        *)            
            break
    esac

    shift
done

LABEL=${1:-testing}

if [[ -n "$SOURCE_IMAGE_TAG_OVERRIDE" ]]; then
    SOURCE_IMAGE_TAG="$SOURCE_IMAGE_TAG_OVERRIDE"
fi

readonly SOURCE_IMAGE_TAG
readonly TARGET_IMAGE
readonly LABEL

set -x
docker tag "${SOURCE_IMAGE_TAG}" "${TARGET_IMAGE}:${LABEL}"
gcloud docker -- push "${TARGET_IMAGE}:${LABEL}"