#!/bin/bash
# Cloud-init startup script for GPU droplet (RTX 6000 Ada, etc.)
# Paste this into DigitalOcean "User data" field when creating droplet

set -e
export DEBIAN_FRONTEND=noninteractive
exec > /var/log/artdag-setup.log 2>&1

echo "=== ArtDAG GPU Setup Started $(date) ==="

# Update system (non-interactive, keep existing configs)
apt-get update
apt-get -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" upgrade

# Install essentials
apt-get install -y \
    python3 python3-venv python3-pip \
    git curl wget \
    ffmpeg \
    vulkan-tools \
    build-essential

# Create venv
VENV_DIR="/opt/artdag-gpu"
python3 -m venv "$VENV_DIR"
source "$VENV_DIR/bin/activate"

# Install Python packages
pip install --upgrade pip
pip install \
    numpy \
    opencv-python-headless \
    wgpu \
    httpx \
    pyyaml \
    celery[redis] \
    fastapi \
    uvicorn \
    asyncpg

# Create code directory
mkdir -p "$VENV_DIR/celery/sexp_effects/effects"
mkdir -p "$VENV_DIR/celery/sexp_effects/primitive_libs"
mkdir -p "$VENV_DIR/celery/streaming"

# Add SSH key for easier access (optional - add your key here)
# echo "ssh-ed25519 AAAA... your-key" >> /root/.ssh/authorized_keys

# Test GPU
echo "=== GPU Info ==="
nvidia-smi || echo "nvidia-smi not available yet"

echo "=== NVENC Check ==="
ffmpeg -encoders 2>/dev/null | grep -E "nvenc|cuda" || echo "NVENC not detected"

echo "=== wgpu Check ==="
"$VENV_DIR/bin/python3" -c "
import wgpu
try:
    adapter = wgpu.gpu.request_adapter_sync(power_preference='high-performance')
    print(f'GPU: {adapter.info}')
except Exception as e:
    print(f'wgpu error: {e}')
" || echo "wgpu test failed"

# Add environment setup
cat >> /etc/profile.d/artdag-gpu.sh << 'ENVEOF'
export WGPU_BACKEND_TYPE=Vulkan
export PATH="/opt/artdag-gpu/bin:$PATH"
ENVEOF

# Mark setup complete
touch /opt/artdag-gpu/.setup-complete
echo "=== Setup Complete $(date) ==="
echo "Venv: /opt/artdag-gpu"
echo "Activate: source /opt/artdag-gpu/bin/activate"
echo "Vulkan: export WGPU_BACKEND_TYPE=Vulkan"