#!/bin/bash

# ===========================================================================================================
# title         : create_JupyterLab_kernel.sh
# usage         : ./create_JupyterLab_kernel.sh $INSTALLATION_ROOT $EBRAINS_SPACK_ENV $RELEASE_NAME $LAB_KERNEL_ROOT
# description   : creates a spackified JupyterLab kernel conf and places it to the NFS where it can be
#                 loaded by all users.
# ===========================================================================================================

set -euo pipefail

INSTALLATION_ROOT=$1
EBRAINS_SPACK_ENV=$2
RELEASE_NAME=$3
LAB_KERNEL_ROOT=$4

# kernel name is lowercase release name
KERNEL_NAME=$(echo "$RELEASE_NAME" | tr '[:upper:]' '[:lower:]')
# kernel dir is a directory inside LAB_KERNEL_ROOT named after the kernel
KERNEL_PATH=$LAB_KERNEL_ROOT/$KERNEL_NAME

# prepare the base env.sh script required for the JupyterLab kernels:
mkdir -p $KERNEL_PATH/bin

# (1) create env.sh bash script
cat <<EOF > $KERNEL_PATH/bin/env.sh
#!/usr/bin/env bash
set -euxo pipefail
EOF

# (2) append the necessary env variables for spack env and tools
cat $INSTALLATION_ROOT/spack/var/spack/environments/$EBRAINS_SPACK_ENV/load_env.sh >> $KERNEL_PATH/bin/env.sh
# also copy the script so that it can be used independently:
cp $INSTALLATION_ROOT/spack/var/spack/environments/$EBRAINS_SPACK_ENV/load_env.sh $KERNEL_PATH/bin/

# (3) also add the user's .local locations to allow package installation at runtime,
# and the location of python modules installed in the base docker Collab image
cat <<EOF >> $KERNEL_PATH/bin/env.sh
export PATH=\$PATH:/opt/app-root/src/.local/bin
export PYTHONPATH=\$PYTHONPATH:/opt/conda/lib/python3.11/site-packages
export R_LIBS_USER=/opt/app-root/src/.local/lib/R/site-library
mkdir -p \$R_LIBS_USER
export R_LIBS=\$R_LIBS_USER:\$R_LIBS
EOF

# ... and set the SYSTEMNAME env var, to use Spack commands inside the kernel
cat <<EOF >> $KERNEL_PATH/bin/env.sh
export SYSTEMNAME=ebrainslab
EOF

# (4) create startup script for python kernel
cat $KERNEL_PATH/bin/env.sh > $KERNEL_PATH/bin/env_python.sh
echo "python -m ipykernel_launcher -f \$@" >> $KERNEL_PATH/bin/env_python.sh
chmod +x $KERNEL_PATH/bin/env_python.sh

# ... and the new kernel's configuration file
mkdir -p $KERNEL_PATH/$KERNEL_NAME
cat <<EOF > $KERNEL_PATH/$KERNEL_NAME/kernel.json
{
 "argv": ["$KERNEL_PATH/bin/env_python.sh", "{connection_file}", "--profile=default"],
 "display_name": "$RELEASE_NAME",
 "language": "python",
 "env": { "LAB_KERNEL_NAME": "$RELEASE_NAME", "LAB_KERNEL_RELEASE_DATE": "$(date +"%Y-%m-%d")" }
}
EOF

# ... and add EBRAINS logo to kernel
cp $(dirname ${BASH_SOURCE[0]})/etc/logo-64x64-python.png $KERNEL_PATH/$KERNEL_NAME/logo-64x64.png

# (5) create startup script for R kernel
cat $KERNEL_PATH/bin/env.sh > $KERNEL_PATH/bin/env_r.sh
echo "R --slave -e \"IRkernel::main()\" --args \$@" >> $KERNEL_PATH/bin/env_r.sh
chmod +x $KERNEL_PATH/bin/env_r.sh

# ... and the new R kernel's configuration file
mkdir -p $KERNEL_PATH/${KERNEL_NAME}-R
cat <<EOF > $KERNEL_PATH/${KERNEL_NAME}-R/kernel.json
{
 "argv": ["$KERNEL_PATH/bin/env_r.sh", "{connection_file}"],
 "display_name": "R-$RELEASE_NAME",
 "language": "R",
 "env": { "LAB_KERNEL_NAME": "R-$RELEASE_NAME", "LAB_KERNEL_RELEASE_DATE": "$(date +"%Y-%m-%d")" }
}
EOF

# ... and add EBRAINS logo and kernel.js to kernel
cp $(dirname ${BASH_SOURCE[0]})/etc/logo-64x64-R.png $KERNEL_PATH/${KERNEL_NAME}-R/logo-64x64.png
cp $(dirname ${BASH_SOURCE[0]})/etc/kernel.js $KERNEL_PATH/${KERNEL_NAME}-R/