#!/usr/bin/env bash
# Purpose: mount the cluster's shared model library on this machine, so that every
#          worker loads the same weights from one place. Prints exactly what it would
#          do and changes nothing until you pass --apply
# Platform: spark, strix, nvidia (Linux with the Ubuntu NFS client package).
#           Track M mounts the same export through the Finder; see the lab page.
# Minimum memory: 8 GB
# Assumes: an .env beside this script with NFS_SERVER, NFS_EXPORT and MODELS_DIR
#          filled in, sudo available, and setup-nfs-server.sh already run on the
#          machine holding the library. It creates one directory and mounts one
#          export; it never writes to /etc/fstab.
set -euo pipefail

HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ENV_FILE="${ENV_FILE:-${HERE}/.env}"
APPLY=0

for arg in "$@"; do
    case "$arg" in
        --apply) APPLY=1 ;;
        -h|--help)
            echo "Usage: $0 [--apply]"
            echo "  Without --apply it prints the plan and changes nothing."
            exit 0
            ;;
        *)
            echo "Unknown argument: ${arg}" >&2
            exit 1
            ;;
    esac
done

if [ -f "$ENV_FILE" ]; then
    set -a
    # shellcheck source=/dev/null
    . "$ENV_FILE"
    set +a
else
    echo "No ${ENV_FILE}. Copy env-example.txt to .env and fill it in." >&2
    exit 1
fi

: "${NFS_SERVER:=}"
: "${NFS_EXPORT:=}"
: "${MODELS_DIR:=}"

if [ "$(uname -s)" != "Linux" ]; then
    echo "This script uses the Linux NFS client packages and $(uname -s) is not Linux." >&2
    echo "On macOS, mount the same export from the Finder with Go > Connect to Server," >&2
    echo "then record the mount point in your notebook by hand." >&2
    exit 1
fi

for name in NFS_SERVER NFS_EXPORT MODELS_DIR; do
    if [ -z "${!name}" ]; then
        echo "${name} is empty in ${ENV_FILE}." >&2
        exit 1
    fi
done

if mountpoint -q "$MODELS_DIR" 2>/dev/null; then
    echo "${MODELS_DIR} is already a mount point. Nothing to do."
    exit 0
fi

if [ -d "$MODELS_DIR" ] && [ -n "$(ls -A "$MODELS_DIR" 2>/dev/null)" ]; then
    echo "${MODELS_DIR} already exists and is not empty." >&2
    echo "Mounting over it would hide what is there. Move it aside or choose another" >&2
    echo "path in .env before continuing." >&2
    exit 1
fi

echo "Plan for this machine:"
echo ""
echo "  1. install the client package:"
echo "       sudo apt install nfs-common"
echo ""
echo "  2. create the mount point:"
echo "       sudo mkdir -p ${MODELS_DIR}"
echo ""
echo "  3. mount the shared library:"
echo "       sudo mount ${NFS_SERVER}:${NFS_EXPORT} ${MODELS_DIR}"
echo ""
echo "  The export is read-only on the server side, so nothing on this machine can"
echo "  change the library by accident. The mount does not survive a reboot; add an"
echo "  /etc/fstab entry yourself once you are happy with it, as the lab page says."
echo ""

if [ "$APPLY" -ne 1 ]; then
    echo "Nothing was changed. Re-run with --apply to carry out the plan above."
    exit 0
fi

echo "==> installing nfs-common"
sudo apt install nfs-common

echo "==> creating ${MODELS_DIR}"
sudo mkdir -p "$MODELS_DIR"

echo "==> mounting ${NFS_SERVER}:${NFS_EXPORT}"
sudo mount "${NFS_SERVER}:${NFS_EXPORT}" "$MODELS_DIR"

echo ""
echo "Mounted. Confirm the machine can read the library with:"
echo "  ls ${MODELS_DIR}"
echo "Unmount again with:"
echo "  sudo umount ${MODELS_DIR}"