#!/usr/bin/env bash
# Extract all 12 EasyCom sessions with three independent workers per GPU.
#
# Sessions are assigned round-robin across four GPUs:
#   GPU 0: Session 1, 5, 9
#   GPU 1: Session 2, 6, 10
#   GPU 2: Session 3, 7, 11
#   GPU 3: Session 4, 8, 12
#
# Worker launches are deliberately serialized. The launcher waits until one
# worker has loaded Qwen and produced its first 0.2-second state before starting
# the next worker, preventing all model checkpoints from loading at once.
set -uo pipefail

PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}"
VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}"
TARGET_ROOT="${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}"
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"

GPUS="${GPUS:-0 1 2 3}"
SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}"
PROCESSES_PER_GPU="${PROCESSES_PER_GPU:-3}"
MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
VIDEO_FPS="${VIDEO_FPS:-2}"
FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}"
DTYPE="${DTYPE:-bfloat16}"
OMP_THREADS_PER_PROCESS="${OMP_THREADS_PER_PROCESS:-4}"

# Maximum time allowed for one worker to index media, load Qwen, and produce
# its first state. A small gap after readiness further separates CUDA startup.
START_TIMEOUT_SECONDS="${START_TIMEOUT_SECONDS:-1800}"
START_GAP_SECONDS="${START_GAP_SECONDS:-10}"
POLL_SECONDS="${POLL_SECONDS:-5}"

read -r -a GPU_ARRAY <<<"${GPUS}"
read -r -a SESSION_ARRAY <<<"${SESSIONS}"

if [[ "${#GPU_ARRAY[@]}" -eq 0 ]]; then
  echo "No GPUs supplied through GPUS" >&2
  exit 2
fi
if [[ "${PROCESSES_PER_GPU}" -le 0 ]]; then
  echo "PROCESSES_PER_GPU must be positive" >&2
  exit 2
fi

capacity=$((${#GPU_ARRAY[@]} * PROCESSES_PER_GPU))
if [[ "${#SESSION_ARRAY[@]}" -gt "${capacity}" ]]; then
  echo "${#SESSION_ARRAY[@]} sessions exceed capacity ${capacity}" >&2
  exit 2
fi

mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
cd "${PROJECT_ROOT}"

declare -a CHILD_PIDS=()
declare -a CHILD_NAMES=()
declare -a CHILD_PIDFILES=()
declare -a EXTERNAL_PIDS=()
declare -a EXTERNAL_NAMES=()
declare -a EXTERNAL_OUTPUTS=()
declare -a EXTERNAL_PIDFILES=()
declare -A GPU_JOB_COUNTS=()

for gpu in "${GPU_ARRAY[@]}"; do
  GPU_JOB_COUNTS["${gpu}"]=0
done

wait_until_ready() {
  local pid="$1"
  local log="$2"
  local session="$3"
  local started_at="${SECONDS}"

  while true; do
    # tqdm writes progress updates with carriage returns; grep -a handles the
    # resulting log as text. This marker appears after the first model forward.
    if grep -aq "media=0.2s" "${log}"; then
      return 0
    fi
    if grep -aq "Traceback (most recent call last)" "${log}"; then
      echo "${session} failed during startup; inspect ${log}" >&2
      return 1
    fi
    if ! kill -0 "${pid}" 2>/dev/null; then
      echo "${session} exited before producing its first state; inspect ${log}" >&2
      return 1
    fi
    if ((SECONDS - started_at >= START_TIMEOUT_SECONDS)); then
      echo "${session} startup timed out after ${START_TIMEOUT_SECONDS}s" >&2
      return 1
    fi
    sleep "${POLL_SECONDS}"
  done
}

startup_status=0
for index in "${!SESSION_ARRAY[@]}"; do
  sid="${SESSION_ARRAY[${index}]}"
  gpu="${GPU_ARRAY[$((index % ${#GPU_ARRAY[@]}))]}"
  session="Session_${sid}"
  session_dir="${VIDEO_ROOT}/${session}"
  target="${TARGET_ROOT}/${session}.npy"
  output="${FEATURE_ROOT}/${session}.pt"
  log="${LOG_ROOT}/${session}.log"
  pidfile="${LOG_ROOT}/${session}.pid"

  GPU_JOB_COUNTS["${gpu}"]=$((GPU_JOB_COUNTS["${gpu}"] + 1))
  if [[ "${GPU_JOB_COUNTS[${gpu}]}" -gt "${PROCESSES_PER_GPU}" ]]; then
    echo "GPU ${gpu} was assigned too many workers" >&2
    startup_status=1
    break
  fi
  if [[ ! -d "${session_dir}" ]]; then
    echo "Missing session directory: ${session_dir}" >&2
    startup_status=1
    break
  fi
  if [[ ! -s "${target}" ]]; then
    echo "Missing target: ${target}" >&2
    startup_status=1
    break
  fi
  if [[ -s "${output}" ]]; then
    echo "Skipping completed ${session}: ${output}"
    continue
  fi

  # Also detect a worker launched manually before this manager. This is useful
  # when Session_1 is already running on GPU 0. Its round-robin slot still
  # counts toward that GPU's three-worker limit.
  running_pid=""
  if command -v pgrep >/dev/null 2>&1; then
    running_pid=$(pgrep -f \
      "streaming_turn_model_directly.py.*${session_dir}" | head -n 1 || true)
  fi
  if [[ -n "${running_pid}" ]] && kill -0 "${running_pid}" 2>/dev/null; then
    echo "Skipping already-running ${session}: PID=${running_pid}"
    echo "${running_pid}" >"${pidfile}"
    EXTERNAL_PIDS+=("${running_pid}")
    EXTERNAL_NAMES+=("${session}")
    EXTERNAL_OUTPUTS+=("${output}")
    EXTERNAL_PIDFILES+=("${pidfile}")
    if [[ -f "${log}" ]]; then
      echo "  Confirming the existing worker reached its first state..."
      if ! wait_until_ready "${running_pid}" "${log}" "${session}"; then
        startup_status=1
        break
      fi
    else
      echo "  WARNING: no log is available to confirm startup; waiting " \
           "${START_GAP_SECONDS}s before the next load" >&2
    fi
    sleep "${START_GAP_SECONDS}"
    continue
  fi
  if [[ -s "${pidfile}" ]]; then
    existing_pid=$(<"${pidfile}")
    if kill -0 "${existing_pid}" 2>/dev/null; then
      echo "Skipping already-running ${session}: PID=${existing_pid}"
      continue
    fi
  fi

  : >"${log}"
  echo "Starting ${session} on physical GPU ${gpu} " \
       "(${GPU_JOB_COUNTS[${gpu}]}/${PROCESSES_PER_GPU} workers)"

  env \
    CUDA_VISIBLE_DEVICES="${gpu}" \
    PYTHONUNBUFFERED=1 \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
    OMP_NUM_THREADS="${OMP_THREADS_PER_PROCESS}" \
    python streaming_turn_model_directly.py \
      --model-id "${MODEL}" \
      --easycom-session-dir "${session_dir}" \
      --target "${target}" \
      --use-audio-in-video \
      --mode open-stream-replay \
      --memory-seconds "${MEMORY_SECONDS}" \
      --video-fps "${VIDEO_FPS}" \
      --frame-cache-size "${FRAME_CACHE_SIZE}" \
      --dtype "${DTYPE}" \
      --flash-attention-2 \
      --device cuda:0 \
      --output "${output}" \
      >"${log}" 2>&1 &

  pid=$!
  echo "${pid}" >"${pidfile}"
  CHILD_PIDS+=("${pid}")
  CHILD_NAMES+=("${session}")
  CHILD_PIDFILES+=("${pidfile}")
  echo "  PID=${pid}; waiting for first state before next model load..."

  if ! wait_until_ready "${pid}" "${log}" "${session}"; then
    startup_status=1
    tail -n 40 "${log}" >&2
    break
  fi
  echo "  ${session} is running; next worker starts in ${START_GAP_SECONDS}s"
  sleep "${START_GAP_SECONDS}"
done

echo "Launcher finished starting workers; waiting for child processes."
final_status="${startup_status}"
for index in "${!CHILD_PIDS[@]}"; do
  pid="${CHILD_PIDS[${index}]}"
  session="${CHILD_NAMES[${index}]}"
  pidfile="${CHILD_PIDFILES[${index}]}"
  if wait "${pid}"; then
    echo "Completed ${session}"
  else
    echo "Failed ${session}; inspect ${LOG_ROOT}/${session}.log" >&2
    final_status=1
  fi
  rm -f "${pidfile}"
done

# Non-child workers cannot be handled by bash's wait builtin. Poll any workers
# that were already running when this launcher started, then validate that they
# produced their expected output before reporting global success.
for index in "${!EXTERNAL_PIDS[@]}"; do
  pid="${EXTERNAL_PIDS[${index}]}"
  session="${EXTERNAL_NAMES[${index}]}"
  output="${EXTERNAL_OUTPUTS[${index}]}"
  pidfile="${EXTERNAL_PIDFILES[${index}]}"
  while kill -0 "${pid}" 2>/dev/null; do
    sleep 30
  done
  if [[ -s "${output}" ]]; then
    echo "Completed pre-existing ${session}"
  else
    echo "Pre-existing ${session} exited without ${output}" >&2
    final_status=1
  fi
  rm -f "${pidfile}"
done

if [[ "${final_status}" -ne 0 ]]; then
  echo "One or more workers failed, or worker startup was interrupted." >&2
  exit 1
fi

echo "All selected EasyCom sessions completed: ${FEATURE_ROOT}"