#!/bin/bash
# OCR Worker Deployment Script
# Run on each GPU server to deploy workers

set -e

echo "=========================================="
echo "OCR Worker Deployment"
echo "=========================================="

# Configuration
BACKEND_DIR="/home/local/translynx/backend"
SYSTEMD_DIR="/etc/systemd/system"
DB_HOST="192.168.1.161"

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo "Please run with sudo"
    exit 1
fi

# Step 1: Check GPU
echo ""
echo "[1/5] Checking GPU..."
GPU_COUNT=$(nvidia-smi -L 2>/dev/null | wc -l || echo 0)
echo "  Found $GPU_COUNT GPU(s)"
nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader 2>/dev/null || echo "  No NVIDIA GPU found"

# Step 2: Check database connectivity
echo ""
echo "[2/5] Testing database connectivity to $DB_HOST..."
if PGPASSWORD=translynx2025 psql -h $DB_HOST -U translynx -d translynx -c "SELECT 1" > /dev/null 2>&1; then
    echo "  Database connection OK"
else
    echo "  ERROR: Cannot connect to database at $DB_HOST"
    echo "  Make sure PostgreSQL is configured to accept connections"
    exit 1
fi

# Step 3: Check paddlex environment
echo ""
echo "[3/5] Checking paddlex conda environment..."
PADDLEX_PYTHON="/home/ubuntu/miniconda3/envs/paddlex/bin/python"
if [ -f "$PADDLEX_PYTHON" ]; then
    echo "  paddlex environment found at $PADDLEX_PYTHON"
else
    echo "  ERROR: paddlex environment not found at $PADDLEX_PYTHON"
    exit 1
fi

# Step 4: Install systemd services
echo ""
echo "[4/5] Installing systemd services..."

# Copy service files
cp $BACKEND_DIR/ocr_pipeline/systemd/ocr-worker-stage2@.service $SYSTEMD_DIR/
cp $BACKEND_DIR/ocr_pipeline/systemd/ocr-worker-stage3.service $SYSTEMD_DIR/

# Reload systemd
systemctl daemon-reload
echo "  Service files installed"

# Step 5: Enable and start workers
echo ""
echo "[5/5] Starting workers..."

# Start Stage 2 workers (one per GPU)
for i in $(seq 0 $((GPU_COUNT - 1))); do
    echo "  Starting Stage 2 worker for GPU $i..."
    systemctl enable ocr-worker-stage2@$i.service
    systemctl start ocr-worker-stage2@$i.service
done

# Start Stage 3 worker on first server only (pass --stage3 flag)
if [ "$1" == "--stage3" ]; then
    echo "  Starting Stage 3 worker..."
    systemctl enable ocr-worker-stage3.service
    systemctl start ocr-worker-stage3.service
fi

echo ""
echo "=========================================="
echo "Deployment Complete!"
echo "=========================================="
echo ""
echo "Worker status:"
systemctl status ocr-worker-stage2@*.service --no-pager 2>/dev/null | grep -E "Active:|ocr-worker" || echo "  No Stage 2 workers"
systemctl status ocr-worker-stage3.service --no-pager 2>/dev/null | grep -E "Active:|ocr-worker" || echo "  No Stage 3 workers"
echo ""
echo "To check logs:"
echo "  journalctl -u ocr-worker-stage2@0 -f"
echo "  journalctl -u ocr-worker-stage3 -f"
