#!/bin/bash

# Function to handle errors
handle_error() {
    echo "Error: $1"
    exit 1
}

# Verify and create the /data directory if it doesn't exist
if [ ! -d "/data" ]; then
    mkdir -p /data || handle_error "Could not create /data directory"
    echo "/data directory created for persistent storage"
fi

# Verify write permissions for /data directory
if [ ! -w "/data" ]; then
    echo "Warning: No write permissions for /data. Some data may not be persistent."
fi

# Use the AUX_TOKEN environment variable if JUPYTER_TOKEN is not set
if [ -z "$JUPYTER_TOKEN" ]; then
    export JUPYTER_TOKEN=${AUX_TOKEN:-}
fi

# Verify if the token is empty
if [ -z "$JUPYTER_TOKEN" ]; then
    handle_error "JupyterLab token is empty"
fi

# Check for GPU availability
if command -v nvidia-smi &> /dev/null; then
    echo "GPU detected. Configuring environment for GPU usage."
    export NVIDIA_VISIBLE_DEVICES=all
    export NVIDIA_DRIVER_CAPABILITIES=compute,utility
else
    echo "No GPU detected. CPU will be used."
fi

# Start JupyterLab in the background
jupyter lab --ip=0.0.0.0 --port=${JUPYTERLAB_PORT} --no-browser --allow-root \
    --NotebookApp.token=${JUPYTER_TOKEN} \
    --notebook-dir=/data &

# Start Nginx in the foreground
nginx -g "daemon off;" || handle_error "Failed to start Nginx"