#!/bin/bash
# Robust birefnet batch runner — handles rembg crashes gracefully
# Processes in small batches, waits for service recovery between batches

MODEL="birefnet-general"
BATCH_SIZE=1
TOTAL=17000
PROCESSED=0
ERRORS=0
LOG="/var/log/sneakerpicks/rembg-birefnet-runner.log"

cd /var/www/sneakerpicks

echo "=== Birefnet Batch Runner started at $(date) ===" | tee -a "$LOG"
echo "Model: $MODEL | Batch: $BATCH_SIZE per run | Target: $TOTAL" | tee -a "$LOG"

wait_for_rembg() {
    local retries=0
    while [ $retries -lt 30 ]; do
        if curl -s -o /dev/null -w "%{http_code}" http://localhost:5100/api/remove 2>/dev/null | grep -q "422"; then
            return 0
        fi
        echo "  Waiting for rembg service... (attempt $((retries+1)))" | tee -a "$LOG"
        sleep 10
        retries=$((retries + 1))
    done
    echo "  rembg service not responding after 5 minutes, restarting..." | tee -a "$LOG"
    systemctl restart rembg
    sleep 30
}

while [ $PROCESSED -lt $TOTAL ]; do
    # Wait for rembg to be ready
    wait_for_rembg

    # Run single batch
    OUTPUT=$(python3 scripts/process-cutouts.py --batch $BATCH_SIZE --model "$MODEL" --priority all --delay 0 2>&1)
    echo "$OUTPUT" | tee -a "$LOG"
    
    # Check if anything was processed
    if echo "$OUTPUT" | grep -q "Processed: 0.*Errors: 0"; then
        echo "No more products to process. Done!" | tee -a "$LOG"
        break
    fi
    
    DONE=$(echo "$OUTPUT" | grep -oP 'Processed: \K[0-9]+')
    ERRS=$(echo "$OUTPUT" | grep -oP 'Errors: \K[0-9]+')
    PROCESSED=$((PROCESSED + ${DONE:-0}))
    ERRORS=$((ERRORS + ${ERRS:-0}))
    
    echo "[$(date)] Total processed: $PROCESSED | Total errors: $ERRORS" | tee -a "$LOG"
    
    # If errors, give rembg time to restart
    if [ "${ERRS:-0}" -gt 0 ]; then
        echo "  Errors detected, waiting 30s for service recovery..." | tee -a "$LOG"
        sleep 30
    else
        # Brief pause between successful batches
        sleep 3
    fi
done

echo "=== Batch Runner finished at $(date) ===" | tee -a "$LOG"
echo "Total: Processed=$PROCESSED Errors=$ERRORS" | tee -a "$LOG"