#/bin/bash

# pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=12.1
# cuda=11.6（这个版本的空镜像，有些框架需要低版本）
# paddlepaddle-gpu==2.6.1.post120  cuda12
# tensroflow 安装最高版本 cuda12

rm -r ./tmp

framework=$1
newtag=$2
tmp_dockerfile="Dockerfile.${RANDOM}"

if [ ! -d "./tmp" ];then
mkdir tmp
fi
if [ "${framework}" == "jupyterlab" ];then
  cp ./jupyter-extension-install-v2/*.whl ./tmp/
  cp ./jupyter-extension-install-v2/requirements.txt ./tmp/
  cp ./extension.sh ./tmp/
  cp ./jupyter_lab_config.py ./tmp/
  cp -r ./jupyter-extension-install-v2/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}
# temp_image="$(grep -n '^FROM' ./tmp/${tmp_dockerfile} | tac | head -1 | awk '{print $2}')"
# sed -i "s?${temp_image}?${base_tag}?g" ./tmp/${tmp_dockerfile}
docker build -f ./tmp/${tmp_dockerfile} -t ${newtag} ./tmp/
if [[ $? -eq 0 ]];then
        echo -e "\033[32mBuild Image Successfully !\033[0m"
else
        echo -e "\033[32mBuild Image unsuccessful !\033[0m"
        exit 1
fi