#!/bin/bash
# Deploy Maya3 Pipeline to a remote cluster
#
# Usage:
#   ./deploy_cluster.sh <host> <ssh_key> [user]
#
# Example:
#   ./deploy_cluster.sh 10.0.0.1 ~/.ssh/cluster01 root

set -e

HOST=$1
SSH_KEY=$2
USER=${3:-root}

if [ -z "$HOST" ] || [ -z "$SSH_KEY" ]; then
    echo "Usage: $0 <host> <ssh_key> [user]"
    exit 1
fi

SSH_OPTS="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -i $SSH_KEY"
SSH_CMD="ssh $SSH_OPTS $USER@$HOST"
SCP_CMD="scp $SSH_OPTS"

echo "=========================================="
echo "Deploying Maya3 Pipeline to $HOST"
echo "=========================================="

# 1. Check connectivity
echo "[1/6] Testing SSH connection..."
$SSH_CMD "echo 'SSH OK' && nvidia-smi --query-gpu=name,memory.total --format=csv,noheader"

# 2. Create workspace
echo "[2/6] Creating workspace..."
$SSH_CMD "mkdir -p /workspace/maya3_data"

# 3. Sync codebase (excluding large files)
echo "[3/6] Syncing codebase..."
rsync -avz --progress \
    -e "ssh $SSH_OPTS" \
    --exclude '.git' \
    --exclude '.venv' \
    --exclude 'output' \
    --exclude 'data' \
    --exclude '__pycache__' \
    --exclude '*.tar' \
    --exclude '*.webm' \
    --exclude '*.flac' \
    /workspace/maya3_data/ \
    $USER@$HOST:/workspace/maya3_data/

# 4. Copy .env file
echo "[4/6] Copying environment file..."
$SCP_CMD /workspace/maya3_data/.env $USER@$HOST:/workspace/maya3_data/.env

# 5. Setup Python environment
echo "[5/6] Setting up Python environment..."
$SSH_CMD << 'REMOTE_SCRIPT'
cd /workspace/maya3_data

# Create venv if not exists
if [ ! -d ".venv" ]; then
    python3 -m venv .venv
fi

# Install dependencies
.venv/bin/pip install --upgrade pip
.venv/bin/pip install -r requirements.txt

# Verify GPU access
.venv/bin/python -c "import torch; print(f'GPUs: {torch.cuda.device_count()}')"
REMOTE_SCRIPT

# 6. Verify installation
echo "[6/6] Verifying installation..."
$SSH_CMD << 'VERIFY_SCRIPT'
cd /workspace/maya3_data
source .env

# Test Supabase connection
.venv/bin/python -c "
from src.supabase_client import SupabaseClient
client = SupabaseClient()
stats = client.get_stats()
print(f'Supabase OK - Pending: {stats.get(\"PENDING\", 0)}')
"

# Test R2 connection
.venv/bin/python -c "
from src.r2_client import R2Client
r2 = R2Client(bucket_type='source')
print(f'R2 OK - Bucket: {r2.bucket}')
"
VERIFY_SCRIPT

echo "=========================================="
echo "Deployment complete: $HOST"
echo "=========================================="
echo ""
echo "To start processing:"
echo "  $SSH_CMD 'cd /workspace/maya3_data && .venv/bin/python massive_process.py --supabase-queue --r2-upload'"