#!/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 using Python
echo "Logging in to Hugging Face..."
python3 -c "from huggingface_hub import login; login(token='$HF_TOKEN', add_to_git_credential=True)"

# 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 using Python
    python3 << 'PYTHON_SCRIPT'
import os
from huggingface_hub import hf_hub_download
import tarfile

try:
    backup_dir = "/tmp/hf_backup"
    os.makedirs(backup_dir, exist_ok=True)
    
    # Download backup file
    file_path = hf_hub_download(
        repo_id=os.environ["HF_DATASET"],
        filename="backup.tar.gz",
        repo_type="dataset",
        local_dir=backup_dir
    )
    
    print(f"Found existing backup at {file_path}, restoring...")
    
    # Extract the backup
    with tarfile.open(file_path, "r:gz") as tar:
        tar.extractall("/home/coder")
    
    print("Restoration complete!")
except Exception as e:
    print(f"No existing backup found or restoration failed: {e}")
    print("Starting fresh.")
PYTHON_SCRIPT
}

# 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 \
        .continue \
        workspace 2>/dev/null || true
    
    # Upload to HF dataset using Python
    python3 << 'PYTHON_SCRIPT'
import os
from huggingface_hub import HfApi

try:
    api = HfApi()
    backup_file = "/tmp/hf_backup/backup.tar.gz"
    
    # Create repo if it doesn't exist and upload
    api.create_repo(
        repo_id=os.environ["HF_DATASET"],
        repo_type="dataset",
        private=True,
        exist_ok=True
    )
    
    api.upload_file(
        path_or_fileobj=backup_file,
        path_in_repo="backup.tar.gz",
        repo_id=os.environ["HF_DATASET"],
        repo_type="dataset"
    )
    
    print("Upload successful!")
except Exception as e:
    print(f"Upload failed: {e}")
    print("Continuing anyway...")
PYTHON_SCRIPT
}

# 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."