#!/usr/bin/env bash

set -e # if any command fails, stop immediately

{ # this ensures the entire script is downloaded #

CVMT_PKGDIR="/opt/cvmt"
CVMT_CDNURL="https://svsupport.blob.core.windows.net/files/VisionSense_Trainer/Packages"
CVMT_APTPKG="ca-certificates curl gpg"
CVMT_DOTENV=""
CVMT_BASEON=""
CVMT_DBTYPE=""
CVMT_DBNAME="cvmt"
CVMT_DBUSER="cvmt"
CVMT_DBPASS=
CVMT_DBHOST="pg"
CVMT_DBPORT="5432"
CVMT_GWTYPE=""

cvmt_console_color() {
    local color_code="$1"
    local text="$2"
    echo -e "\033[${color_code}m${text}\033[0m"
}

cvmt_console_error() {
    local text="$1"
    cvmt_console_color "31" "$text"
}

cvmt_console_title() {
    local text="$1"
    cvmt_console_color "36" "$text"
}

cvmt_console_light() {
    local text="$1"
    cvmt_console_color "96" "$text"
}

cvmt_console_event() {
    local text="$1"
    cvmt_console_color "32" "$text"
}

cvmt_network_ipaddress() {
    local interface="$1"
    ip addr show "$interface" | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1
}

cvmt_welcome() {
    local accept_eula="n"

    cvmt_console_title "Welcome to CVMT Installer"
    echo "This script will help you install CVMT on your system."
    echo "Please follow the prompts to complete the installation."
    echo

    cvmt_console_title "End User License Agreement (EULA)"
    echo "Please read the EULA carefully before proceeding."
    echo ""
    echo "https://svsupport.blob.core.windows.net/files/VisionSense_Trainer/Packages/EULA.txt"
    echo ""
    echo "By continuing, you agree to the terms and conditions of the EULA."
    echo
    read -r -p "Do you accept the EULA? (y/N): " accept_eula </dev/tty
    if [ "$accept_eula" != "y" ] && [ "$accept_eula" != "Y" ]; then
        cvmt_console_error "You must accept the EULA to proceed with the installation. The installation will now abort."
        exit 1
    fi
}

cvmt_require_is_advenv() {
    if [[ ! -f ~/advantech/advantech/nginx/docker-compose.yml ]]; then
        return 1
    else
        return 0
    fi
}

cvmt_require_is_wslenv() {
    case "$(uname -r)" in
	*microsoft* ) true ;; # WSL 2
	*Microsoft* ) true ;; # WSL 1
	* ) false;;
	esac
}

cvmt_require() {
    if cvmt_require_is_advenv; then
        cvmt_console_event "Detected Advantech Environment."
        CVMT_BASEON="ADVENV"
        CVMT_APTPKG+=" ubuntu-drivers-common htop net-tools"
    elif cvmt_require_is_wslenv; then
        cvmt_console_event "Detected WSL Environment."
        CVMT_BASEON="WSLENV"
    else
        cvmt_console_error "Unsupported Environment. This installer only supports Advantech Environment or WSL Environment."
        exit 1
    fi

    if [ "$CVMT_BASEON" = "WSLENV" ]; then
        nvidia-smi || {
            cvmt_console_error "NVIDIA GPU not detected in WSL. Please ensure your system has a compatible NVIDIA GPU and the latest drivers are installed on the Windows host."
        }
    fi
}

cvmt_setting_search() {
    local envfile="$CVMT_PKGDIR/.env"

    if [ -f "$envfile" ]; then
        while IFS='=' read -r key value; do
            case "$key" in
                'POSTGRESQL_PASSWORD') CVMT_DBPASS=$(echo $value | sed "s/'//g") ;;
            esac
        done < "$envfile"
    fi
}

cvmt_setting_advenv() {
    echo "Using advantech database."
    CVMT_DBTYPE="advantech"
    CVMT_DBHOST="pg14"
    CVMT_DBPORT="5432"
    CVMT_DBNAME="cvmt"
    CVMT_DBUSER="advantech"
    CVMT_DBPASS="\$t@k@Pp\$"
    echo "Using advantech ingress."
    CVMT_GWTYPE="advantech"
}

cvmt_setting_wslenv() {
    local newpass=$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 32 ; echo '')

    echo "Using internal database."
    CVMT_DBTYPE="internal"
    CVMT_DBHOST="pg"
    CVMT_DBPORT="5432"
    CVMT_DBNAME="cvmt"
    CVMT_DBUSER="cvmt"
    CVMT_DBPASS="$newpass"
    echo "Using internal ingress."
    CVMT_GWTYPE="internal"
    cvmt_setting_search
}

cvmt_setting() {
    if [ "$CVMT_BASEON" = "ADVENV" ]; then
        cvmt_setting_advenv
    elif [ "$CVMT_BASEON" = "WSLENV" ]; then
        cvmt_setting_wslenv
    fi
}

cvmt_install_docker_engine() {
    cvmt_console_title "Installing Docker"
    echo "This will install Docker CE and Docker Compose on your system."
    echo "If you already have Docker installed, it will be updated to the latest version."
    echo

    curl -fsSL https://get.docker.com | sh
    sudo usermod -aG docker $USER
}

cvmt_install_docker_before() {
    cvmt_console_title "Installing Prerequisites"
    echo "Installing required packages/tools for CVMT."
    echo

    sudo apt-get update
    sudo apt --fix-broken install
    sudo apt-get install -y --no-install-recommends $CVMT_APTPKG
}

cvmt_install_docker() {
    cvmt_install_docker_before
    cvmt_install_docker_engine
}

cvmt_install_nvidia_docker() {
    cvmt_console_title "Installing NVIDIA Container Toolkit"
    echo "This will install the NVIDIA Container Toolkit to enable GPU support in Docker."
    echo "If you already have the NVIDIA Container Toolkit installed, it will be updated to the latest version."
    echo

    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --yes --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
    curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
        sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
        sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

    sudo apt-get update
    sudo apt-get install -y --no-install-recommends nvidia-container-toolkit

    sudo nvidia-ctk runtime configure --runtime=docker
    sudo systemctl restart docker
}

cvmt_install_nvidia() {
    if [ "$CVMT_BASEON" = "ADVENV" ]; then
        cvmt_console_error "NVIDIA installation for Advantech Environment is not able to do now."
        echo
        echo "Please install NVIDIA drivers and Docker NVIDIA Container Toolkit manually after installation."
        echo
    elif [ "$CVMT_BASEON" = "WSLENV" ]; then
        cvmt_install_nvidia_docker
    else
        cvmt_console_error "NVIDIA installation for this environment is not implemented yet."
        exit 1
    fi
}

cvmt_install_option_forall() {
    CVMT_DOTENV=""
    CVMT_DOTENV+="NVIDIA_VISIBLE_DEVICES=all\n"
    CVMT_DOTENV+="WORKFLOW_SHMSIZE=4g\n"
    CVMT_DOTENV+="WORKFLOW_MEMLIMIT=4g\n"
    CVMT_DOTENV+="WORKFLOW_MEMLIMIT_ANNO=8g\n"
    CVMT_DOTENV+="WORKFLOW_MEMLIMIT_TRAIN=12g\n"
    CVMT_DOTENV+="WORKFLOW_MEMLIMIT_EXPORT=8g\n"
    CVMT_DOTENV+="WORKFLOW_CPULIMIT=30%\n"
    CVMT_DOTENV+="WORKFLOW_CPULIMIT_ANNO=50%\n"
    CVMT_DOTENV+="WORKFLOW_CPULIMIT_TRAIN=75%\n"
    CVMT_DOTENV+="WORKFLOW_CPULIMIT_EXPORT=50%\n"
}

cvmt_install_option_advenv() {
    cvmt_install_option_forall
    
    CVMT_DOTENV+="POSTGRE_HOST=$CVMT_DBHOST:$CVMT_DBPORT/$CVMT_DBNAME\n"
    CVMT_DOTENV+="POSTGRE_USER=$CVMT_DBUSER\n"
    CVMT_DOTENV+="POSTGRE_PWD='$CVMT_DBPASS'\n"

    CVMT_DOTENV+="CVMT_LINKNETWORK=advantech\n"
    CVMT_DOTENV+="CVMT_BUILDTAG=latest\n"
    CVMT_DOTENV+="CVMT_INGRESS=false\n"

    CVMT_DOTENV+="COMPOSE_PROJECT_NAME=cvmt\n"
    CVMT_DOTENV+="COMPOSE_PATH_SEPARATOR=:\n"
    CVMT_DOTENV+="COMPOSE_PROFILES=cv\n"
    CVMT_DOTENV+="COMPOSE_FILE=$CVMT_PKGDIR/compose-cv.yml:$CVMT_PKGDIR/compose-hw.yml\n"
}

cvmt_install_option_wslenv() {
    cvmt_install_option_forall

    CVMT_DOTENV+="POSTGRE_HOST=$CVMT_DBHOST:$CVMT_DBPORT/$CVMT_DBNAME\n"
    CVMT_DOTENV+="POSTGRE_USER=$CVMT_DBUSER\n"
    CVMT_DOTENV+="POSTGRE_PWD='$CVMT_DBPASS'\n"

    CVMT_DOTENV+="POSTGRESQL_DATABASE=$CVMT_DBNAME\n"
    CVMT_DOTENV+="POSTGRESQL_USERNAME=$CVMT_DBUSER\n"
    CVMT_DOTENV+="POSTGRESQL_PASSWORD='$CVMT_DBPASS'\n"

    CVMT_DOTENV+="PGADMIN_DEFAULT_EMAIL=pg@admin.local\n"
    CVMT_DOTENV+="PGADMIN_DEFAULT_PASSWORD=pg@dm1n\n"

    CVMT_DOTENV+="CVMT_BUILDTAG=lts\n"
    CVMT_DOTENV+="CVMT_LINKNETWORK=db\n"
    CVMT_DOTENV+="CVMT_DOMAINNAME=local\n"
    CVMT_DOTENV+="CVMT_INGRESS=true\n"

    CVMT_DOTENV+="COMPOSE_PROJECT_NAME=cvmt\n"
    CVMT_DOTENV+="COMPOSE_PATH_SEPARATOR=:\n"
    CVMT_DOTENV+="COMPOSE_PROFILES=cv\n"
    CVMT_DOTENV+="COMPOSE_FILE=$CVMT_PKGDIR/compose-cv.yml:$CVMT_PKGDIR/compose-db.yml:$CVMT_PKGDIR/compose-ep.yml\n"
}

cvmt_install_option() {
    if [ -f $CVMT_PKGDIR/.env ]; then
        cvmt_console_event "Stopping existing CVMT services..."
        sudo docker compose --env-file $CVMT_PKGDIR/.env down
    fi

    if [ "$CVMT_BASEON" = "ADVENV" ]; then
        cvmt_install_option_advenv
    elif [ "$CVMT_BASEON" = "WSLENV" ]; then
        cvmt_install_option_wslenv
    fi

    if [ -z "$CVMT_DOTENV" ]; then
        cvmt_console_error "Generated .env content is empty. Installation aborted."
        exit 1
    else
        local Out="$CVMT_PKGDIR/.env"
        sudo mkdir -p $CVMT_PKGDIR
        echo -e "$CVMT_DOTENV" | sudo tee $Out > /dev/null
        sudo chown $(whoami):$(whoami) $Out
        sudo chmod 600 $Out
    fi
}

cvmt_install_config_fetch() {
    local Url="$CVMT_CDNURL/$1"
    local Out="$CVMT_PKGDIR/$1"

    sudo curl -L $Url -o $Out
    sudo chown $(whoami):$(whoami) $Out
    sudo chmod 600 $Out

    if [ ! -f $Out ]; then
        local delay=30
        cvmt_console_error "Failed to download $1 from $Url, Retrying in $delay seconds..."
        sleep $delay
        cvmt_install_config_fetch "$1"
    fi
}

cvmt_install_config_advenv() {
    cvmt_install_config_fetch "compose-cv.yml"
    cvmt_install_config_fetch "compose-hw.yml"
}

cvmt_install_config_wslenv() {
    cvmt_install_config_fetch "compose-cv.yml"
    cvmt_install_config_fetch "compose-db.yml"
    cvmt_install_config_fetch "compose-ep.yml"
}

cvmt_install_config() {
    cvmt_console_title "Downloading Compose Stack Configuration"
    echo "This will download the Docker Compose stack configuration files for CVMT."
    echo

    if [ "$CVMT_BASEON" = "ADVENV" ]; then
        cvmt_install_config_advenv
    elif [ "$CVMT_BASEON" = "WSLENV" ]; then
        cvmt_install_config_wslenv
    else
        cvmt_console_error "Compose stack installation for this environment is not implemented yet."
        exit 1
    fi

    if [ ! -f $CVMT_PKGDIR/.env ]; then
        cvmt_console_error ".env file not found in $CVMT_PKGDIR. Installation aborted."
        exit 1
    fi

    while ! sudo docker compose --env-file $CVMT_PKGDIR/.env pull; do
        local delay=30
        cvmt_console_error "Fail to pull docker images. Retrying in $delay seconds..."
        sleep $delay
    done
}

cvmt_install() {
    cvmt_install_docker
    cvmt_install_nvidia
    cvmt_install_option
    cvmt_install_config
}

cvmt_cleanup_advenv() {
    echo "Continue to setup other packages for Advantech Environment..."
}

cvmt_cleanup_wslenv() {
    cd $CVMT_PKGDIR
    sudo docker compose --env-file $CVMT_PKGDIR/.env up -d
    echo

    cvmt_console_title "Thank you for trying VisionSense Trainer!"
    local ipaddr=$(cvmt_network_ipaddress "eth0")
    echo
    echo "To access Trainer in Windows Host, please run the following command in PowerShell as Administrator:"
    echo
    cvmt_console_light "Add-Content -Path \"C:/Windows/System32/drivers/etc/hosts\" -Value \"$ipaddr trainer.local\""
    echo
    echo "You can verify the entry has been added by running:"
    echo
    cvmt_console_light "Resolve-DnsName trainer.local"
    echo
    echo "Then open your web browser in Windows Host and go to:"
    echo
    cvmt_console_light "http://trainer.local"
    echo
}

cvmt_cleanup() {
    cvmt_console_title "Installation Completed"
    echo "CVMT has been successfully installed on your system."
    echo

    if [ "$CVMT_BASEON" = "ADVENV" ]; then
        cvmt_cleanup_advenv
    elif [ "$CVMT_BASEON" = "WSLENV" ]; then
        cvmt_cleanup_wslenv
    else
        cvmt_console_error "Cleanup for this environment is not implemented yet."
        exit 1
    fi
}

cvmt_welcome

cvmt_require

cvmt_setting

cvmt_install

cvmt_cleanup

} # this ensures the entire script is downloaded #