#!/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-v3/*.whl ./tmp/
  cp ./jupyter-extension-install-v3/requirements.txt ./tmp/
  cp ./extension.sh ./tmp/
  cp ./jupyter_lab_config.py ./tmp/
  cp -r ./jupyter-extension-install-v3/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 --build-arg BASE_IMAGE_IS_TORCH=1 --build-arg IMAGE="pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime" -f ./tmp/${tmp_dockerfile} -t ${newtag} ./tmp/
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/
#echo "docker build -f ./tmp/${tmp_dockerfile} -t ${image_tag} $build_args ./tmp/"
build_status=$?
rm -r ./tmp
if [[ build_status -eq 0 ]];then
        echo -e "\033[32mBuild Image Successfully !\033[0m"
else
        echo -e "\033[32mBuild Image fail!\033[0m"
        exit 1
fi