#!/bin/bash
# Worker 启动脚本 — #1714 队列隔离
#
# 部署约束：worker 容器单实例（replicas=1），容器内启动两个 celery 进程：
#   1. generation-worker：独占消费 generation 队列（用户视频生成，高优先级），
#      内嵌 celery beat（-B），定时清理任务只在一个进程里跑，避免重复执行；
#   2. transcode-worker：消费 transcode + celery 默认队列（素材转码/分类/查重/
#      配音/下载等后台任务）。
# 转码队列积压时，generation 队列仍有独立 worker 立即领取视频生成任务。
#
# 环境变量：
#   WORKER_CONCURRENCY            总并发槽参考（默认 4）；生成 worker 并发默认 2，
#                                 可用 GENERATION_CONCURRENCY 覆盖
#   GENERATION_CONCURRENCY        generation worker 并发（默认 2）
#   TRANSCODE_CONCURRENCY         transcode worker 并发（默认 = WORKER_CONCURRENCY - 2，最小 1）
#   WORKER_MAX_TASKS_PER_CHILD    每个子进程最大任务数（默认 100）

set -e

CONCURRENCY="${WORKER_CONCURRENCY:-4}"
MAX_TASKS="${WORKER_MAX_TASKS_PER_CHILD:-100}"

GEN_CONCURRENCY="${GENERATION_CONCURRENCY:-2}"
if [ -z "$TRANSCODE_CONCURRENCY" ]; then
  TRANS_CONCURRENCY=$((CONCURRENCY - GEN_CONCURRENCY))
  if [ "$TRANS_CONCURRENCY" -lt 1 ]; then
    TRANS_CONCURRENCY=1
  fi
else
  TRANS_CONCURRENCY="$TRANSCODE_CONCURRENCY"
fi

echo "Starting generation worker (queue=generation, concurrency=$GEN_CONCURRENCY, beat embedded)"
celery \
  -A worker_app.celery_app \
  worker \
  --loglevel=info \
  "-B" \
  -s /tmp/celerybeat-schedule \
  -Q generation \
  "--concurrency=${GEN_CONCURRENCY}" \
  "--max-tasks-per-child=${MAX_TASKS}" \
  -n generation@%h &
GEN_PID=$!

echo "Starting transcode worker (queues=transcode,celery, concurrency=$TRANS_CONCURRENCY)"
celery \
  -A worker_app.celery_app \
  worker \
  --loglevel=info \
  -Q transcode,celery \
  "--concurrency=${TRANS_CONCURRENCY}" \
  "--max-tasks-per-child=${MAX_TASKS}" \
  -n transcode@%h &
TRANS_PID=$!

# 任一进程退出则终止另一个，让容器整体重启（restart: unless-stopped）
trap 'echo "Shutting down workers..."; kill -TERM $GEN_PID $TRANS_PID 2>/dev/null || true' TERM INT

wait -n $GEN_PID $TRANS_PID
EXIT_CODE=$?
echo "One worker exited (code=$EXIT_CODE), stopping the other..."
kill -TERM $GEN_PID $TRANS_PID 2>/dev/null || true
exit $EXIT_CODE