#!/bin/bash
# Worker 启动脚本 — #1714 + #2073 队列分流
#
# 容器内启动三个独立进程（任一退出则整体退出由 docker restart 拉起）：
#   1. beat：celery beat 调度器，不消费任何任务，只发定时任务到 celery 默认队列
#   2. generation-worker：独占消费 generation 队列（用户实时任务，高优先级）
#   3. transcode-worker：消费 transcode + celery 默认队列（后台/清理任务）
#
# 环境变量：
#   GENERATION_CONCURRENCY        generation worker 并发（默认 2）
#   TRANSCODE_CONCURRENCY         transcode worker 并发（默认 2）
#   WORKER_MAX_TASKS_PER_CHILD    每个子进程最大任务数（默认 100）
#   WORKER_CONCURRENCY            兼容旧变量：若未显式设置 GENERATION_CONCURRENCY /
#                                 TRANSCODE_CONCURRENCY，则按比例分配（gen=ceil(total*1/2)，
#                                 trans=剩余，各至少 1）；已显式设置时忽略此变量。
#   BEAT_ENABLED                  是否在本容器内启动 beat 进程（默认 1）；
#                                 若独立 beat 容器部署设为 0。

set -e

# #2129: 幂等执行数据库迁移（watchtower 自动部署兜底）
# worker 容器独立启动，不能依赖 API 容器先跑迁移
cd /app
echo "[entrypoint-worker] Running alembic upgrade head..."
if alembic upgrade head; then
  echo "[entrypoint-worker] Migrations ok."
else
  echo "[entrypoint-worker] WARNING: alembic upgrade failed, continuing to start workers..." >&2
fi
cd - >/dev/null

MAX_TASKS="${WORKER_MAX_TASKS_PER_CHILD:-100}"

# ── 并发计算：显式 env 优先；否则从 WORKER_CONCURRENCY 按比例推导 ──
if [ -n "$GENERATION_CONCURRENCY" ]; then
  GEN_CONCURRENCY="$GENERATION_CONCURRENCY"
else
  TOTAL="${WORKER_CONCURRENCY:-4}"
  GEN_CONCURRENCY=$(( (TOTAL + 1) / 2 ))
  if [ "$GEN_CONCURRENCY" -lt 1 ]; then GEN_CONCURRENCY=1; fi
fi

if [ -n "$TRANSCODE_CONCURRENCY" ]; then
  TRANS_CONCURRENCY="$TRANSCODE_CONCURRENCY"
else
  if [ -n "$WORKER_CONCURRENCY" ] && [ -z "$GENERATION_CONCURRENCY" ]; then
    # 两个都没显式设置，按 WORKER_CONCURRENCY 分配剩余
    TOTAL="$WORKER_CONCURRENCY"
    TRANS_CONCURRENCY=$(( TOTAL - GEN_CONCURRENCY ))
    if [ "$TRANS_CONCURRENCY" -lt 1 ]; then TRANS_CONCURRENCY=1; fi
  else
    # 默认 2（#2073：独立伸缩，不再依赖 WORKER_CONCURRENCY 差值）
    TRANS_CONCURRENCY=2
  fi
fi

BEAT_ENABLED="${BEAT_ENABLED:-1}"

PIDS=()

# ── 1. Beat 调度器（独立进程，不消费任务）──
if [ "$BEAT_ENABLED" = "1" ] || [ "$BEAT_ENABLED" = "true" ]; then
  echo "Starting beat scheduler (schedule file=/tmp/celerybeat-schedule)"
  celery \
    -A worker_app.celery_app \
    beat \
    --loglevel=info \
    -s /tmp/celerybeat-schedule &
  PIDS+=($!)
fi

# ── 2. Generation worker（实时高优队列）──
echo "Starting generation worker (queue=generation, concurrency=$GEN_CONCURRENCY)"
celery \
  -A worker_app.celery_app \
  worker \
  --loglevel=info \
  -Q generation \
  "--concurrency=${GEN_CONCURRENCY}" \
  "--max-tasks-per-child=${MAX_TASKS}" \
  -n generation@%h &
PIDS+=($!)
GEN_PID=${PIDS[1]:-${PIDS[0]}}

# ── 3. Transcode worker（后台 + 清理队列）──
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 &
PIDS+=($!)
TRANS_PID=${PIDS[2]:-${PIDS[1]}}

# 任一进程退出则终止其他进程，让容器整体重启
cleanup() {
  echo "Shutting down all celery processes..."
  for pid in "${PIDS[@]}"; do
    kill -TERM "$pid" 2>/dev/null || true
  done
}
trap cleanup TERM INT

# wait -n 等待任意一个子进程退出（bash 4.3+）
# 容器镜像基础为 python:3.11-slim，bash 版本满足
wait -n "${PIDS[@]}"
EXIT_CODE=$?
echo "One celery process exited (code=$EXIT_CODE), stopping the rest..."
cleanup
exit "$EXIT_CODE"