#!/bin/bash
# Setup script for GPU droplet with NVENC support
# Run as root on a fresh Ubuntu droplet with NVIDIA GPU

set -e

echo "=== ArtDAG GPU Droplet Setup ==="

# 1. System updates
echo "[1/7] Updating system..."
apt-get update
apt-get upgrade -y

# 2. Install NVIDIA drivers (if not already installed)
echo "[2/7] Checking NVIDIA drivers..."
if ! command -v nvidia-smi &> /dev/null; then
    echo "Installing NVIDIA drivers..."
    apt-get install -y nvidia-driver-535 nvidia-utils-535
    echo "NVIDIA drivers installed. Reboot required."
    echo "After reboot, run this script again."
    exit 0
fi

nvidia-smi
echo "NVIDIA drivers OK"

# 3. Install FFmpeg with NVENC support
echo "[3/7] Installing FFmpeg with NVENC..."
apt-get install -y ffmpeg

# Verify NVENC
if ffmpeg -encoders 2>/dev/null | grep -q nvenc; then
    echo "NVENC available:"
    ffmpeg -encoders 2>/dev/null | grep nvenc
else
    echo "WARNING: NVENC not available. GPU may not support hardware encoding."
fi

# 4. Install Python and create venv
echo "[4/7] Setting up Python environment..."
apt-get install -y python3 python3-venv python3-pip git

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

# 5. Install Python dependencies
echo "[5/7] Installing Python packages..."
pip install --upgrade pip
pip install \
    numpy \
    opencv-python-headless \
    wgpu \
    httpx \
    pyyaml \
    celery[redis] \
    fastapi \
    uvicorn

# 6. Clone/update art-dag code
echo "[6/7] Setting up art-dag code..."
ARTDAG_DIR="$VENV_DIR/celery"
if [ -d "$ARTDAG_DIR" ]; then
    echo "Updating existing code..."
    cd "$ARTDAG_DIR"
    git pull || true
else
    echo "Cloning art-dag..."
    git clone https://git.rose-ash.com/art-dag/celery.git "$ARTDAG_DIR" || {
        echo "Git clone failed. You may need to copy code manually."
    }
fi

# 7. Test GPU compute
echo "[7/7] Testing GPU compute..."
"$VENV_DIR/bin/python3" << 'PYTEST'
import sys
try:
    import wgpu
    adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance")
    print(f"GPU Adapter: {adapter.info.get('device', 'unknown')}")
    device = adapter.request_device_sync()
    print("wgpu device created successfully")

    # Check for NVENC via FFmpeg
    import subprocess
    result = subprocess.run(['ffmpeg', '-encoders'], capture_output=True, text=True)
    if 'h264_nvenc' in result.stdout:
        print("NVENC H.264 encoder: AVAILABLE")
    else:
        print("NVENC H.264 encoder: NOT AVAILABLE")
    if 'hevc_nvenc' in result.stdout:
        print("NVENC HEVC encoder: AVAILABLE")
    else:
        print("NVENC HEVC encoder: NOT AVAILABLE")

except Exception as e:
    print(f"Error: {e}")
    sys.exit(1)
PYTEST

echo ""
echo "=== Setup Complete ==="
echo "Venv: $VENV_DIR"
echo "Code: $ARTDAG_DIR"
echo ""
echo "To activate: source $VENV_DIR/bin/activate"
echo "To test: cd $ARTDAG_DIR && python -c 'from streaming.backends import get_backend; print(get_backend(\"wgpu\"))'"