#!/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

# Configuration of package versions:
loki_ant_version=1.10.13

hpc2020_java_version=11.0.6
hpc2020_python_version=3.10.10-01
hpc2020_cmake_version=3.25.2

# 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
proxy_certificate=""
jdk_certificate=""
venv_path=
with_jdk=false
with_ant=false
with_claw=false
claw_install_env=""
with_ofp=false
with_docs=false
with_dace=true
with_tests=true
with_examples=true

# Helper functions
print_usage() {
  echo "Usage: $0 [-v] [--hpc2020] [--*-certificate[=]<path>] [--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 "  --proxy-certificate[=]<path> Provide https proxy certificate, disable cert verification for JDK/ant/OFP downloads"
  echo "  --jdk-certificate[=]<path>   Provide trusted certificate for JDK runtime"
  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]-ant              Install ant instead of using system version (default: use system version)"
  echo "  --with[out]-claw             Install CLAW and OMNI Compiler (default: disabled)"
  echo "  --with[out]-ofp              Install Open Fortran Parser (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)"
  echo "  --claw-install-env[=]<...>   Specify environmental variables for CLAW build and install"
}

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
        proxy-certificate)
          proxy_certificate=$(realpath "${!OPTIND}")
          OPTIND=$(( OPTIND + 1 ))
          ;;
        proxy-certificate=*)
          proxy_certificate=$(realpath "${!OPTARG#*=}")
          ;;
        jdk-certificate)
          jdk_certificate=$(realpath "${!OPTIND}")
          OPTIND=$(( OPTIND + 1 ))
          ;;
        jdk-certificate=*)
          jdk_certificate=$(realpath "${!OPTARG#*=}")
          ;;
        claw-install-env)    # Provide custom environment variables for CLAW build
          claw_install_env="${!OPTIND}"
          OPTIND=$(( OPTIND + 1 ))
          ;;
        claw-install-env=*)  # Provide custom environment variables for CLAW build
          claw_install_env="${!OPTARG#*=}"
          ;;
        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-ant)    # Enable installation of ant
          with_ant=true
          ;;
        without-ant) # Disable installation of ant
          with_ant=false
          ;;
        with-ofp)    # Enable installation of OFP
          with_ofp=true
          ;;
        without-ofp) # Disable installation of OFP
          with_ofp=false
          ;;
        with-claw)   # Enable installation of OMNI+CLAW
          with_claw=true
          ;;
        without-claw) # Disable installation of OMNI+CLAW
          with_claw=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_ant" == true ]]  && echo "    --with-ant"
  [[ "$with_claw" == true ]] && echo "    --with-claw"
  [[ "$with_ofp" == false ]] && echo "    --without-ofp"
  [[ "$with_dace" == false ]] && echo "    --without-dace"
  [[ "$with_tests" == false ]] && echo "    --without-tests"
  [[ "$with_docs" == false ]] && echo "    --without-docs"
  [[ "$with_examples" == false ]] && echo "    --without-examples"
  [[ "$proxy_certificate" != "" ]] && echo "    --proxy-certificate='$proxy_certificate'"
  [[ "$jdk_certificate" != "" ]]   && echo "    --jdk-certificate='$jdk_certificate'"
  [[ "$claw_install_env" != "" ]]  && echo "    --claw_install_env='$claw_install_env'"
fi

if [ "x$proxy_certificate" != "x" ]; then
  PIPPROXYOPTIONS="--cert  $proxy_certificate"
  WGETPROXYOPTIONS="--no-check-certificate"
else
  PIPPROXYOPTIONS=""
  WGETPROXYOPTIONS=""
fi

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

  module unload cmake
  module load cmake/${hpc2020_cmake_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"

if [ "x$proxy_certificate" != "x" ]
then
  export SSL_CERT_FILE=$proxy_certificate
fi

#
# Install Loki with Python dependencies
#

print_step "Installing Loki and Python dependencies"

cd "$loki_path"

pip_opts=()
[[ "$with_tests" == true ]] && pip_opts+=(tests)
[[ "$with_ofp" == true ]] && pip_opts+=(ofp)
[[ "$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

pip install $PIPPROXYOPTIONS --upgrade pip
pip install $PIPPROXYOPTIONS -e .$pip_opts  # Installs Loki dev copy in editable mode
pip install $PIPPROXYOPTIONS -e ./transformations
pip install $PIPPROXYOPTIONS -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 $WGETPROXYOPTIONS -O "${JDK_ARCHIVE}" "${JDK_URL}"
  tar -xzf "${JDK_ARCHIVE}"

  if [ "x$jdk_certificate" != "x" ]; then
    f1=$jdk_certificate
    f2=${JAVA_HOME}/lib/security/cacerts
    mv "$f2" "$f2.bak"
    cp "$f1" "$f2"
  fi
fi

#
# Install ant
#

if [ "$with_ant" == true ]; then
  print_step "Downloading and installing ANT"

  ANT_INSTALL_DIR=${venv_path}/opt/ant
  export ANT_HOME=${ANT_INSTALL_DIR}/apache-ant-${loki_ant_version}

  # Cache NetRexx if it doesn't exist (Download fails from time to time)
  NETREXX_TEMP=${HOME}/.ant/tempcache/NetRexx.zip
  mkdir -p "${HOME}/.ant/tempcache"
  if [[ $(shasum -a 1 "${NETREXX_TEMP}" | cut -d ' ' -f1) != "1a47bf7b5d0055d4a94befc999c593d15b66c119" ]]
  then
    for mirror in https://public.dhe.ibm.com ftp://ftp.software.ibm.com
    do
      wget --tries 3 $WGETPROXYOPTIONS -O "${NETREXX_TEMP}" $mirror/software/awdtools/netrexx/NetRexx.zip
      if [ $? -eq 0 ]
      then
        break
      fi
    done
  fi

  # Download, unpack and install ant
  ANT_ARCHIVE=apache-ant-${loki_ant_version}-bin.tar.gz
  mkdir -p "${ANT_INSTALL_DIR}"
  rm -rf "${ANT_HOME}" "${ANT_INSTALL_DIR}/${ANT_ARCHIVE}"
  cd "${ANT_INSTALL_DIR}"

  set +e
  for mirror in http://ftp.fau.de/apache http://mirror.ox.ac.uk/sites/rsync.apache.org http://archive.apache.org/dist
  do
    wget --tries 3 $WGETPROXYOPTIONS -O "${ANT_ARCHIVE}" $mirror/ant/binaries/${ANT_ARCHIVE}
    if [ $? -eq 0 ]
    then
      break
    fi
  done
  set -e

  tar -xzf "${ANT_ARCHIVE}"

  export PATH="${ANT_HOME}/bin:${PATH}"
  ant -f "${ANT_HOME}/fetch.xml" -Ddest=optional
fi

#
# Install CLAW+OMNI
#

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

  CLAW_INSTALL_DIR=${venv_path}/opt/claw
  mkdir -p "${CLAW_INSTALL_DIR}"
  cd "${CLAW_INSTALL_DIR}"
  rm -rf claw-compiler
  # git clone --recursive https://github.com/claw-project/claw-compiler.git claw-compiler

  # Note that our current regression tests (CLOUDSC) will fail due to faulty offload
  # directives with current CLAW-master. The fixes exist in this pinned branch:
  git clone --recursive --single-branch --branch=mlange-dev https://github.com/mlange05/claw-compiler.git claw-compiler
  cd claw-compiler

  (
    set -a
    eval ${claw_install_env}
    cmake -DCMAKE_INSTALL_PREFIX="${CLAW_INSTALL_DIR}" .
  )

  make
  make install
fi

#
# Install OFP
#

OFP_HOME=${venv_path}/src/open-fortran-parser

if [ "$with_ofp" == true ]; then
  print_step "Installing OFP"

  # Reinstall OFP editable to enable below hack
  pip install $PIPPROXYOPTIONS -e git+https://github.com/mlange05/open-fortran-parser-xml@mlange05-dev#egg=open-fortran-parser

  # HACK: Force OFP version
  echo "VERSION = '0.5.3'" > "${OFP_HOME}/open_fortran_parser/_version.py"

  if [ "x$proxy_certificate" != "x" ]; then
  # Bypass certificate verification
  perl -e '

    use FileHandle;
    my $py = shift;

    my $code = do { local $/ = undef; my $fh = "FileHandle"->new ("<$py"); <$fh> };

    "FileHandle"->new (">$py")->print (<< "EOF");
  import ssl

  ssl._create_default_https_context = ssl._create_unverified_context

  $code
  EOF

  ' "${OFP_HOME}/open_fortran_parser/__main__.py"
  fi

  # Install Java dependencies
  python3 -m open_fortran_parser --deps

  # Copy downloaded binaries to lib
  cd "${OFP_HOME}"
  mkdir -p lib
  cp open_fortran_parser/*.jar lib

  # Rebuild OFP binaries to include custom changes
  ant
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"
fi

# Setup Java for OFP
if [ "$with_ofp" == true ]; then
  echo "
export _OLD_CLASSPATH=\"\${CLASSPATH}\"
export CLASSPATH=\"${OFP_HOME}/open_fortran_parser/*:\${CLASSPATH}\"
" >> "${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 CLAW into env
if [ "$with_claw" == true ]; then
  path_var=${CLAW_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
echo "You can test the Loki installation by running 'py.test transformations lint_rules .'"
echo
