#!/bin/bash

# Default values
export GKFS_MNT=${GKFS_MNT:-/dev/shm/gkfs_mnt}
export GKFS_ROOT=${GKFS_ROOT:-/dev/shm/gkfs_root}
export GKFS_DAEMON_LOG_LEVEL=${GKFS_DAEMON_LOG_LEVEL:-0}
export GKFS_HOSTS_FILE_TIMEOUT=${GKFS_HOSTS_FILE_TIMEOUT:-600}

# Helper function to print usage
usage() {
    echo "Usage: $0 {start|stop|start_fuse|stop_fuse|help}"
    echo
    echo "Commands:"
    echo "  start       Start GekkoFS daemons on allocated nodes"
    echo "  stop        Stop GekkoFS daemons and clean up"
    echo "  start_fuse  Start GekkoFS FUSE clients (requires servers running)"
    echo "  stop_fuse   Stop GekkoFS FUSE clients"
    echo "  help        Show this help message"
}

# Function to start daemons
start_servers() {
    if [ -z "$SLURM_JOB_NUM_NODES" ]; then
        echo "Error: SLURM_JOB_NUM_NODES not set. Are you in a Slurm allocation?"
        exit 1
    fi

    echo "Cleaning up previous run..."
    srun -n ${SLURM_JOB_NUM_NODES} -N ${SLURM_JOB_NUM_NODES} --oversubscribe \
         bash -c "rm -rf ${GKFS_MNT} ${GKFS_ROOT} ; mkdir -p ${GKFS_MNT} ${GKFS_ROOT}"

    # Use GKFS_HOSTS_FILE from environment or default to HOME
    if [ -z "$GKFS_HOSTS_FILE" ]; then
        export GKFS_HOSTS_FILE=${HOME}/gkfs_hosts.txt
    fi
    export LIBGKFS_HOSTS_FILE=${GKFS_HOSTS_FILE}
    
    # Create empty hosts file
    touch ${GKFS_HOSTS_FILE}

    COMM="-P ofi+verbs"
    CMD="${GKFS_DAEMON} --mountdir=${GKFS_MNT} --rootdir=${GKFS_ROOT} ${COMM} -l ib0 -c"

    echo "Starting GekkoFS Daemons on ${SLURM_JOB_NUM_NODES} nodes..."
    echo "Command: ${CMD}"

    # Unset conflicting env vars
    unset I_MPI_PMI_LIBRARY
    export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0

    # Launch daemons
    srun -N ${SLURM_JOB_NUM_NODES} -n ${SLURM_JOB_NUM_NODES} \
         -c 112 --overlap --overcommit --mem=0 --oversubscribe --export="ALL" \
         /bin/bash -c "${CMD}" &
    
    # Wait for servers to be ready. Calibration adds comment lines to the
    # hostfile, so count only non-empty, non-comment daemon entries.
    echo "Waiting for servers to start..."
    WAITED=0
    while true; do
        if [ -f "$GKFS_HOSTS_FILE" ]; then
            HOST_COUNT=$(awk '
                /^[[:space:]]*#/ { next }
                /^[[:space:]]*$/ { next }
                { count++ }
                END { print count + 0 }
            ' "$GKFS_HOSTS_FILE")
            if [ "$HOST_COUNT" -ge "${SLURM_JOB_NUM_NODES}" ]; then
                echo "All ${SLURM_JOB_NUM_NODES} servers started."
                break
            fi
        fi
        if [ "$WAITED" -ge "$GKFS_HOSTS_FILE_TIMEOUT" ]; then
            echo "Error: timed out after ${GKFS_HOSTS_FILE_TIMEOUT}s waiting for ${SLURM_JOB_NUM_NODES} servers (found ${HOST_COUNT:-0} host entries)."
            return 1
        fi
        sleep 1
        WAITED=$((WAITED + 1))
    done
    
    echo "GekkoFS servers are ready."
    echo "Mount point: ${GKFS_MNT}"
    echo "Hosts file: ${GKFS_HOSTS_FILE}"
}

start_fuse() {
    # Check if servers are running (simplistic check via hosts file)
    # We allow user to run this independently, but warn if servers might not be ready
    if [ -z "$GKFS_HOSTS_FILE" ]; then
         export GKFS_HOSTS_FILE=${HOME}/gkfs_hosts.txt
    fi
    
    if [ ! -f "$GKFS_HOSTS_FILE" ]; then
        echo "Warning: Hosts file $GKFS_HOSTS_FILE not found. Are servers running?"
    fi

    echo "Starting GekkoFS FUSE Clients..."
    if [ -z "$GKFS_FUSE" ]; then
        echo "Error: GKFS_FUSE not set. Did you load the fuse module?"
        exit 1
    fi
    
    CMD2="${GKFS_FUSE} -o max_idle_threads=32 -o direct_io -f -o fifo -o auto_unmount $GKFS_MNT"
    
    srun -N ${SLURM_JOB_NUM_NODES} \
         -n ${SLURM_JOB_NUM_NODES} \
         --export="ALL" --oversubscribe --overlap \
         /bin/bash -c "${CMD2}" &
         
    echo "FUSE clients launched."
}

# Function to stop daemons
stop_servers() {
    if [ -z "$SLURM_JOB_NUM_NODES" ]; then
        echo "Error: SLURM_JOB_NUM_NODES not set. Are you in a Slurm allocation?"
        exit 1
    fi

    echo "Stopping GekkoFS daemons..."
    srun -n ${SLURM_JOB_NUM_NODES} -N ${SLURM_JOB_NUM_NODES} \
         -c 1 --mem=0 --oversubscribe --export="ALL" \
         /bin/bash -c "pkill --signal SIGINT gkfs_daemon ; rm -rf ${GKFS_MNT} ${GKFS_ROOT}"
    
    echo "Server cleanup complete."
}

stop_fuse() {
    if [ -z "$SLURM_JOB_NUM_NODES" ]; then
        echo "Error: SLURM_JOB_NUM_NODES not set. Are you in a Slurm allocation?"
        exit 1
    fi
    
    echo "Stopping GekkoFS FUSE clients..."
    srun -n ${SLURM_JOB_NUM_NODES} -N ${SLURM_JOB_NUM_NODES} \
         -c 1 --mem=0 --oversubscribe --export="ALL" \
         /bin/bash -c "pkill --signal SIGINT fuse_client"

    echo "FUSE cleanup complete."
}

# Main logic
case "$1" in
    start)
        start_servers
        ;;
    stop)
        stop_servers
        ;;
    start_fuse)
        start_fuse
        ;;
    stop_fuse)
        stop_fuse
        ;;
    help)
        usage
        ;;
    *)
        usage
        exit 1
        ;;
esac
