#!/usr/bin/env bash
script_version='1.6'
ubuntu_ver=$(cat /etc/os-release | grep 'VERSION_ID'|cut -f2 -d "\""|cut -f1 -d ".")

function os_support_check() {
    LIST=$1
    DELIMITER=$2
    VALUE=$3
    echo $LIST | tr "$DELIMITER" '\n' | grep -F -q -x "$VALUE"
}

if [[ $ubuntu_ver = "23" || $ubuntu_ver = "22" || $ubuntu_ver = "20" ]]; then
    pkgs='libgtk-3-0 libnotify4 libnss3 libxss1 libxtst6 xdg-utils libatspi2.0-0 kde-cli-tools trash-cli libglib2.0-bin xvfb libsqlcipher-dev v4l-utils xdotool'
elif [ $ubuntu_ver = "18" ]; then
    pkgs='libgtk-3-0 libnotify4 libnss3 libxss1 libxtst6 xdg-utils libatspi2.0-0 kde-cli-tools trash-cli libglib2.0-bin gvfs-bin xvfb libsqlcipher-dev v4l-utils xdotool'
fi

sudo /usr/bin/printf "Welcome to Cordatus Installation\n"
sudo apt update && sudo apt install curl -y
ARCH=$(uname -m)
L4T_SUPPORT=("36.2.0 35.5.0 35.4.1 35.3.1 35.2.1 35.1.0 32.7.4 32.7.3 32.7.2 32.7.1")
JP4_LIST=("32.7.4 32.7.3 32.7.2 32.7.1")
JP5_LIST=("35.5.0 35.4.1 35.3.1 35.2.1 35.1.0")
JP6_LIST=("36.2.0")

# Determine if L4T basics are in place and installed correctly
if [ "$ARCH" = "aarch64" ]; then
    ARCH="arm64"
    if [ -f /etc/nv_tegra_release ]; then
        JETSON_RELEASE=$(head -n 1 /etc/nv_tegra_release | cut -f 2 -d ' ' | grep -Po '(?<=R)[^;]+')
        JETSON_REVISION=$(head -n 1 /etc/nv_tegra_release | cut -f 2 -d ',' | grep -Po '(?<=REVISION: )[^;]+')
        if [[ -z "$JETSON_RELEASE" ]] || [[ -z "$JETSON_REVISION" ]]; then
            echo "❌ We couldn't determine your L4T version. Please check if /etc/tegra_release or /etc/nv_tegra_release file does exist."
            exit 1
        fi
        jetpack="$JETSON_RELEASE.$JETSON_REVISION"
    else
        echo "❌ We couldn't locate the /etc/nv_tegra_release file. This file should exist inside the given location in a healthy L4T installation."
        echo "ℹ️  Please re-flash the device and make sure /etc/nv_tegra_release file is in place. Exiting..."
    fi

    echo "ℹ️  Checking if the nvidia-l4t-core package is installed and matching the L4T versions..."
    if [ ! -z "$(dpkg -l | grep nvidia-l4t-core)" ]; then
        echo "✔️  L4T Core package detected."
        JETSON_L4T_STRING=$(dpkg-query --showformat='${Version}' --show nvidia-l4t-core)
        JETSON_L4T_ARRAY=$(echo $JETSON_L4T_STRING | cut -f 1 -d '-')
        # JETSON_L4T_RELEASE=$(echo $JETSON_L4T_ARRAY | cut -f 1 -d '.')
        echo "L4T Release: $JETSON_L4T_ARRAY"
        if os_support_check "$L4T_SUPPORT" " " $JETSON_L4T_ARRAY; then
            echo "✔️ This version of L4T $JETSON_L4T_ARRAY is supported by the Cordatus AI."
            if [ $JETSON_L4T_ARRAY = $jetpack ]; then
                echo "ℹ️  L4T versions from /etc/nv_tegra_release and nvidia-l4t-core matches."
            else
                echo "ℹ️  L4T versions from /etc/nv_tegra_release and nvidia-l4t-core doesn't match. This is not normal but continueing the installation..."
            fi
        else
            echo "❌ $JETSON_L4T_ARRAY is not supported by the Cordatus AI. Exiting the installation..."
            echo "ℹ️  Supported L4T versions: ${L4T_SUPPORT[*]}"
            exit 1
        fi
    else
        echo "❌ L4T Core package (nvidia-l4t-core) is not installed. Please check your JetPack SDK Components installation."
        echo "ℹ️  Exiting..."
        exit 1
    fi
fi

# Global parameters
jetpack="None"
docker_check=false
docker_compose_check=false
nvdocker_check=false
PULL_TEST_CONTAINER="cdcr.dev/cordatushello/hello:linux"
DOCKER_INSTALLED=false

# Jetson variables
cuda_check=false
cudnn_check=false

# Desktop variables
nvgpu_check=false
nvgpu_driver_check=false

# Check docker and nvdocker installation
/usr/bin/printf "🐋 Checking Docker Installation...\n"
DOCKER_CHECK=$(sudo docker pull $PULL_TEST_CONTAINER -q)
if [[ $DOCKER_CHECK = $PULL_TEST_CONTAINER ]]; then
    /usr/bin/printf "🐋 Docker is up and working. Skipping Docker installation...\n"
    DOCKER_INSTALLED=true
    docker_check=true
else
    if [ -x "$(command -v docker)" ]; then
        /usr/bin/printf "✔️  Docker detected.\n"
        docker_check=true
    else
        /usr/bin/printf "❌ Docker is not installed.\n"

        while true; do
        read -p "ℹ️ Docker is required to run Cordatus with all the capabilities and this script will try to install Docker. Do you want to proceed? [y / n]: " yn
            case $yn in
                [Yy]* ) /usr/bin/printf "⚙️ Installing Docker...\n" && curl https://get.docker.com | sh \
                      && sudo systemctl --now enable docker \
                      && /usr/bin/printf "✔️  Docker is installed and enabled...\n" \
                      docker_check=true; break;;
                [Nn]* ) /usr/bin/printf "❌ Exiting Cordatus installation since Docker is required...\n" exit;;
                * ) echo "Please answer yes or no.";;
            esac
        done
    fi
fi

/usr/bin/printf "🐳 Checking nvidia-docker2 Installation...\n"

if [ -x "$(command -v nvidia-container-cli)" ]; then
    echo "✔️  nvidia-docker2 is installed. Skipping NVIDIA-Docker installation...\n"
    nvdocker_check=true
else
    /usr/bin/printf "❌ nvidia-docker2 is not installed.\n"

    while true; do
    read -p "ℹ️  NVIDIA-Docker is required to run GPU accelerated containers with Cordatus and this script will try to install NVIDIA-Docker. Do you want to proceed? [y / n]: " yn
        case $yn in
            [Yy]* ) distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
                    && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
                    && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.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 nvidia-docker2 \
                    && sudo systemctl restart docker \
                    && /usr/bin/printf "✔️  nvidia-docker2 is installed and enabled...\n" \
                    nvdocker_check=true; break;;
            [Nn]* ) /usr/bin/printf "❌ Exiting Cordatus installation since NVIDIA-Docker is required...\n" exit;;
            * ) echo "Please answer yes or no.";;
        esac
    done
fi

/usr/bin/printf "🐋 Checking Docker Compose Installation...\n"

if [ -x "$(command -v docker-compose)" ]; then
    /usr/bin/printf "✔️  docker-compose detected. Skipping Docker Compose installation...\n"
    docker_compose_check=true
else
    /usr/bin/printf "⚙️  Installing docker-compose...\n"
    sudo curl -L https://github.com/docker/compose/releases/download/v2.5.1/docker-compose-`uname -s`-`uname -m` -o /usr/local/bin/docker-compose
    sudo chmod +x /usr/local/bin/docker-compose
    /usr/bin/printf "✔️  docker-compose is installed...\n"
    docker_compose_check=true
fi

if [ -f /etc/nv_tegra_release ]; then
    JETSON_RELEASE=$(head -n 1 /etc/nv_tegra_release | cut -f 2 -d ' ' | grep -Po '(?<=R)[^;]+')
    JETSON_REVISION=$(head -n 1 /etc/nv_tegra_release | cut -f 2 -d ',' | grep -Po '(?<=REVISION: )[^;]+')
    if [[ -z "$JETSON_RELEASE" ]] || [[ -z "$JETSON_REVISION" ]]; then
        /usr/bin/printf "❌ We couldn't determine your L4T version. Please check if /etc/tegra_release or /etc/nv_tegra_release file is exist.\n"
        exit;
    fi
    jetpack="$JETSON_RELEASE.$JETSON_REVISION"

    /usr/bin/printf "Jetpack version: $jetpack\n"
fi

LATEST=$( curl -s -d '{"arch":"'$ARCH'", "jetpack":"'$jetpack'"}' -H "Content-Type: application/json" \
            -X POST https://api.cordatus.ai/latest-version | tr -d '"' )

if [[ ! -f "/home/$USER/.config/autostart/cordatus.desktop" ]]; then
    /usr/bin/printf "Autostart Cordatus [y / n]? "
    read AUTO_START
    AUTO_START="$(echo $AUTO_START | tr '[A-Z]' '[a-z]')"
    if [ "$AUTO_START" = "y" ] || [ "$AUTO_START" = "yes" ]; then
        sudo mkdir -p /home/$USER/.config/autostart
        echo "[Desktop Entry]
        Type=Application
        Version= $LATEST
        Name=cordatus
        Comment=cordatusstartup script
        Exec=/usr/bin/cordatus
        StartupNotify=false
        Terminal=false" | sudo tee -a /home/$USER/.config/autostart/cordatus.desktop
    fi
fi

/usr/bin/printf "Latest version: $LATEST\n"

if [[ $ARCH == "x86_64" ]]; then
    /usr/bin/printf "❔ Checking GPU...\n"

    gpu_model=`sudo lshw -c video | grep 'product'|cut -f2 -d ":"`

    if [ -z "$gpu_model" ]; then
        /usr/bin/printf "No GPU detected.\n"
    else
        /usr/bin/printf "✔️  GPU detected.\n"
        nvgpu_check=true
    fi

    /usr/bin/printf "❔️ Checking GPU driver...\n"

    if [ -d "/proc/driver/nvidia" ]; then
        /usr/bin/printf "✔️  GPU driver path detected.\n"
        nvgpu_driver_check=true
    else
        /usr/bin/printf "❌ GPU driver couldn't be detected in /proc/driver/nvidia.\n"
        /usr/bin/printf "ℹ️  Please follow the steps provided in the Cordatus Documentation from: https://docs.cordatus.ai/installation/#install-nvidia-cuda-driver \n"
    fi

    if "$nvgpu_check" && "$nvgpu_driver_check" && "$docker_check" && "$nvdocker_check" && "$docker_compose_check"; then
        /usr/bin/printf "⚙️  Proceeding with Cordatus installation...\n"
        /usr/bin/printf "⚙️  Installing requirements... \n"
        sleep 0.5
        sudo apt update
        
        for pkg in $pkgs; do
            status="$(dpkg-query -W --showformat='${db:Status-Status}' "$pkg" 2>&1)"
            if [ ! $? = 0 ] || [ ! "$status" = installed ]; then
                /usr/bin/printf "ℹ️  Installing $pkg"
                sudo apt install -y $pkg
            fi
        done
        while true; do
            read -p "🐋 Docker privileges needs to be modified in order to use Cordatus. Do you accept granting root privileges to Docker? [y / n]: " yn
            case $yn in
                [Yy]* ) sudo groupadd docker; sudo usermod -aG docker $USER; break;;
                [Nn]* ) exit;;
                * ) echo "Please answer yes or no.";;
            esac
        done
    else
        /usr/bin/printf "❌ Some dependencies has not been installed. Please check the logs above."
        exit
    fi

    wget https://download.cordatus.ai/cordatus_${LATEST}_x64.deb -O /tmp/cordatus_${LATEST}_x64.deb
    sudo dpkg -i /tmp/cordatus_${LATEST}_x64.deb

    while true; do
        read -p "ℹ️  Your system needs to be rebooted to complete the installation. Do you want to reboot now? [y / n]: " yn
        case $yn in
            [Yy]* ) sudo reboot; break;;
            [Nn]* ) exit;;
            * ) echo "Please answer yes or no.";;
        esac
    done

else
    # CUDA Check
    /usr/bin/printf "❔️ Checking CUDA Installation...\n"

    # CUDA version (path) mapping with respect to L4T version
    if os_support_check "$JP4_LIST" " " $JETSON_L4T_ARRAY; then
        cuda_ver="cuda-10.2"
    elif os_support_check "$JP5_LIST" " " $JETSON_L4T_ARRAY; then
        cuda_ver="cuda-11.4"
    elif os_support_check "$JP6_LIST" " " $JETSON_L4T_ARRAY; then
        cuda_ver="cuda-12.2"
    else
        echo "Couldn't match the L4T Revision with CUDA Version. Exiting..."
        exit 1
    fi

    # Check if CUDA directory exists. If yes, gather CUDA version via nvcc command.
    if [ -d "/usr/local/$cuda_ver" ]; then
        echo "✔️  CUDA directory found at /usr/local/$cuda_ver"
        export PATH=/usr/local/$cuda_ver/bin${PATH:+:${PATH}}
        export LD_LIBRARY_PATH=/usr/local/$cuda_ver/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
        echo "❔️ Checking if CUDA installed correctly..."
        if [[ -x "$(command -v nvcc)" || -f "/usr/local/$cuda_ver/bin/nvcc" ]]; then
            nvcc_check=`/usr/local/$cuda_ver/bin/nvcc --version | grep 'release'|cut -f5 -d " "| tr -d ','`
            echo "✔️  CUDA version $nvcc_check found. Your CUDA installation is verified."
            cuda_check=true
        else
            echo "❌ CUDA version couldn't be verified. Check your JetPack Components installation."
            echo "Exiting..."
            exit 1
        fi
    else
        echo "❌ CUDA directory doesn't exist. Check your JetPack Components installation."
        echo "Exiting..."
        exit 1
    fi

    /usr/bin/printf "❔️ Checking cuDNN Installation...\n"

    cudnnhdr_check=`find /usr -name "cudnn.h"`
    if [ -z "$cudnnhdr_check" ]; then
        echo "❌ cuDNN is not installed."
        cudnn_check=false
    else
        echo "✔️ cuDNN detected in $cudnnhdr_check."
        cudnn_check=true
    fi

    /usr/bin/printf "❔️ Checking TensorRT Installation...\n"

    trtlibs_check=`dpkg -l | grep nvinfer`
    if [ -z "$trtlibs_check" ]; then
        echo "❌ TensorRT is not installed."
        trt_check=false
    else
        echo "✔️  TensorRT detected."
        trt_check=true
        trt_version=`dpkg -l | grep libnvinfer-bin | tr -s ' ' | cut -d ' ' -f3`
    fi

    if "$cuda_check" && "$cudnn_check" && "$trt_check" && "$docker_check" && "$nvdocker_check" && "$docker_compose_check"; then
        /usr/bin/printf "🚀 Your JetPack installation is valid. Proceeding with Cordatus installation...\n"
        /usr/bin/printf "⚙️  Installing required dependencies for Cordatus...\n"
        sudo apt update
        
        for pkg in $pkgs; do
            status="$(dpkg-query -W --showformat='${db:Status-Status}' "$pkg" 2>&1)"
            if [ ! $? = 0 ] || [ ! "$status" = installed ]; then
                /usr/bin/printf "ℹ️ Installing $pkg"
                sudo apt install -y $pkg
            fi
        done
        while true; do
            read -p "🐋 Docker privileges needs to be modified in order to use Cordatus. Do you accept granting root privileges to Docker? [y / n]: " yn
            case $yn in
                [Yy]* ) sudo groupadd docker; sudo usermod -aG docker $USER; break;;
                [Nn]* ) exit;;
                * ) echo "Please answer yes or no.";;
            esac
        done

    else
        /usr/bin/printf "❌ You have missing components. Please install JetPack Components using NVIDIA SDKManager.\n"
        exit;
    fi
    wget https://download.cordatus.ai/cordatus_${LATEST}_arm64.deb -O /tmp/cordatus_${LATEST}_arm64.deb
    sudo dpkg -i /tmp/cordatus_${LATEST}_arm64.deb

    while true; do
        read -p "ℹ️  Your system needs to be rebooted to complete the installation. Do you want to reboot now? [y / n]: " yn
        case $yn in
            [Yy]* ) sudo reboot; break;;
            [Nn]* ) exit;;
            * ) echo "Please answer yes or no.";;
        esac
    done
fi
