#!/usr/bin/env bash

# (C) Copyright 2018- ECMWF.
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

set -euo pipefail

hpc2020_java_version=11.0.6
hpc2020_python_version=3.11.10-01
hpc2020_cmake_version=3.28.3
hpc2020_meson_version=1.2.1
hpc2020_ninja_version=1.11.1

# Determine base path for loki
# Either take the root of the current git tree or, if not inside a git repository, then
# use the path of this install script
if [ $(git rev-parse --git-dir > /dev/null 2>&1) ]; then
  loki_path=$(git rev-parse --show-toplevel)
else
  loki_path=$(realpath $(dirname "$0"))
fi

# Configuration default values
verbose=false
is_hpc2020=false
venv_path=
with_jdk=false
with_omni=false
with_docs=false
with_dace=true
with_tests=true
with_examples=true

# Helper functions
print_usage() {
  echo "Usage: $0 [-v] [--hpc2020] [--use-venv[=]<path>] [--with-*] [...]" >&2
}

print_usage_with_options() {
  echo "Loki install script. This installs Loki and selected dependencies."
  echo
  print_usage
  echo
  echo "Available options:"
  echo "  -h / --help                  Display this help message"
  echo "  -v                           Enable verbose output"
  echo "  --hpc2020                    Load HPC2020 (Atos) specific modules and settings"
  echo "  --use-venv[=]<path>          Use existing virtual environment at <path>"
  echo "  --with[out]-jdk              Install JDK instead of using system version (default: use system version)"
  echo "  --with[out]-omni             Install OMNI Compiler (default: disabled)"
  echo "  --with[out]-dace             Install DaCe (default: enabled)"
  echo "  --with[out]-tests            Install dependencies to run tests (default: enabled)"
  echo "  --with[out]-docs             Install dependencies to generate documentation (default: disabled)"
  echo "  --with[out]-examples         Install dependencies to run the example notebooks (default: enabled)"
}

print_step() {
  echo "------------------------------------------------------"
  echo "  $1"
  echo "------------------------------------------------------"
}

# Parse arguments
# (see https://stackoverflow.com/a/7680682)
optspec=":hv-:"
while getopts "$optspec" optchar; do
  case "${optchar}" in
    -)
      case "${OPTARG}" in
        hpc2020)     # Load ECMWF HPC2020 (Atos) specific modules and settings
          is_hpc2020=true
          ;;
        use-venv)    # Specify existing virtual environment
          venv_path=$(realpath "${!OPTIND}")
          OPTIND=$(( OPTIND + 1 ))
          ;;
        use-venv=*)  # Specify existing virtual environment
          venv_path=$(realpath "${OPTARG#*=}")
          ;;
        with-jdk)    # Enable installation of Java
          with_jdk=true
          ;;
        without-jdk) # Disable installation of Java
          with_jdk=false
          ;;
        with-omni)   # Enable installation of OMNI
          with_omni=true
          ;;
        without-omni) # Disable installation of OMNI
          with_omni=false
          ;;
        with-dace)    # Enable installation of DaCe
          with_dace=true
          ;;
        without-dace) # Disable installation of DaCe
          with_dace=false
          ;;
        with-tests)    # Enable installation of dependencies for running tests
          with_tests=true
          ;;
        without-tests) # Disable installation of dependencies for running tests
          with_tests=false
          ;;
        with-docs)    # Enable installation of dependencies for docs generation
          with_docs=true
          ;;
        without-docs) # Disable installation of dependencies for docs generation
          with_docs=false
          ;;
        with-examples)    # Enable installation of dependencies for notebook examples
          with_examples=true
          ;;
        without-examples) # Disable installation of dependencies for notebook examples
          with_examples=false
          ;;
        help)
          print_usage_with_options
          exit 2
          ;;
        *)
          echo "Unknown option '--${OPTARG}'." >&2
          print_usage
          echo "Try '$0 -h' for more options."
          exit 1
          ;;
      esac
      ;;
    h)
      print_usage_with_options
      exit 2
      ;;
    v)
      verbose=true
      ;;
    *)
      echo "Unknown option '-${OPTARG}'." >&2
      print_usage
      echo "Try '$0 -h' for more options."
      exit 1
      ;;
  esac
done


# Print configuration
if [ "$verbose" == true ]; then
  print_step "Installation configuration:"

  [[ "$is_hpc2020" == true ]]  && echo "    --hpc2020"
  [[ "$venv_path" != "" ]]   && echo "    --use-venv='$venv_path'"
  [[ "$with_jdk" == true ]]  && echo "    --with-jdk"
  [[ "$with_omni" == true ]] && echo "    --with-omni"
  [[ "$with_dace" == false ]] && echo "    --without-dace"
  [[ "$with_tests" == false ]] && echo "    --without-tests"
  [[ "$with_docs" == false ]] && echo "    --without-docs"
  [[ "$with_examples" == false ]] && echo "    --without-examples"
fi

# Load modules
if [ "$is_hpc2020" == true ]; then
  print_step "Loading HPC2020 modules and settings"

  module unload cmake
  module load cmake/${hpc2020_cmake_version}

  module unload meson
  module load meson/${hpc2020_meson_version}

  module unload ninja
  module load ninja/${hpc2020_ninja_version}

  if [ "$with_jdk" == false ]; then
    module unload java
    module load java/${hpc2020_java_version}
  fi

  if [ "$venv_path" == "" ]; then
    module unload python3
    module load python3/${hpc2020_python_version}
  fi

fi

#
# Create Python virtualenv
#

if [ "$venv_path" == "" ]; then
  print_step "Creating virtualenv"
  venv_path=${loki_path}/loki_env
  for activate_file in activate activate.csh activate.fish Activate.ps1; do
    if [ -f "${loki_path}/loki_env/bin/${activate_file}" ]; then
      chmod u+w "${loki_path}/loki_env/bin/${activate_file}"
    fi
  done
  python3 -m venv "${venv_path}"
fi

#
# Activate Python virtualenv
#

print_step "Activating virtualenv"
source "${venv_path}/bin/activate"

#
# Install Loki with Python dependencies
#

print_step "Installing Loki and Python dependencies"

cd "$loki_path"

pip_opts=()
[[ "$with_tests" == true ]] && pip_opts+=(tests)
[[ "$with_dace" == true ]] && pip_opts+=(dace)
[[ "$with_docs" == true ]] && pip_opts+=(docs)
[[ "$with_examples" == true ]] && pip_opts+=(examples)
pip_opts=$(printf ",%s" "${pip_opts[@]}")

if [ "$pip_opts" == "," ]; then
  pip_opts=
else
  pip_opts=[${pip_opts:1}]
fi

# Supply pretend version if not a git worktree
if [ ! -e .git ]; then
  export "SETUPTOOLS_SCM_PRETEND_VERSION=$(cat VERSION)"
fi


pip install --upgrade pip
pip install -e .$pip_opts  # Installs Loki dev copy in editable mode
pip install -e ./lint_rules

#
# Install Java
#

if [ "$with_jdk" == true ]; then
  print_step "Downloading and installing JDK"

  JDK_ARCHIVE=openjdk-11.0.2_linux-x64_bin.tar.gz
  JDK_URL=https://download.java.net/java/GA/jdk11/9/GPL/${JDK_ARCHIVE}
  JAVA_INSTALL_DIR=${venv_path}/opt/java
  export JAVA_HOME=${JAVA_INSTALL_DIR}/jdk-11.0.2

  mkdir -p "${JAVA_INSTALL_DIR}"
  rm -rf "${JAVA_HOME}" "${JAVA_INSTALL_DIR}/${JDK_ARCHIVE}"
  cd "${JAVA_INSTALL_DIR}"
  wget -O "${JDK_ARCHIVE}" "${JDK_URL}"
  tar -xzf "${JDK_ARCHIVE}"
fi

#
# Install OMNI
#

if [ "$with_omni" == true ]; then
  print_step "Downloading and installing OMNI Compiler"

  OMNI_INSTALL_DIR=${venv_path}/opt/omni
  mkdir -p "${OMNI_INSTALL_DIR}"
  cd "${OMNI_INSTALL_DIR}"
  rm -rf xcodeml-tools
  git clone --recursive --single-branch https://github.com/omni-compiler/xcodeml-tools.git xcodeml-tools

  cd xcodeml-tools

  omni_opts=()
  [[ ! -z "${JAVA_HOME}" ]] && omni_opts+=("JAVA_HOME=${JAVA_HOME}")

  # A CMake install would be cleaner but they inject without good reason a -Werror for
  # GNU and Clang without writing good enough code that actually avoids any warnings...
  # cmake -S . -B build -DCMAKE_INSTALL_PREFIX="${OMNI_INSTALL_DIR}" "${omni_opts[@]}"
  # cmake --build build
  # cmake --install build
  ./configure --prefix="${OMNI_INSTALL_DIR}" "${omni_opts[@]}"
  make && make install
fi

#
# Writing loki-activate script
#

print_step "Writing loki-activate script"

path_var=\${PATH}

echo "
# This script activates Loki's virtual environment, loads additional modules and sets dependend paths.
#
# Run as 'source loki-activate'

# Load virtualenv
. ${venv_path}/bin/activate
" > "${loki_path}/loki-activate"

# Load ECMWF modules, if required
if [ "${is_hpc2020}" == true ]; then
  if [ "$with_jdk" == false ]; then
    echo "
module unload java
module load java/${hpc2020_java_version}
" >> "${loki_path}/loki-activate"
  fi

  echo "
module unload cmake
module load cmake/${hpc2020_cmake_version}
" >> "${loki_path}/loki-activate"

  echo "
module unload meson
module load meson/${hpc2020_meson_version}
" >> "${loki_path}/loki-activate"

  echo "
module unload ninja
module load ninja/${hpc2020_ninja_version}
" >> "${loki_path}/loki-activate"
fi

# Inject self-installed JDK into env
if [ "$with_jdk" == true ]; then
  echo "
export JAVA_HOME=\"\${JAVA_HOME}\"
" >> "${loki_path}/loki-activate"
  path_var=${JAVA_HOME}/bin:$path_var
fi

# Inject OMNI into env
if [ "$with_omni" == true ]; then
  path_var=${OMNI_INSTALL_DIR}/bin:$path_var
fi

# Update path variable
echo "
export PATH=\"$path_var\"

echo \"Activated loki environment. Unload with 'deactivate'.\"
" >> "${loki_path}/loki-activate"

#
# Finish
#

print_step "Installation finished"
echo
echo "Activate the Loki environment with"
echo
echo "    source loki-activate"
echo

if [ "$with_tests" == true ]; then
  echo "You can test the Loki installation by running"
  echo
  echo "    pytest --pyargs loki lint_rules"
  echo
fi