#!/bin/bash
set -e

echo "Starting VS Code with Hugging Face persistence..."

# Check if HF_TOKEN and HF_DATASET are set
if [ -z "$HF_TOKEN" ] || [ -z "$HF_DATASET" ]; then
    echo "Warning: HF_TOKEN or HF_DATASET not set. Running without persistence."
    exec code-server --bind-addr 0.0.0.0:7860 --auth password /home/coder/workspace
    exit 0
fi

# Login to Hugging Face
echo "Logging in to Hugging Face..."
huggingface-cli login --token "$HF_TOKEN" --add-to-git-credential

# Directories to persist
CONFIG_DIR="/home/coder/.config/code-server"
LOCAL_DIR="/home/coder/.local/share/code-server"
WORKSPACE_DIR="/home/coder/workspace"
BACKUP_DIR="/tmp/hf_backup"

# Create backup directory
mkdir -p "$BACKUP_DIR"

# Function to download from HF dataset
download_from_hf() {
    echo "Downloading persisted data from Hugging Face dataset: $HF_DATASET"
    
    # Try to download the backup archive
    if huggingface-cli download "$HF_DATASET" backup.tar.gz --repo-type dataset --local-dir "$BACKUP_DIR" 2>/dev/null; then
        echo "Found existing backup, restoring..."
        
        # Extract the backup
        cd /home/coder
        tar -xzf "$BACKUP_DIR/backup.tar.gz" 2>/dev/null || echo "No previous backup found or extraction failed"
        
        echo "Restoration complete!"
    else
        echo "No existing backup found. Starting fresh."
    fi
}

# Function to upload to HF dataset
upload_to_hf() {
    echo "Saving data to Hugging Face dataset..."
    
    # Create backup archive
    cd /home/coder
    tar -czf "$BACKUP_DIR/backup.tar.gz" \
        --exclude='*.log' \
        --exclude='*.sock' \
        --exclude='CachedData' \
        --exclude='Cache' \
        --exclude='GPUCache' \
        .config/code-server \
        .local/share/code-server \
        workspace 2>/dev/null || true
    
    # Upload to HF dataset
    huggingface-cli upload "$HF_DATASET" "$BACKUP_DIR/backup.tar.gz" backup.tar.gz --repo-type dataset --create || \
        echo "Upload failed, continuing anyway..."
}

# Setup trap to save on exit
trap 'upload_to_hf' EXIT SIGTERM SIGINT

# Download existing data
download_from_hf

# Start background sync every 5 minutes
(
    while true; do
        sleep 300
        upload_to_hf
    done
) &

SYNC_PID=$!

echo "Starting code-server..."
# Start code-server
code-server --bind-addr 0.0.0.0:7860 --auth password /home/coder/workspace &
CODE_SERVER_PID=$!

# Wait for code-server to exit
wait $CODE_SERVER_PID

# Kill background sync
kill $SYNC_PID 2>/dev/null || true

# Final upload
upload_to_hf

echo "Shutdown complete."