#!/bin/bash
JUPYTER_TOKEN="${JUPYTER_TOKEN:=huggingface}"

NOTEBOOK_DIR="/data/workspaces"

# Set up NVIDIA EGL library links at runtime (in case they're mounted differently)
echo "🔗 Setting up NVIDIA EGL library links..."
for nvidia_lib_dir in /usr/local/nvidia/lib64 /usr/local/cuda/lib64 /usr/lib/nvidia; do
    if [ -f "$nvidia_lib_dir/libEGL_nvidia.so.0" ]; then
        echo "Found NVIDIA EGL library at $nvidia_lib_dir/libEGL_nvidia.so.0"
        ln -sf "$nvidia_lib_dir/libEGL_nvidia.so.0" /usr/lib/x86_64-linux-gnu/libEGL_nvidia.so.0
        break
    fi
done

# perform checks on the GPU configuration
python init_gpu.py

# this will download stuff used by Mujoco (the collection of models)
python init_mujoco.py

# Copy sample notebooks to persistent storage
echo "📚 Setting up sample workspaces in /data/workspaces..."

# Create the workspaces directory if it doesn't exist
mkdir -p /data/workspaces

# Create JupyterLab workspace state directory (for UI preferences, layout, etc.)
mkdir -p /data/.jupyter/workspaces
mkdir -p /data/.jupyter/settings
echo "✓ JupyterLab workspace state directory: /data/.jupyter/workspaces"
echo "✓ JupyterLab settings directory: /data/.jupyter/settings"

# Loop through each sample directory
for sample_dir in $HOME/app/samples/*/; do
    # Get the sample name (e.g., "locomotion" from "samples/locomotion/")
    sample_name=$(basename "$sample_dir")

    # Create the workspace directory if it doesn't exist
    workspace_dir="/data/workspaces/$sample_name"
    if [ ! -d "$workspace_dir" ]; then
        echo "  Creating workspace: $workspace_dir"
        mkdir -p "$workspace_dir"
    else
        echo "  Workspace already exists: $workspace_dir"
    fi

    # Copy the .ipynb file only if it doesn't already exist (to preserve user changes)
    ipynb_file="$sample_dir$sample_name.ipynb"
    dest_ipynb="$workspace_dir/$sample_name.ipynb"

    if [ -f "$ipynb_file" ]; then
        if [ ! -f "$dest_ipynb" ]; then
            echo "  Copying: $sample_name.ipynb → $workspace_dir/"
            cp "$ipynb_file" "$dest_ipynb"
        else
            echo "  Preserving existing: $dest_ipynb (user may have made changes)"
        fi
    fi

    # Copy any other files from the sample directory (excluding .ipynb files)
    # This allows samples to include datasets, scripts, etc.
    for file in "$sample_dir"*; do
        filename=$(basename "$file")
        # Skip if it's the .ipynb file (already handled) or if it's a directory
        if [ "$filename" != "$sample_name.ipynb" ] && [ -f "$file" ]; then
            dest_file="$workspace_dir/$filename"
            if [ ! -f "$dest_file" ]; then
                echo "  Copying additional file: $filename → $workspace_dir/"
                cp "$file" "$dest_file"
            fi
        fi
    done
done

echo "✅ Sample workspaces ready!"
echo ""

# Initialize OpenTrack (download datasets, create symlinks)
if [ -f "$HOME/app/samples/opentrack/init_opentrack.sh" ]; then
    bash "$HOME/app/samples/opentrack/init_opentrack.sh"
fi

jupyter labextension disable "@jupyterlab/apputils-extension:announcements"

jupyter-lab \
    --ip 0.0.0.0 \
    --port 7860 \
    --no-browser \
    --allow-root \
    --ServerApp.token="$JUPYTER_TOKEN" \
    --ServerApp.tornado_settings="{'headers': {'Content-Security-Policy': 'frame-ancestors *'}}" \
    --ServerApp.cookie_options="{'SameSite': 'None', 'Secure': True}" \
    --ServerApp.disable_check_xsrf=True \
    --LabApp.news_url=None \
    --LabApp.check_for_updates_class="jupyterlab.NeverCheckForUpdate" \
    --notebook-dir=$NOTEBOOK_DIR \
    --ServerApp.default_url="/data/workspaces/opentrack/opentrack_quickstart.ipynb"