#!/usr/bin/env bash
set -euo pipefail

HOST=$(uname -n)
case "$HOST" in
  *neon*)  DEFAULT_WORKERS=12 ;;
  *)       DEFAULT_WORKERS=4 ;;
esac

WORKERS="$DEFAULT_WORKERS"
ROLE="worker"
BUILD=false
PLATFORM=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --workers) WORKERS="$2"; shift 2 ;;
    --role)    ROLE="$2";    shift 2 ;;
    --build)   BUILD=true;   shift ;;
    --platform) PLATFORM="$2 --platform $2"; shift 2 ;;
    *) echo "Usage: $0 [--workers N] [--role seed|worker|all] [--build] [--platform linux/arm64|linux/amd64]"
       echo "  Default workers: $DEFAULT_WORKERS"; exit 1 ;;
  esac
done

DIR="$(cd "$(dirname "$0")" && pwd)"
IMAGE="lkb-worker:latest"

if [ "$BUILD" = true ] || ! podman image exists "$IMAGE" 2>/dev/null; then
  echo "[build] $IMAGE ..."
  podman build -t "$IMAGE" -f "$DIR/Dockerfile.lkb" "$DIR"
  echo "[build] done"
fi

run_seed() {
  podman rm -f lkb-seed 2>/dev/null || true
  podman run -d --rm --name lkb-seed --network host "$IMAGE" seed
  echo "[seed] started. Logs: podman logs -f lkb-seed"
}

run_workers() {
  echo "[worker] launching $WORKERS containers ..."
  for i in $(seq 1 "$WORKERS"); do
    podman rm -f "lkb-worker-$i" 2>/dev/null || true
    podman run -d --rm --name "lkb-worker-$i" --network host "$IMAGE" worker
  done
  echo "[worker] $WORKERS launched"
}

case "$ROLE" in
  seed)   run_seed ;;
  worker) run_workers ;;
  all)    run_seed; sleep 3; run_workers ;;
esac

echo "=== Running ==="
podman ps --format "table {{.Names}} {{.Status}}" | grep lkb-
echo "=== DB queue ==="
PGPASSWORD=kaleidoscope psql -h 100.79.14.103 -U kaleidoscope -d research_stack \
  -c "SELECT status, count(*) FROM lkb.compute_queue GROUP BY status ORDER BY status" 2>/dev/null || true