#!/bin/bash

# script that builds docker image with credentials baked in
# usage:
# studio-add-credentials [--base-image=<image>] [--tag=<tag>] [--check-gpu]
#
# --base-image specifies the base image to add credentials to. Default is peterzhokhoff/tfstudio 
# --tag specifes the tag of the resulting image. Default is <image>_creds
# --check-gpu option (if specified) checks if nvidia-smi works correctly on a current machine, and if it is not, uninstalls tensorflow-gpu from docker image. 
# Without --check-gpu option the built docker image may not work on a current machine. 

base_img=peterzhokhoff/tfstudio

while [[ $# -gt 0 ]] 
do 
    key="$1"
    case ${key%%=*} in 

            -b|--base-image)
            base_img="${1##*=}"
            ;;

            -t|--tag)
            output_img="${1##*=}"
            ;;

    esac
    shift
done

if [ -z $output_img ]; then
        output_img=$base_img"_creds"
fi

# mypath=$(pwd)/${0%/*}
dockerfile=".Dockerfile_bake_creds"
awspath=$HOME/.aws

# cd $mypath/../../

echo "Base image: $base_img"
echo "Tag: $output_img"


echo "Uninstall tensorflow-gpu: $uninstall_tfgpu"

contextdir=$TMPDIR/studioml_container_context
mkdir $contextdir
cd $contextdir

if [ -d $HOME/.studioml/keys ]; then
    cp -r $HOME/.studioml/keys .keys
fi

if [ -n $GOOGLE_APPLICATION_CREDENTIALS ]; then
        cp $GOOGLE_APPLICATION_CREDENTIALS .gac_credentials
fi

if [ -d $awspath ]; then
        cp -r $awspath .aws
fi

# build dockerfile
echo "Constructing dockerfile..."
echo "FROM $base_img" > $dockerfile

if [ -d $HOME/.studioml/keys/ ]; then
    echo "ADD .keys /root/.studioml/keys" >> $dockerfile
fi

if [ -n $GOOGLE_APPLICATION_CREDENTIALS ]; then
    echo "ADD .gac_credentials /root/gac_credentials" >> $dockerfile
    echo "ENV GOOGLE_APPLICATION_CREDENTIALS /root/gac_credentials" >> $dockerfile
fi

if [ -d $awspath ]; then
   echo "ADD .aws /root/.aws" >> $dockerfile
fi

echo "Done. Resulting dockerfile: "
cat $dockerfile


# build docker image
echo "Building image..."
docker build -t $output_img -f $dockerfile .
echo "Done"

# cleanup 
echo "Cleaning up..."
rm -rf .keys
rm -rf $dockerfile

if [ -n $GOOGLE_APPLICATION_CREDENTIALS ]; then
    rm -rf $gac_file
fi

if [ -d $awspath ]; then
   rm -rf .aws
fi
echo "Done"
