#!/bin/bash

# 框架
framework=$1
# 输出镜像tag
image_tag=$2
# 基础镜像
base_image=$3
build_args=" --build-arg BASE_IMAGE=$base_image"
if [ ${base_image%%:*} = "pytorch/pytorch" ]; then
    build_args="$build_args --build-arg BASE_IMAGE_IS_TORCH=1 "
fi
for arg in ${*:4}
do
    build_args="$build_args --build-arg $arg "
done

tmp_dockerfile="Dockerfile.${RANDOM}"

if [ ! -d "./tmp" ];then
mkdir tmp
fi
if [ "${framework}" == "jupyterlab" ];then
  cp ./jupyter-extension-install-v4/*.whl ./tmp/
  cp ./jupyter-extension-install-v4/requirements.txt ./tmp/
  cp ./extension.sh ./tmp/
  cp ./jupyter_lab_config.py ./tmp/
  cp -r ./jupyter-extension-install-v4/static ./tmp/
fi
if [ "${framework}" == "codeserver" ];then
  cp ./code-server_4.7.0_amd64.deb ./tmp/
fi
cp ./python-requirements.txt ./tmp/
cp -f ./Dockerfile.${framework}_ubuntu ./tmp/${tmp_dockerfile}
echo "docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/"
docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/
rm -r ./tmp
if [[ $? -eq 0 ]];then
        echo -e "\033[32mBuild Image Successfully !\033[0m"
else
        echo -e "\033[32mBuild Image fail!\033[0m"
        exit 1
fi