#!/bin/bash
ARCH=$(uname -m)

if [ "$ARCH" = "x86_64" ]; then
	ARCH=x64
fi

if [ ! -f "bin/micromamba" ]; then
	if [ "$ARCH" = "x64" ]; then
		 curl -Ls https://anaconda.org/conda-forge/micromamba/1.5.3/download/linux-64/micromamba-1.5.3-0.tar.bz2 | tar -xvj bin/micromamba
	elif [ "$ARCH" = "aarch64" ]; then
		 curl -Ls https://anaconda.org/conda-forge/micromamba/1.5.3/download/linux-aarch64/micromamba-1.5.3-0.tar.bz2 | tar -xvj bin/micromamba
	else
		 echo "CPU Architecture $ARCH is not supported by this script, please try compiling manually."
		 exit 1
	fi
fi

NVIDIA_GPU=0
if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L 2>/dev/null | grep -qE '^GPU [0-9]+:'; then
	NVIDIA_GPU=1
fi

if [[ ! -f "conda/envs/linux/bin/python" || $1 == "rebuild" ]] && [ -z "$KCPP_CUDA" ]; then
	if [ "$NVIDIA_GPU" = 1 ]; then
		if nvidia-smi | grep -qE 'CUDA Version: (11|12\.0)'; then
			KCPP_CUDA=11.4.0
			ARCHES_CU11=true
		else
			KCPP_CUDA=12.1.0
		fi
	elif command -v rocmsmi >/dev/null 2>&1 || [ -d /opt/rocm ]; then
		KCPP_CUDA=rocm
	else
		KCPP_CUDA=12.1.0
	fi
fi

if [[ ! -f "conda/envs/linux/bin/python" && $KCPP_CUDA != "rocm" || $1 == "rebuild" && $KCPP_CUDA != "rocm" ]]; then
	cp environment.yaml environment.tmp.yaml
	sed -i -e "s/nvidia\/label\/cuda-12.1.0/nvidia\/label\/cuda-$KCPP_CUDA/g" environment.tmp.yaml
	bin/micromamba create --no-rc --no-shortcuts -r conda -p conda/envs/linux -f environment.tmp.yaml -y
	bin/micromamba run -r conda -p conda/envs/linux make clean
	echo $KCPP_CUDA > conda/envs/linux/cudaver
	rm environment.tmp.yaml
fi

if [[ ! -f "conda/envs/linux/bin/python" && $KCPP_CUDA == "rocm" || $1 == "rebuild" && $KCPP_CUDA == "rocm" ]]; then
	bin/micromamba create --no-rc --no-shortcuts -r conda -p conda/envs/linux -f environment-nocuda.yaml -y
	bin/micromamba run -r conda -p conda/envs/linux make clean
	echo "rocm" > conda/envs/linux/cudaver
fi

KCPP_CUDA=$(<conda/envs/linux/cudaver)
KCPP_CUDAAPPEND=-cuda${KCPP_CUDA//.}$KCPP_APPEND

if [[ "$KCPP_CUDA" == 11.* && -z "$ARCHES_CU11$ARCHES_CU12$ARCHES_CU13" ]]; then
	ARCHES_CU11=1
fi

LLAMA_NOAVX1_FLAG=""
LLAMA_NOAVX2_FLAG=""
ARCHES_FLAG=""
if [ -n "$NOAVX2" ]; then
	LLAMA_NOAVX2_FLAG="LLAMA_NOAVX2=1"
fi
if [ -n "$NOAVX1" ]; then
	LLAMA_NOAVX1_FLAG="LLAMA_NOAVX1=1"
fi
if [ -n "$ARCHES_CU11" ]; then
	ARCHES_FLAG="LLAMA_ARCHES_CU11=1"
fi
if [ -n "$ARCHES_CU12" ]; then
	ARCHES_FLAG="LLAMA_ARCHES_CU12=1"
fi
if [ -n "$ARCHES_CU13" ]; then
	ARCHES_FLAG="LLAMA_ARCHES_CU13=1"
fi

# CUDA's native architecture flag needs a visible GPU; old-CPU builds need the fallback libraries.
if [[ "$KCPP_CUDA" != "rocm" && "$NVIDIA_GPU" = 0 ]] || [[ "$ARCH" = x64 && -n "$NOAVX1$NOAVX2" ]]; then
	KCPP_PORTABLE=1
fi

if [ -n "$KCPP_PORTABLE" ]; then
	LLAMA_PORTABLE_FLAG="LLAMA_PORTABLE=1"
	# The Makefile generates these fallback libraries only on x86.
	if [ "$ARCH" = x64 ]; then
		PORTABLE_SO="--add-data ./koboldcpp_failsafe.so:. --add-data ./koboldcpp_noavx2.so:. --add-data ./koboldcpp_vulkan_noavx2.so:."
		VULKAN_FAILSAFE_SO="--add-data ./koboldcpp_vulkan_failsafe.so:."
	fi
fi

if [ "$KCPP_CUDA" = "rocm" ]; then
	bin/micromamba run -r conda -p conda/envs/linux make -j$(nproc) LLAMA_VULKAN=1 LLAMA_HIPBLAS=1 LLAMA_USE_BUNDLED_GLSLC=1 LLAMA_ADD_CONDA_PATHS=1 $LLAMA_PORTABLE_FLAG $LLAMA_NOAVX1_FLAG $LLAMA_NOAVX2_FLAG $ARCHES_FLAG
else
	bin/micromamba run -r conda -p conda/envs/linux make -j$(nproc) LLAMA_VULKAN=1 LLAMA_CUBLAS=1 LLAMA_USE_BUNDLED_GLSLC=1 LLAMA_ADD_CONDA_PATHS=1 $LLAMA_PORTABLE_FLAG $LLAMA_NOAVX1_FLAG $LLAMA_NOAVX2_FLAG $ARCHES_FLAG
fi

if [ $? -ne 0 ]; then
    echo "Error: make failed."
    exit 1
fi

if [[ $1 == "rebuild" ]]; then
	echo Rebuild complete, you can now try to launch Koboldcpp.
elif [[ $1 == "dist" ]]; then
	if [ ! -n "$KCPP_PORTABLE" ]; then
		echo "WARNING: KCPP_PORTABLE NOT SPECIFIED, THIS BINARY WILL ONLY RUN ON YOUR SYSTEM!!"
		sleep 5
	fi
	bin/micromamba remove --no-rc -r conda -p conda/envs/linux --force ocl-icd -y
	bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onedir --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --clean --console koboldcpp.py -n "koboldcpp-launcher"
	if [ "$KCPP_CUDA" = "rocm" ]; then
		if [ ! -n "$ROCM_PATH" ]; then
			ROCM_PATH=/opt/rocm
		fi
		if [[ "$ARCH" = x64 && -n "$NOAVX1" ]]; then
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_hipblas.so:.' $PORTABLE_SO $VULKAN_FAILSAFE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --add-data "$ROCM_PATH/lib/rocblas:." --add-data "$ROCM_PATH/lib/libamd_comgr.so:." --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH-rocm"
		elif [[ "$ARCH" = x64 && -n "$NOAVX2" ]]; then
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_hipblas.so:.' $PORTABLE_SO $VULKAN_FAILSAFE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --add-data "$ROCM_PATH/lib/rocblas:." --add-data "$ROCM_PATH/lib/libamd_comgr.so:." --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH-rocm"
		else
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_default.so:.' --add-data './koboldcpp_hipblas.so:.' --add-data './koboldcpp_vulkan.so:.' $PORTABLE_SO $VULKAN_FAILSAFE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --add-data "$ROCM_PATH/lib/rocblas:." --add-data "$ROCM_PATH/lib/libamd_comgr.so:." --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH-rocm"
		fi
	else
		bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onedir --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --clean --console koboldcpp.py -n "koboldcpp-launcher"
		if [[ "$ARCH" = x64 && -n "$NOAVX1" ]]; then
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_cublas.so:.' $PORTABLE_SO $VULKAN_FAILSAFE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH$KCPP_CUDAAPPEND"
		elif [[ "$ARCH" = x64 && -n "$NOAVX2" ]]; then
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_cublas.so:.' $PORTABLE_SO $VULKAN_FAILSAFE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH$KCPP_CUDAAPPEND"
		else
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_default.so:.' --add-data './koboldcpp_cublas.so:.' --add-data './koboldcpp_vulkan.so:.' $PORTABLE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH$KCPP_CUDAAPPEND"
			bin/micromamba run -r conda -p conda/envs/linux pyinstaller --noconfirm --onefile --collect-all customtkinter --collect-all jinja2 --collect-all psutil --add-data './dist/koboldcpp-launcher/koboldcpp-launcher:.' --add-data './koboldcpp_default.so:.' --add-data './koboldcpp_vulkan.so:.' $PORTABLE_SO --add-data './kcpp_adapters:./kcpp_adapters' --add-data './koboldcpp.py:.' --add-data './kcpp_agent.py:.' --add-data './json_to_gbnf.py:.' --add-data './LICENSE.md:.' --add-data './MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md:.' --add-data './embd_res:./embd_res' --clean --console koboldcpp.py -n "koboldcpp-linux-$ARCH-nocuda$KCPP_APPEND"
		fi
	fi
	bin/micromamba install --no-rc -r conda -p conda/envs/linux ocl-icd -c conda-forge -y
else
	bin/micromamba run -r conda -p conda/envs/linux python koboldcpp.py "$@"
fi