#!/usr/bin/env bash

# Copyright (c) 2010-2025, Lawrence Livermore National Security, LLC. Produced
# at the Lawrence Livermore National Laboratory. All Rights reserved. See files
# LICENSE and NOTICE for details. LLNL-CODE-806117.
#
# This file is part of the MFEM library. For more information and source code
# availability visit https://mfem.org.
#
# MFEM is free software; you can redistribute it and/or modify it under the
# terms of the BSD-3 license. We welcome feedback and contributions, see file
# CONTRIBUTING.md for details.

# Initialize modules for users not using bash as a default shell
if test -e /usr/share/lmod/lmod/init/bash
then
  . /usr/share/lmod/lmod/init/bash
fi

set -o errexit
set -o nounset

script_name=$(basename $0)

function usage()
{
  echo "MFEM CI build and test driver"
  echo ""
  echo "Syntax:"
  echo "> ${script_name} --spec \"spack spec\" [--deps-only] [--data]"
  echo "                 [--data-dir=/path/to/mfem/data]"
  echo ""
  echo "> ${script_name} --build-only [--data]"
  echo "                 [--data-dir=/path/to/mfem/data]"
  echo ""
  echo "> ${script_name} --test-only [--data]"
  echo "                 [--data-dir=/path/to/mfem/data]"
  echo ""
  echo "Options:"
  echo " --spec"
  echo "    The most common usage is to pass a spack spec to the script."
  echo "    This spec may or may not specify \"mfem\" as the package, since it"
  echo "    is the default."
  echo ""
  echo " --deps-only"
  echo "    Only install MFEM dependencies. Will clone mfem-uberenv to drive a"
  echo "    local instance on spack with which we install the dependencies."
  echo "    This local spack instance is tuned to be completely independent"
  echo "    from the user environment."
  echo ""
  echo " --build-only"
  echo "    Only build MFEM, based on the existing configuration."
  echo ""
  echo " --test-only"
  echo "    Only test MFEM, based on the existing configuration."
  echo ""
  echo " --data"
  echo "    Make running tests with \"--data\" a requirement: will fail in the"
  echo "    data directory is not present in the parent of the mfem root directory."
  echo "    Note: default behavior is to run data tests if data dir is present."
  echo ""
  echo " --data-dir=/path/to/mfem/data"
  echo "    Path to a clone of the MFEM/data repo: https://github.com/mfem/data"
  echo "    The default path is: '../data'."
  echo ""
}

project_dir="$(pwd)"
hostname="$(hostname)"

spec=""
mode=""
data_dir=""
with_data=false

sys_type=${SYS_TYPE:-""}
threads=${THREADS:-""}
module_list=${MODULE_LIST:-""}
job_unique_id=${CI_JOB_ID:-""}
use_dev_shm=${USE_DEV_SHM:-true}
spack_debug=${SPACK_DEBUG:-false}
debug_mode=${DEBUG_MODE:-false}

# CI_REGISTRY_USER, CI_REGISTRY_IMAGE, and CI_JOB_TOKEN are automatically set by
# Gitlab, see https://docs.gitlab.com/ee/ci/variables/predefined_variables.html.
# The Spack build cache files can be viewed from the Gitlab web interface under
# "Deploy" -> "Container Registry".
#
# Since the CI_<VARIABLES> are only set in CI, we provide sensible defaults for
# usage outside CI context (e.g. to reproduce a CI scenario locally).
#
# REGISTRY_TOKEN when set allows to provide our own personal access token to
# the CI registry. Be sure to set the token with at least read access to the
# registry. It is optional outside CI but will speed up the build if many
# dependencies are needed.
#
registry_token=${REGISTRY_TOKEN:-""}
ci_registry_user=${CI_REGISTRY_USER:-"${USER}"}
ci_registry_image=${CI_REGISTRY_IMAGE:-"czregistry.llnl.gov:5050/mfem/mfem-autotest"}
ci_registry_token=${CI_JOB_TOKEN:-"${registry_token}"}

spec_tab=()

# Options
while [[ $# -gt 0 ]]
do
    key="$1"

    case $key in
        --spec)
            shift
            arg="${1}"
            if [[ "${arg}" =~ ^\" ]]; then
              arg="${arg#\"}"
              while [[ ! "${arg}" =~ [^\\]\"$ ]]; do
                spec_tab+=("${arg}")
                shift
                arg="${1}"
              done
              arg="${arg%\"}"
            fi
            spec_tab+=("${arg}")
            shift
            ;;
        --deps-only|--build-only|--test-only)
            mode="$key"
            shift # past argument
            ;;
        --data)
            with_data=true
            shift # past argument
            ;;
        --data-dir)
            data_dir="$2"
            shift # past argument
            shift # past value
            ;;
        -h|--help)
            usage
            exit 0;
            ;;
        *)  # unknown option
            echo "[Error]: option $key is unknown"
            exit 1;
            ;;
    esac
done

spec="${spec_tab[@]}"

timed_message ()
{
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
    echo "~ $(date --rfc-3339=seconds) ~ ${1}"
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
}

at_exit()
{
  if [[ "${debug_mode}" != "true" ]]; then
    if [[ "${#exit_rm_list[@]}" != 0 ]]; then
      echo ${script_name}:at_exit: rm -rf "${exit_rm_list[@]}"
      rm -rf "${exit_rm_list[@]}"
    fi
    if [[ "${exit_distclean}" == 1 ]]; then
      echo ${script_name}:at_exit: \
        make -C "${project_dir}" distclean "> /dev/null 2>&1"
      make -C "${project_dir}" distclean > /dev/null 2>&1
    fi
  fi
}

if [[ ${debug_mode} == true ]]
then
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
    echo "~~~~~ Debug mode:"
    echo "~~~~~ - Spack debug mode."
    echo "~~~~~ - Deactivated shared memory."
    echo "~~~~~ - Do not push to buildcache."
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
    use_dev_shm=false
    spack_debug=true
fi

if [[ -n ${module_list} ]]
then
    timed_message "Modules to load: ${module_list}"
    module load ${module_list}
fi

prefix=""

if [[ -d /dev/shm && ${use_dev_shm} == true ]]
then
    prefix="/dev/shm/${hostname}"
    if [[ -z ${job_unique_id} ]]; then
      job_unique_id=manual_job_$(date +%s)
      while [[ -d ${prefix}-${job_unique_id} ]] ; do
          sleep 1
          job_unique_id=manual_job_$(date +%s)
      done
    fi

    prefix="${prefix}-${job_unique_id}"
else
    # We set the prefix in the parent directory so that spack dependencies are
    # not installed inside the source tree.
    prefix="${project_dir}/../spack-and-build-root"
fi

echo "Creating directory ${prefix}"
echo "project_dir: ${project_dir}"

mkdir -p ${prefix}

# List of files/directories to remove at exit; more entries are added below.
# This list will be removed even if there's an error. In debug mode, this
# cleanup is not performed.
exit_rm_list=("${prefix}")
exit_distclean=
trap 'at_exit' EXIT HUP INT QUIT ABRT KILL PIPE TERM XCPU

# Prefix tree:
#   <prefix>
#    ├── spack
#    ├── spack_env
#    └── spack-user-cache

spack_cmd="${prefix}/spack/bin/spack"
spack_env_path="${prefix}/spack_env"
command -v python3 > /dev/null && python_cmd=python3 || python_cmd=python
uberenv_cmd="${python_cmd} tests/uberenv/uberenv.py"
if [[ ${spack_debug} == true ]]
then
    spack_cmd="${spack_cmd} --debug --stacktrace"
    uberenv_cmd="${uberenv_cmd} --spack-debug"
fi

# Dependencies
if [[ "${mode}" != "--build-only" && "${mode}" != "--test-only" ]]
then
    timed_message "Building dependencies"

    if [[ -z ${spec} ]]
    then
        echo "[Error]: SPEC is undefined, aborting..."
        exit 1
    fi

    prefix_opt="--prefix=${prefix}"

    echo ${spec} > spec.txt

    # We force Spack to put all generated files (cache and configuration of
    # all sorts) in a unique location so that there can be no collision
    # with existing or concurrent Spack.
    spack_user_cache="${prefix}/spack-user-cache"
    export SPACK_DISABLE_LOCAL_CONFIG=""
    export SPACK_USER_CACHE_PATH="${spack_user_cache}"
    mkdir -p ${spack_user_cache}

    echo "Fetching uberenv ..."
    # FIXME: the messages from the next clone appear out-of-place, at the end of
    #        the logs on Gitlab. (On Lassen the messages are where expected.)
    tests/gitlab/get_mfem_uberenv || { echo "Error fetching Uberenv"; exit 1; }

    # Remove "${project_dir}/tests/uberenv" at exit (even if there's an error)
    exit_rm_list=("${exit_rm_list[@]}" "${project_dir}/tests/uberenv")

    echo "Removing existing configuration"
    make distclean

    # Add -O2 flag for debug builds to speed up the tests.
    sed -e 's/^\(DEBUG_FLAGS =.*\)$/\1 -O2/g' config/defaults.mk > \
        config/defaults.mk.new
    mv -f config/defaults.mk.new config/defaults.mk

    # generate cmake cache file with uberenv and radiuss spack package
    timed_message "Spack setup and environment"
    ${uberenv_cmd} --setup-and-env-only --spec="${spec}" ${prefix_opt}

    if [[ -n ${ci_registry_token} ]]
    then
        timed_message "GitLab registry as Spack Buildcache"
        ${spack_cmd} -D ${spack_env_path} mirror add --unsigned \
          --oci-username ${ci_registry_user} \
          --oci-password ${ci_registry_token} \
          gitlab_ci oci://${ci_registry_image}
    fi

    timed_message "Spack build of dependencies"
    ${uberenv_cmd} --skip-setup-and-env --spec="${spec}" ${prefix_opt} || \
    {
      # Save Spack logs
      cd ${SPACK_USER_CACHE_PATH}
      log_archive="${project_dir}/spack-logs-${hostname}-${CI_JOB_ID}.tgz"
      tar zvcf "${log_archive}" stage/*/*/spack*.txt
      timed_message "Spack build failed! See log archive: ${log_archive}"
      exit 1
    }

    if [[ -n ${ci_registry_token} && ${debug_mode} == false ]]
    then
        timed_message "Push dependencies to buildcache"
        ${spack_cmd} -D ${spack_env_path} buildcache push \
          --only dependencies gitlab_ci
    fi

    timed_message "Dependencies built"

    # Make sure that a configuration was generated by spack (part 1).
    cp config/config.mk config/spack-config.mk
    cp config/_config.hpp config/spack_config.hpp
fi

# Configuration
if [[ "${mode}" != "--deps-only" ]]
then
# Host config file
# Make sure that a configuration was generated by spack (part 2).
# Without this running the rest of the script would generate a new config.
    if [[ -f "config/spack-config.mk" && -f "config/spack_config.hpp" ]]
    then
        cp config/spack-config.mk config/config.mk
        cp config/spack_config.hpp config/_config.hpp
    else
        echo "[Error]: No result for at least one of"
        echo "[Error]:     ${project_dir}/config/spack-config.mk"
        echo "[Error]:     ${project_dir}/config/spack_config.hpp"
        echo "[Error]: Spack generated configuration not found."
        exit 1
    fi

    # Setup the MFEM/data repository directory
    # Some additional unit tests are enabled when '../data' is present
    if [[ -z "${data_dir}" ]]
    then
        # By default, data_dir is ../data.
        data_dir="../data"
    else
        # data_dir is specified, so we need to link its content into the
        # project parent dir.
        if [[ -e "../data" ]]; then
            if [[ -L "../data" ]]; then
                echo "[Information]: '../data' link already exists. Deleting."
                rm "../data"
            else
                echo "[Error]: '../data' already exists and it's NOT a link"
                exit 1
            fi
        fi
        ln -sf "${data_dir}" "../data"
    fi
    # The PUMI examples expect the PUMI datafiles to be in 'data/pumi'
    if [[ -d "../data/pumi" ]]; then
        ln -sf "../../data/pumi" "data"
    fi

    if [[ "$with_data" == "true" && ! -d "../data" ]]
    then
        echo "[Error]: '$data_dir' is not a directory while asking for --data"
        exit 1
    fi
fi

# Build (also build when asked for testing, to make sure we test what we want)
if [[ "${mode}" != "--deps-only" ]]
then
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
    echo "~ Project Dir: ${project_dir}"
    echo "~ Data Dir: ${data_dir}"
    echo "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
    echo ""

    timed_message "MFEM Configuration"

    make info

    timed_message "Building MFEM"

    # If the build is terminated, e.g. job timeout, run 'make distclean'
    exit_distclean=1

    make all -j ${threads}

    exit_distclean=
fi

# Test
if [[ "${mode}" != "--deps-only" && "${mode}" != "--build-only" ]]
then
    timed_message "Testing MFEM"

    # If the build is terminated, e.g. job timeout, run 'make distclean'
    exit_distclean=1

    make test
    test_status=$?

    timed_message "Cleaning MFEM"

    if make distclean > distclean.log 2>&1; then
      echo "[Information]: make distclean OK"
    else
      cat distclean.log
      echo
      echo "[Information]: make distclean FAILED (see above)"
    fi
    rm -f distclean.log

    exit_distclean=

    exit $test_status
fi