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

# Define directories
NOTEBOOK_DIR="/data"
ORIGINAL_NOTEBOOKS="/home/user/app"  # This is where your notebooks likely are
BACKUP_DIR="/home/user/original_notebooks"

# Add debugging to see what's happening
echo "Starting notebook setup process..."
echo "Looking for notebooks in: $ORIGINAL_NOTEBOOKS"
ls -la $ORIGINAL_NOTEBOOKS/*.ipynb || echo "No notebooks found in $ORIGINAL_NOTEBOOKS"

# Create backup directory if it doesn't exist
if [ ! -d "$BACKUP_DIR" ]; then
    echo "Creating backup directory for original notebooks..."
    mkdir -p "$BACKUP_DIR"
    
    # Try to find notebooks in various possible locations
    if ls $ORIGINAL_NOTEBOOKS/*.ipynb 1> /dev/null 2>&1; then
        echo "Found notebooks in $ORIGINAL_NOTEBOOKS"
        cp -v $ORIGINAL_NOTEBOOKS/*.ipynb "$BACKUP_DIR/"
    elif ls /notebooks/*.ipynb 1> /dev/null 2>&1; then
        echo "Found notebooks in /notebooks"
        cp -v /notebooks/*.ipynb "$BACKUP_DIR/"
    elif ls /home/user/*.ipynb 1> /dev/null 2>&1; then
        echo "Found notebooks in /home/user"
        cp -v /home/user/*.ipynb "$BACKUP_DIR/"
    else
        echo "No notebooks found in expected locations. Checking all directories..."
        find / -name "*.ipynb" -not -path "*/\.*" 2>/dev/null
    fi
fi

# Clear the notebook directory and copy fresh notebooks for this session
echo "Preparing fresh notebooks for this session..."
mkdir -p "$NOTEBOOK_DIR"
rm -f "$NOTEBOOK_DIR"/*.ipynb
ls -la "$BACKUP_DIR"/*.ipynb || echo "No notebooks found in backup directory"
cp -v "$BACKUP_DIR"/*.ipynb "$NOTEBOOK_DIR/" || echo "Failed to copy notebooks"

# Set proper permissions
chmod -R 755 "$NOTEBOOK_DIR"
chown -R $(whoami) "$NOTEBOOK_DIR"

# Show what we've got
echo "Contents of notebook directory ($NOTEBOOK_DIR):"
ls -la "$NOTEBOOK_DIR"

# Start JupyterLab
echo "Starting JupyterLab..."
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