#!/bin/bash
set -e

usage() {
cat << EOF
Usage: $0 [OPTIONS]
Compare a given Docker image package versions against the prod image.

Options:
    -g, --gpu       Compare GPU images.
    -b, --base      The base image to diff against.
    -t, --target    The image to diff against the base image.
                    Default is the locally built image.
    -p, --package   Only show diff for this package and its dependencies.
EOF
}


BASE_IMAGE_TAG='gcr.io/kaggle-images/python:latest'
BASE_IMAGE_TAG_OVERRIDE=''
TARGET_IMAGE_TAG='kaggle/python-build'
TARGET_IMAGE_TAG_OVERRIDE=''
PACKAGE_NAME=''

while :; do
    case "$1" in 
        -h|--help)
            usage
            exit
            ;;
        -g|--gpu)
            BASE_IMAGE_TAG='gcr.io/kaggle-private-byod/python:latest'
            TARGET_IMAGE_TAG='kaggle/python-gpu-build'
            ;;
        -b|--base)
            if [[ -z "$2" ]]; then
                usage
                printf 'ERROR: No IMAGE specified after the %s flag.\n' "$1" >&2
                exit
            fi
            BASE_IMAGE_TAG_OVERRIDE="$2"
            shift # skip the flag value
            ;;
        -t|--target)
            if [[ -z "$2" ]]; then
                usage
                printf 'ERROR: No IMAGE specified after the %s flag.\n' "$1" >&2
                exit
            fi
            TARGET_IMAGE_TAG_OVERRIDE="$2"
            shift # skip the flag value
            ;;
        -p|--package)
            if [[ -z "$2" ]]; then
                usage
                printf 'ERROR: No PACKAGE specified after the %s flag.\n' "$1" >&2
                exit
            fi
            PACKAGE_NAME="$2"
            shift # skip the flag value
            ;;
        -?*)
            usage
            printf 'ERROR: Unknown option: %s\n' "$1" >&2
            exit
            ;;
        *)            
            break
    esac

    shift
done

if [[ -n "$BASE_IMAGE_TAG_OVERRIDE" ]]; then
    BASE_IMAGE_TAG="$BASE_IMAGE_TAG_OVERRIDE"
fi

if [[ -n "$TARGET_IMAGE_TAG_OVERRIDE" ]]; then
    TARGET_IMAGE_TAG="$TARGET_IMAGE_TAG_OVERRIDE"
fi

readonly BASE_IMAGE_TAG
readonly TARGET_IMAGE_TAG

echo "Base: $BASE_IMAGE_TAG"
echo "Target: $TARGET_IMAGE_TAG"

if [[ "$BASE_IMAGE_TAG" == "gcr.io/"* ]]; then
    docker pull "$BASE_IMAGE_TAG"
fi


if [[ -n "$PACKAGE_NAME" ]]; then
    echo "Package: $PACKAGE_NAME"
    CMDS=("python /tools/pip_list_versions.py $PACKAGE_NAME | sort")
else
    CMDS=("pip freeze" 'cat /etc/os-release | grep -oP "PRETTY_NAME=\"\K([^\"]*)"' "uname -r" "dpkg --list | awk '{print \$2\"==\"\$3}'")
fi

for cmd in "${CMDS[@]}"; do
    echo "== Comparing $cmd =="
    diff --suppress-common-lines --side-by-side \
        <(docker run -v $PWD/tools:/tools --rm "$BASE_IMAGE_TAG" /bin/bash -c "$cmd") \
        <(docker run -v $PWD/tools:/tools --rm "$TARGET_IMAGE_TAG" /bin/bash -c "$cmd") \
        && echo 'No diff' || true
done