chbh

#!/usr/bin/env bash

# Ensures the script crashes rather than creates a bad environment
set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$PWD"
VENV_DIR="$PROJECT_DIR/.venv"

TEMPLATE_DIR="$HOME/.local/share/chbh"
TEMPLATE_PYPROJECT="$TEMPLATE_DIR/pyproject.toml"

usage() {
cat <<EOF
CHBH environment manager

Usage:
  chbh install
      Install the chbh command and environment template.

  chbh version
      Show the version of the installed CHBH environment template.

  chbh bear-modules
      Load the standard BlueBEAR modules required for CHBH.

  chbh summarise-logs <log-directory>
      Summarise SLURM logs in the specified directory.

  chbh venv <group> [group ...]
      Create or update a UV environment in the current directory.

Available venv groups:
  meeg     EEG, MEG and electrophysiology tools
  mri      MRI and neuroimaging tools
  sleep    Sleep analysis tools
  dev      Development and code-quality tools

Examples:
  chbh install

  chbh version

  chbh bear-modules

  chbh summarise-logs ./logs

  chbh venv meeg

  chbh venv mri

  chbh venv meeg mri

  chbh venv meeg sleep

  chbh venv meeg mri sleep dev

Notes:
  - A .venv is created in the current working directory.
  - If no pyproject.toml exists, the CHBH template is copied into
    the current directory.
  - Existing pyproject.toml files are preserved and used as-is.
  - Activate the environment with:

        source .venv/bin/activate

EOF
}

# %% ---------------------
# Helper functions

cmd_install() {
    echo "Installing from: $SCRIPT_DIR to: $HOME/.local/bin"

    mkdir -p "$HOME/.local/bin"
    mkdir -p "$HOME/.local/share/chbh"

    install -v -m 755 \
        "$SCRIPT_DIR/chbh" \
        "$HOME/.local/bin/chbh"

    install -v -m 644 \
        "$SCRIPT_DIR/pyproject.toml" \
        "$HOME/.local/share/chbh/pyproject.toml"

    install -v -m 644 \
        "$SCRIPT_DIR/chbh_summarise_slurm_logs.py" \
        "$HOME/.local/share/chbh/chbh_summarise_slurm_logs.py"

    echo
    echo "Installed: $HOME/.local/bin/chbh"

    case ":$PATH:" in
        *":$HOME/.local/bin:"*)
            echo "Command available as: chbh"
            ;;
        *)
            echo
            echo "WARNING: $HOME/.local/bin is not on your PATH."
            echo
            echo 'Add the following to ~/.bashrc:'
            echo 'export PATH="$HOME/.local/bin:$PATH"'
            echo
            echo 'Then reload your shell:'
            echo 'source ~/.bashrc'
            ;;
    esac
}

cmd_add_base_modules() {
    echo Copy and run the following:
    echo 
    echo module purge
    echo module load bluebear
    echo module load bear-apps/2023a
    echo module load cairo/1.17.8-GCCcore-12.3.0
    echo module load uv/0.6.5
}

ensure_pyproject() {
    local template="$HOME/.local/share/chbh/pyproject.toml"

    if [[ ! -f pyproject.toml ]]; then
        cp "$template" pyproject.toml
        echo "Created pyproject.toml from CHBH template."
    else
        echo "Using existing pyproject.toml in working directory."
        echo "This may not match the current CHBH specification."
    fi
}

cmd_version() {
    if [[ ! -f "$TEMPLATE_PYPROJECT" ]]; then
        echo "CHBH template not installed."
        exit 1
    fi

    grep '^version = ' "$TEMPLATE_PYPROJECT" \
        | sed 's/version = "\(.*\)"/\1/'
}

cmd_summarise_logs() {
    uv run \
        "$HOME/.local/share/chbh/chbh_summarise_slurm_logs.py" \
        "$1"
}

# %% ---------------------
# Show usage if no arguments

if [[ $# -eq 0 ]]; then
    usage
    exit 1
fi

# %% ---------------------
# Top-level commands

case "${1:-}" in

    install)
        cmd_install
        exit 0
        ;;

    version)
        cmd_version
        exit 0
        ;;

    bear-modules)
        cmd_add_base_modules
        exit 0
        ;;

    summarise-logs)
        if [[ $# -ne 2 ]]; then
            echo "Usage: chbh summarise-logs <log-directory>"
            exit 1
        fi

        cmd_summarise_logs "$2"
        exit 0
        ;;

    venv)
        shift
        ;;

    *)
        echo "Unknown command: ${1:-}"
        echo
        usage
        exit 1
        ;;

esac

# %% ---------------------
# Parse venv groups

if [[ $# -eq 0 ]]; then
    echo "No venv groups specified."
    echo
    usage
    exit 1
fi

ARGS=()

for group in "$@"; do
    case "$group" in
        meeg|mri|sleep|dev)
            ARGS+=(--group "$group")
            ;;
        *)
            echo "Unknown venv group: $group"
            echo
            usage
            exit 1
            ;;
    esac
done

# %% ---------------------
# Create environment

ensure_pyproject

# Avoid warnings if script is run from an active venv
unset VIRTUAL_ENV || true

uv sync "${ARGS[@]}"

cat <<EOF

Environment ready.

Activate with:

    source .venv/bin/activate

Installed groups:

    $*

EOF