#!/bin/bash -eu

update=false
if [[ $# -gt 0 ]]; then
    if [[ $1 != "--update" ]]; then
        echo "Invalid parameter! Use:"
        echo "./create_sys_image"
        echo "or"
        echo "./create_sys_image --update"
        exit 1
    else
        update=true
    fi
fi

if [[ $(basename $(pwd)) == "bin" ]]; then
    cd ..
fi

mkdir -p output

export MPLBACKEND=qt5agg
export JULIA_PKG_SERVER_REGISTRY_PREFERENCE="eager"

if [[ $HOSTNAME == "ufryzen" || $HOSTNAME == "framework" ]]; then
    export NO_MTK=true
fi

if ! command -v juliaup &> /dev/null; then
    echo "Please install the Julia installer 'juliaup'!"
    echo "See: https://github.com/JuliaLang/juliaup"
    exit 1
fi

# if julia is not installed, install Julia 1.11
if ! command -v julia &> /dev/null; then
    juliaup add 1.11
    juliaup default 1.11
fi

julia_version=$(julia --version | awk '{print($3)}')
julia_major=${julia_version:0:3}
if [[ $julia_major == "1.1" ]]; then
    julia_major=${julia_version:0:4} 
fi

# if the current julia version is older than 1.10, exit
if [[ $julia_major == "1.10" ]] || [[ $julia_major == "1.11" ]]; then
    echo OK
else
    echo "The current Julia version is too old, installing Julia 1.11!"
    juliaup add 1.11
    juliaup default 1.11
fi

PYTHON_PATH=$(which python3)
if [ -x "$PYTHON_PATH" ]; then
    echo "Python is found at $PYTHON_PATH"
    if $PYTHON_PATH -c "import matplotlib" &> /dev/null; then
        echo "Matplotlib found. Using existing installation."
        export PYTHON=$PYTHON_PATH
    else
        echo "Matplotlib is not found."
        read -p "Do you want to install matplotlib with Conda? (y/n): " choice
        case "$choice" in 
          y|Y ) 
            export PYTHON=""
            ;;
          n|N ) 
            echo "Exiting without installing matplotlib."
            exit 1
            ;;
          * ) 
            echo "Invalid choice. Exiting."
            exit 1
            ;;
        esac
    fi
else
    echo "Python is not found."
    exit 1
fi

totalmem=$(grep MemTotal /proc/meminfo | awk '{printf "%i", $2 / 1024}')
if [[ $totalmem -lt 27000 ]]; then
    echo "Warning: Less than 27GB of memory detected. Using only one thread for sysimage compilation."
    export JULIA_IMAGE_THREADS=1
fi

if [ -d .git ] || git rev-parse --git-dir > /dev/null 2>&1 ; then
    branch=$(git rev-parse --abbrev-ref HEAD | sed 's/\//-/g')
else
    branch=""
fi
if test -f "kps-image-${julia_major}-${branch}.so"; then
    mv bin/kps-image-${julia_major}-${branch}.so kps-image-${julia_major}-${branch}.so.bak
fi
julia --startup-file=no -e "using Pkg; Pkg.add(\"TestEnv\")"

if [[ $update == true ]]; then
    echo "Updating packages..."
    export JULIA_PKG_SERVER_REGISTRY_PREFERENCE="eager"
    if test -f "Manifest.toml"; then
        mv Manifest.toml Manifest.toml.bak
    fi
    if [[ $julia_major == "1.10" ]]; then
        if test -f "Manifest-v1.10.toml"; then
            mv Manifest-v1.10.toml Manifest-v1.10.toml.bak
        fi
    else
        if test -f "Manifest-v1.11.toml"; then
            mv Manifest-v1.11.toml Manifest-v1.11.toml.bak
        fi
    fi
    julia --startup-file=no --pkgimages=no --project -e "using TestEnv; TestEnv.activate(); using Pkg; Pkg.add(\"PyCall\"); Pkg.build(\"PyCall\")"
    if [[ $PYTHON == "" ]]; then
        julia --startup-file=no --pkgimages=no --project -e "using TestEnv; TestEnv.activate(); using Pkg; Pkg.add(\"Conda\"); using Conda; Conda.add(\"matplotlib\"); using ControlPlots"
    fi
    julia --startup-file=no --pkgimages=no --project -e "using Pkg; Pkg.update()"
    if [[ $julia_major == "1.10" ]]; then
        mv Manifest.toml Manifest-v1.10.toml
    else
        mv Manifest.toml Manifest-v1.11.toml
    fi
else
    if [[ $julia_major == "1.10" ]]; then
        cp Manifest-v1.10.toml.default Manifest-v1.10.toml
        echo "Using Manifest-v1.10.toml.default ..."        
    else
        cp Manifest-v1.11.toml.default Manifest-v1.11.toml
        echo "Using Manifest-v1.11.toml.default ..."
    fi
fi
if [[ $update == true ]]; then
    julia --startup-file=no --pkgimages=no --project -e "include(\"./test/create_sys_image.jl\");"
else
    julia --startup-file=no --pkgimages=no --project -e "using Pkg; Pkg.instantiate();"
    julia --startup-file=no --pkgimages=no --project -e "using TestEnv; TestEnv.activate(); include(\"./test/create_sys_image.jl\");"
fi

if [[ $branch != "" ]]; then
    SOFILE="bin/kps-image-${julia_major}-${branch}.so"
else
    SOFILE="bin/kps-image-${julia_major}.so"
fi
if test -f $SOFILE; then
    mv $SOFILE $SOFILE.bak
fi
mv kps-image_tmp.so $SOFILE
julia --startup-file=no  --project -e "using Pkg; Pkg.precompile(); Pkg.resolve()"
if [ -d src ]; then
    cd src
    touch *.jl # make sure all modules get recompiled in the next step
    cd ..
fi
julia --project -J bin/kps-image-${julia_major}-${branch}.so -e "using KiteUtils, KiteControllers"
julia --project -J bin/kps-image-${julia_major}-${branch}.so -e "using TestEnv; TestEnv.activate(); using ControlPlots, NativeFileDialog"
julia --project -J bin/kps-image-${julia_major}-${branch}.so -e "using Pkg; Pkg.update()"
