#!/bin/bash

# =========================================================================================================================================
# title         : install_spack_env.sh
# usage         : ./install_spack_env.sh $SPACK_JOBS $INSTALLATION_ROOT $SPACK_VERSION $EBRAINS_REPO $EBRAINS_SPACK_ENV $UPSTREAM_INSTANCE
# description   : installs or updates the spack environment defined in the EBRAINS spack repo
#                 (if the specified spack instance doesn't exist, it also creates it)
# =========================================================================================================================================

SPACK_JOBS=$1         # number of jobs
INSTALLATION_ROOT=$2  # where to set up the installation
SPACK_VERSION=$3      # which spack version to use
EBRAINS_REPO=$4       # location of ebrains-spack-builds repository
EBRAINS_SPACK_ENV=$5  # name of EBRAINS Spack environment to be created/updated
UPSTREAM_INSTANCE=$6  # path to Spack instance to use as upstream (optional)

SPACK_REPO=https://gitlab.ebrains.eu/ri/tech-hub/platform/esd/spack.git
SPACK_VERSION_EBRAINS=${SPACK_VERSION}_ebrains24.04

# specify location of .spack dir (by default in ~)
# this is where cache and configuration settings are stored
export SPACK_USER_CACHE_PATH=$INSTALLATION_ROOT/spack/.spack
export SPACK_USER_CONFIG_PATH=$INSTALLATION_ROOT/spack/.spack

# define SYSTEMNAME variable in sites where it's not already defined
export SYSTEMNAME=${SYSTEMNAME:-${HPC_SYSTEM:-$BSC_MACHINE}}

# initial setup: clone spack if spack dir doesn't already exist
if [ ! -d $INSTALLATION_ROOT/spack ]
then
  git clone --depth 1 -c advice.detachedHead=false -c feature.manyFiles=true --branch $SPACK_VERSION_EBRAINS $SPACK_REPO $INSTALLATION_ROOT/spack
  # SPACK PATCH: the post-build logs on install-time-test-logs.txt gets ovewritten by the post-install logs.
  # quick fix for that: (TODO: investigate more and open PR)
  sed -i "s/self.file_like, \"w\"/self.file_like, \"a\"/g" $INSTALLATION_ROOT/spack/lib/spack/llnl/util/tty/log.py
fi

if [[ $UPSTREAM_INSTANCE ]]
then
  cat <<EOF > $INSTALLATION_ROOT/spack/etc/spack/defaults/upstreams.yaml
upstreams:
  upstream-spack-instance:
    install_tree: $UPSTREAM_INSTANCE/spack/opt/spack
EOF
fi

# activate Spack
source $INSTALLATION_ROOT/spack/share/spack/setup-env.sh

# add repo if it does not exist
if [[ ! $(spack repo list | grep ebrains-spack-builds$) ]]
then
  spack repo add $EBRAINS_REPO
fi

# install platform compiler (extract version from packages.yaml)
if [ $SYSTEMNAME == ebrainslab ]
then
  EBRAINS_SPACK_COMPILER=$(grep 'compiler' $EBRAINS_REPO/site-config/$SYSTEMNAME/packages.yaml | awk -F'[][]' '{ print $2 }')
  spack compiler find
  spack load $EBRAINS_SPACK_COMPILER || { spack install $EBRAINS_SPACK_COMPILER; spack load $EBRAINS_SPACK_COMPILER; }
fi

spack compiler find

# create environment if it does not exist
if [ ! -d "$SPACK_ROOT/var/spack/environments/$EBRAINS_SPACK_ENV" ]
then
  spack env create $EBRAINS_SPACK_ENV
fi

# update environment site-configs
rm -rf $SPACK_ROOT/var/spack/environments/$EBRAINS_SPACK_ENV/site-config && cp -r $EBRAINS_REPO/site-config $SPACK_ROOT/var/spack/environments/$EBRAINS_SPACK_ENV
# update spack.yaml: merge top-level and site-specific spack.yaml files
spack-python $EBRAINS_REPO/site-config/ymerge.py $EBRAINS_REPO/spack.yaml $EBRAINS_REPO/site-config/$SYSTEMNAME/spack.yaml > /tmp/spack.yaml
cp /tmp/spack.yaml $SPACK_ROOT/var/spack/environments/$EBRAINS_SPACK_ENV/

# activate environment
spack env activate --without-view $EBRAINS_SPACK_ENV

# fetch all sources
spack concretize --force --fresh --test root
spack-python -c "exit(not len(spack.environment.active_environment().uninstalled_specs()))" && spack fetch --dependencies --missing

# install the environment, use 2 jobs to reduce the amount of required RAM
spack install -y -j$SPACK_JOBS --fresh --test root

# rebuild spack's database
spack reindex

# create load script that when sourced activates and loads the installed spack environment, using views
# this needs deactivating the environment first:
spack env deactivate
unset SPACK_LD_LIBRARY_PATH
spack env activate --sh $EBRAINS_SPACK_ENV > $SPACK_ROOT/var/spack/environments/$EBRAINS_SPACK_ENV/load_env.sh

# create modules files with spack
# spack module tcl refresh -y
# create loads script that when sourced activates and loads the installed spack environment, using modules
# spack env loads -r
