#!/bin/bash

# Ensure /data exists and has the correct permissions
if [ ! -d /data ]; then
    sudo mkdir -p /data
    sudo chown user:user /data
fi

# Copy notebooks to /data, appending _UPDATED to filenames if they have changed
shopt -s globstar
for file in /home/user/notebooks/**/*; do
    if [ -f "$file" ]; then
        relative_path="${file#/home/user/notebooks/}"
        target_file="/data/$relative_path"
        target_dir=$(dirname "$target_file")
        if [ ! -d "$target_dir" ]; then
            mkdir -p "$target_dir"
        fi
        if [ -e "$target_file" ]; then
            if ! cmp -s "$file" "$target_file"; then
                new_filename="${target_file%.*}_UPDATED.${target_file##*.}"
                cp "$file" "$new_filename"
            fi
        else
            cp "$file" "$target_file"
        fi
    fi
done

# GPU detection and package installation
if command -v nvidia-smi &> /dev/null && nvidia-smi -L; then
    echo "GPU detected. Installing GPU-specific Python packages."
    pip install --no-cache-dir -r /home/user/notebooks/requirements_gpu.txt
else
    echo "No GPU detected. Installing CPU-specific Python packages."
    pip install --no-cache-dir -r /home/user/notebooks/requirements_cpu.txt
fi

JUPYTER_TOKEN="${JUPYTER_TOKEN:=huggingface}"

echo "Starting Jupyter Lab with token $JUPYTER_TOKEN"

NOTEBOOK_DIR="/data"

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