#!/usr/bin/env bash

{
	SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
	cd "$SCRIPT_DIR"

	source .colorfunctions.sh

	os_name=$(uname -s)

	declare -i exit_code=0

	function set_debug {
		trap 'echo -e "${CYAN}$(date +"%Y-%m-%d %H:%M:%S")${NC} ${MAGENTA}| Line: $LINENO ${NC}${YELLOW}-> ${NC}${BLUE}[DEBUG]${NC} ${GREEN}$BASH_COMMAND${NC}"' DEBUG
	}

	function unset_debug {
		trap - DEBUG
	}

	install_those=()

	if [ ! -f "$SCRIPT_DIR/requirements.txt" ]; then
		echo "The file $SCRIPT_DIR/requirements.txt doesn't exist."
		exit 21
	fi

	if [ ! -f "$SCRIPT_DIR/test_requirements.txt" ]; then
		echo "The file $SCRIPT_DIR/test_requirements.txt doesn't exist."
		exit 21
	fi

	while IFS= read -r line; do
		install_those+=("$line")
	done < "$SCRIPT_DIR/requirements.txt"

	if [[ -n $install_tests ]]; then
		while IFS= read -r line; do
			install_those+=("$line")
		done < "$SCRIPT_DIR/test_requirements.txt"
	fi

	FROZEN=""

	function displaytime {
		local T=$1
		local D=$((T/60/60/24))
		local H=$((T/60/60%24))
		local M=$((T/60%60))
		local S=$((T%60))
		(( $D > 0 )) && printf '%d days ' $D
		(( $H > 0 )) && printf '%d hours ' $H
		(( $M > 0 )) && printf '%d minutes ' $M
		(( $D > 0 || $H > 0 || $M > 0 )) && printf 'and '
		printf '%d seconds\n' $S
	}

	function error_message {
		if command -v resize 2>/dev/null >/dev/null; then
			eval "$(resize)"
		fi
		MSG=$1
		echo_red "$MSG"

		if command -v whiptail 2>/dev/null >/dev/null; then
			export NEWT_COLORS='
			window=,red
			border=white,red
			textbox=white,red
			button=black,white
			'
			whiptail --title "Error Message" --scrolltext --msgbox "$MSG" $LINES $COLUMNS $(( $LINES - 8 ))
			export NEWT_COLORS=""
		else
			echo_red "Whiptail not found. Try installing it, for example, with apt-get install whiptail"
		fi
	}

	export RUN_VIA_RUNSH=1

	export PYTHONDONTWRITEBYTECODE=1

	IFS=$'\n'

	Green='\033[0;32m'
	Color_Off='\033[0m'
	Red='\033[0;31m'

	set -e

	LMOD_DIR=/software/foundation/$(uname -m)/lmod/lmod/libexec

	myml () {
		if [[ -e $LMOD_DIR/ml_cmd ]]; then
			eval "$($LMOD_DIR/ml_cmd "$@")" 2>/dev/null >/dev/null
		fi
	}

	if [ -z "$LOAD_MODULES" ] || [ "$LOAD_MODULES" -eq 1 ]; then
		if [[ -n "$CLUSTERHOST" && "$CLUSTERHOST" == *capella* ]]; then
			LMOD_DIR=/usr/share/lmod/lmod/libexec

			myml release/24.04 GCCcore/12.3.0 Python/3.11.3 Tkinter/3.11.3 PostgreSQL/16.1
		else
			myml release/23.04 GCCcore/12.2.0 Python/3.10.8 GCCcore/11.3.0 Tkinter/3.10.4 PostgreSQL/14.4

			if [[ $(uname -m) == "ppc64le" ]]; then
				myml zlib/1.2.12 GCC/12.2.0 OpenBLAS/0.3.21
			fi
		fi
	fi

	#if command -v sacct 2>/dev/null >/dev/null; then
	#	export PATH="$SCRIPT_DIR/.tools/:$PATH"
	#fi

	if ! command -v python3 2>/dev/null >/dev/null; then
		red_text "python3 not installed. Cannot continue."
		exit 245
	fi

	VENV_DIR_NAME=".omniax_$(uname -m)_$(python3 --version | sed -e 's# #_#g')"

	ROOT_VENV_DIR=$HOME

	if [[ -n $root_venv_dir ]] && [[ -d $root_venv_dir ]]; then
		ROOT_VENV_DIR=$root_venv_dir
	fi

	VENV_DIR=$ROOT_VENV_DIR/$VENV_DIR_NAME

	INSTALL_STUFF=1
	CUSTOM_VIRTUAL_ENV=0

	if [[ -n $VIRTUAL_ENV ]]; then
		VENV_DIR=$VIRTUAL_ENV
		INSTALL_STUFF=0
		CUSTOM_VIRTUAL_ENV=1
	fi

	export CUSTOM_VIRTUAL_ENV

	export VENV_DIR

	NUMBER_OF_INSTALLED_MODULES=0
	PROGRESSBAR=""

	generate_progress_bar_setup() {
		local total_nr_modules=$1

		NUMBER_OF_INSTALLED_MODULES=$(get_nr_of_already_installed_modules)

		if ! [[ "$NUMBER_OF_INSTALLED_MODULES" =~ ^[0-9]+$ ]]; then
			echo "Error: NUMBER_OF_INSTALLED_MODULES must be a positive integer, but is $NUMBER_OF_INSTALLED_MODULES." >&2
			return 1
		fi

		if ! [[ "$total_nr_modules" =~ ^[0-9]+$ ]]; then
			echo "Error: total_nr_modules must be a positive integer, but is $total_nr_modules." >&2
			return 1
		fi

		if [ "$NUMBER_OF_INSTALLED_MODULES" -gt "$total_nr_modules" ]; then
			echo "Error: the current progress cannot exceed the total progress ($NUMBER_OF_INSTALLED_MODULES/$total_nr_modules)." >&2
			return 1
		fi

		# Call the generate_progress_bar function to print the progress bar
		generate_progress_bar "$NUMBER_OF_INSTALLED_MODULES" "$total_nr_modules"
	}

	generate_progress_bar() {
		local current="$1"
		local max="$2"

		if ! [[ "$current" =~ ^[0-9]+$ ]]; then
			echo "Error: current must be positive integer, but is $current." >&2
			return 1
		fi

		if ! [[ "$max" =~ ^[0-9]+$ ]]; then
			echo "Error: max must be positive integer, but is $max." >&2
			return 1
		fi

		if [ "$current" -gt "$max" ]; then
			echo "Error: the current progress cannot exceed the total progress ($current/$max)." >&2
			return 1
		fi

		local bar_length=30
		local filled_length=$((bar_length * current / max))
		local empty_length=$((bar_length - filled_length))
		local percentage=$((current * 100 / max))

		bar_char="#"

		if [[ "$os_name" == "Linux" ]]; then
			bar_char="━"
		fi

		local bar=""
		for ((i = 0; i < filled_length; i++)); do
			bar="${bar}$bar_char"
		done
		for ((i = 0; i < empty_length; i++)); do
			bar="${bar} "
		done

		printf "[%s] %d%%\n" "$bar" "$percentage"
	}

	function ppip {
		MODULE=$1
		AS_REQUIREMENT_OF=$2
		NUMBER_OF_MAIN_MODULES=$3

		set +e

		FROZEN=$(pip --disable-pip-version-check list --format=freeze)

		PROGRESSBAR=$(generate_progress_bar_setup "$NUMBER_OF_MAIN_MODULES")

		if [[ -z $CI ]]; then
			green_reset_line "${PROGRESSBAR}➤Installing $MODULE "
		fi


		MODULES_WITHOUT_VERSIONS=$(echo "$MODULE" | sed -e 's#[=<>]=.*##' -e 's#~.*##')

		echo "$FROZEN" | grep -i "$MODULES_WITHOUT_VERSIONS" 2>/dev/null >/dev/null
		_exit_code=$?

		if [[ "$_exit_code" != "0" ]]; then
			if [[ "$MODULE" != "$AS_REQUIREMENT_OF" ]] && [[ "$AS_REQUIREMENT_OF" != "-" ]]; then
				k=0

				for i in $(pip3 install --disable-pip-version-check --dry-run "$MODULE" | grep -v "already satisfied" | grep "Collecting" | sed -e 's#Collecting ##' | grep -v "^$MODULE$"); do
					if [[ "$i" != "$MODULE" ]]; then
						if [[ $k -eq 0 ]]; then
							green_reset_line "${PROGRESSBAR}➤Installing requirements for $MODULE"
						fi
						ppip "$i" "$MODULE" "$NUMBER_OF_MAIN_MODULES" || {
							red_reset_line "❌Failed to install $i."

							exit 20
						}

						k=$((k+1))
					fi
				done

				if [[ $k -gt 0 ]]; then
					green_reset_line "${PROGRESSBAR}➤Installed all requirements for $MODULE, now installing the package itself..."
				fi
			fi

			green_reset_line "${PROGRESSBAR}➤Installing $MODULE..."
			mkdir -p logs
			export PIP_DISABLE_PIP_VERSION_CHECK=1
			INSTALL_ERRORS_FILE="logs/install_errors"

			if [[ -n $RUN_UUID ]]; then
				INSTALL_ERRORS_FILE="logs/${RUN_UUID}_install_errors"
			fi

			if [[ -z $DEBUG ]]; then
				pip3 --disable-pip-version-check install -q $MODULE >&2 2>> $INSTALL_ERRORS_FILE || pip3 --disable-pip-version-check install $MODULE >&2 2>> $INSTALL_ERRORS_FILE || {
					red_reset_line "❌Failed to install $MODULE. Check $INSTALL_ERRORS_FILE"

					if [[ -n $CI ]] || { [[ -f /proc/self/cgroup ]] && grep -qE '/docker|/lxc' /proc/self/cgroup; }; then
						cat "$INSTALL_ERRORS_FILE" 
					fi

					exit 20
				}
			else
				pip3 --disable-pip-version-check install $MODULE >&2 2>> $INSTALL_ERRORS_FILE || pip3 --disable-pip-version-check install $MODULE >&2 2>> $INSTALL_ERRORS_FILE || {
					red_reset_line "❌Failed to install $MODULE. Check $INSTALL_ERRORS_FILE"

					if [[ -n $CI ]] || { [[ -f /proc/self/cgroup ]] && grep -qE '/docker|/lxc' /proc/self/cgroup; }; then
						cat "$INSTALL_ERRORS_FILE" 
					fi

					exit 20
				}
			fi

			FROZEN=$(pip --disable-pip-version-check list --format=freeze)

			PROGRESSBAR=$(generate_progress_bar_setup "$NUMBER_OF_MAIN_MODULES")

			if [[ -z $CI ]]; then
				green_reset_line "${PROGRESSBAR}✅$MODULE installed successfully"
			fi
		fi
		set -e
	}

	get_nr_of_already_installed_modules () {
		nr=0
		for key in "${install_those[@]}"; do
			noversion=$(echo "$key" | sed -e 's#[=<>]=.*##' -e 's#~.*##')

			if [[ -z $FROZEN ]]; then
				FROZEN=$(pip --disable-pip-version-check list --format=freeze)
			fi

			if [[ $noversion == "rich_argparse" ]]; then
				noversion="rich-argparse"
			fi

			if echo "$FROZEN" | grep -i "$noversion" 2>/dev/null >/dev/null; then
				nr=$(($nr+1))
			fi
		done

		echo "$nr"
	}

	function install_required_modules {
		green_reset_line "➤Checking environment $VENV_DIR..."
		MAX_NR="${#install_those[@]}"
		NUMBER_OF_INSTALLED_MODULES=$(get_nr_of_already_installed_modules)

		PROGRESSBAR=$(generate_progress_bar_setup "$MAX_NR")

		for key in "${!install_those[@]}"; do
			install_this=${install_those[$key]}
			PROGRESSBAR=$(generate_progress_bar_setup "$MAX_NR")
			if [[ -z $CI ]]; then
				green_reset_line "${PROGRESSBAR}➤Checking if $install_this is installed..."
			fi

			if ! echo "$FROZEN" | grep -q "$install_this"; then
				ppip "$install_this" "-" "$MAX_NR"
			fi
		done

		_tput cr
		_tput el

		green_reset_line "✅Environment checking done!"
		_tput cr
		_tput el
	}

	required_programs=(
		"stdbuf:coreutils"
		"base64:base64"
		"curl:curl"
		"wget:wget"
		"uuidgen:uuid-runtime"
		"python3:python3"
		"gcc:gcc"
		"resize:xterm"
		"cat:coreutils"
		"ls:coreutils"
		"wget:wget"
		"grep:grep"
		"tput:ncurses-bin"
		"sed:sed"
	)

	if command -v apt 2>/dev/null >/dev/null; then
		required_programs+="findmnt:util-linux"
		required_programs+="whiptail:whiptail"
	fi

	not_found_programs=0

	for cmd_pkg in "${required_programs[@]}"; do
		cmd="${cmd_pkg%%:*}"
		pkg="${cmd_pkg##*:}"

		if ! command -v "$cmd" >/dev/null 2>&1; then
			red_text "❌$cmd not found. Try installing it with 'sudo apt-get install $pkg' (depending on your distro)\n"
			not_found_programs=$(($not_found_programs+1))
		fi
	done
	
	if [[ $not_found_programs -ne 0 ]]; then
		exit 11
	fi

	if [[ "$SCRIPT_DIR" != *"$VENV_DIR"* ]]; then
		if [[ ! -d "$VENV_DIR" ]]; then
			if ! python3 -c 'from distutils.sysconfig import get_makefile_filename as m; from os.path import isfile; import sys ; sys.exit(not isfile(m()))' >/dev/null 2>/dev/null; then
				red_text "❌python3 header files not found. Try installing them, for example, with 'sudo apt-get install python3-dev' (depending on your distro)\n"
				if [[ "$OSTYPE" == "darwin"* ]]; then
					red_text "Not exiting because I am not sure if you need it on Macs"
				else
					exit 22
				fi
			fi

			green_reset_line "${PROGRESSBAR}➤Environment $VENV_DIR was not found. Creating it..."
			python3 -mvenv "$VENV_DIR/" || {
				red_text "❌Failed to create Virtual Environment in $VENV_DIR"
				exit 20
			}

			green_reset_line "✅Virtual Environment $VENV_DIR created. Activating it..."

			if [[ -e "$VENV_DIR/bin/activate" ]]; then
				source "$VENV_DIR/bin/activate" || {
					red_text "❌Failed to activate $VENV_DIR"
					exit 20
				}
			else
				red_text "❌Failed to activate $VENV_DIR"
				exit 20
			fi

			green_reset_line "✅Virtual Environment activated. Now installing software. This may take some time."

		fi
	fi

	if [[ -e "$VENV_DIR/bin/activate" ]]; then
		source "$VENV_DIR/bin/activate" || {
			red_reset_line "❌Failed to activate $VENV_DIR. Deleting venv and creating it again..."
			rm -rf "$VENV_DIR"

			python3 -mvenv "$VENV_DIR/" || {
				red_text "❌Failed to create Virtual Environment in $VENV_DIR"
				rm -rf "$VENV_DIR"
				exit 20
			}

			source "$VENV_DIR/bin/activate" || {
				red_reset_line "❌Failed to activate recreated $VENV_DIR. Deleting venv and NOT trying again..."
				exit 20
			}

			install_required_modules
		}
	else
		red_reset_line "❌Failed to activate $VENV_DIR. Deleting venv and creating it again..."
		rm -rf "$VENV_DIR"

		python3 -mvenv "$VENV_DIR/" || {
			red_text "❌Failed to create Virtual Environment in $VENV_DIR"
			exit 20
		}

		if [[ -e "$VENV_DIR/bin/activate" ]]; then
			source "$VENV_DIR/bin/activate" || {
				red_reset_line "❌Failed to activate recreated $VENV_DIR. Deleting venv and NOT trying again..."
				rm -rf "$VENV_DIR"
				exit 20
			}

			downgrade_output=$(pip3 --disable-pip-version-check install -q pip==24.0) || {
				red_text "Failed to downgrade pip. Output:"
				red_text "$downgrade_output"
			}
		else
			red_reset_line "❌Failed to activate recreated $VENV_DIR. Deleting venv and NOT trying again..."
			rm -rf "$VENV_DIR"
			exit 20
		fi

		install_required_modules

	fi

	if [[ $INSTALL_STUFF -eq 1 ]]; then
		if [[ -z $DONT_INSTALL_MODULES ]]; then
			if [[ -z $SLURM_JOB_ID ]]; then
				REQUIREMENTS_HASH_MAIN=$(md5sum requirements.txt | awk '{print $1}')
				REQUIREMENTS_HASH_TEST=$(md5sum test_requirements.txt | awk '{print $1}')

				hash_is_different() {
					hash_file=$1
					required_hash=$2

					if [[ -f "$hash_file" ]]; then
						stored_hash=$(cat "$hash_file")
						if [[ "$stored_hash" == "$required_hash" ]]; then
							return 0
						else
							return 1
						fi
					else
						return 2
					fi
				}


				if ! hash_is_different "$VENV_DIR/hash" "$REQUIREMENTS_HASH_MAIN" || ! hash_is_different "$VENV_DIR/hash_test" "$REQUIREMENTS_HASH_TEST"; then
					set +e
					FROZEN=$(pip --disable-pip-version-check list --format=freeze)
					exit_code_pip=$?
					set -e

					if [[ "$exit_code_pip" -ne "0" ]]; then
						printf "pip list --format=freeze exited with exit code %s\n" $exit_code_pip
						exit 12
					fi

					install_required_modules

					echo "$REQUIREMENTS_HASH_MAIN" > "$VENV_DIR/hash"
					echo "$REQUIREMENTS_HASH_TEST" > "$VENV_DIR/hash_test"
				fi

			fi
		else
			if [[ -z $DONT_SHOW_DONT_INSTALL_MESSAGE ]]; then
				red_text "\$DONT_INSTALL_MODULES is set. Don't install modules.\n"
			fi
		fi
	fi

	export PYTHONPATH=$VENV_DIR:$PYTHONPATH
}
