#!/usr/bin/env bash
source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/block_white_box_common.sh"
if [[ "${RESUME:-false}" == true ]]; then
  [[ -f "$DATASET_ROOT/meta/info.json" ]] || { echo "找不到可续录的数据：$DATASET_ROOT" >&2; exit 1; }
else
  [[ ! -e "$DATASET_ROOT" ]] || { echo "目录已存在；续录请设置 RESUME=true，或指定新的 DATASET_ROOT / DATASET_ID。" >&2; exit 1; }
fi
check_task_devices
if [[ "${DRY_RUN:-0}" != 1 ]]; then
  [[ -r "$LEADER_PORT" && -w "$LEADER_PORT" ]] || { echo "主臂未连接或缺少权限：$LEADER_PORT" >&2; exit 1; }
fi
printf '\n========== 积木放入白盒：准备采集 ==========\n'
echo "本次 ${NUM_EPISODES:-5} 条；每条最多 ${EPISODE_SECONDS:-30} 秒；复位最多 ${RESET_SECONDS:-30} 秒。"
echo "绿色横幅 = 正在录制；黄色横幅 = 复位 / 重录。"
echo "→ 完成当前阶段    ← 重录当前条    Esc 结束采集"
echo "请在此终端操作方向键，并留意阶段横幅。"
if [[ "${DRY_RUN:-0}" != 1 ]]; then
  read -r -p "摆好积木、盒子和主从臂后，按 Enter 连接设备并开始录制：" _ready
fi
run_task python "$TASK_SCRIPT_DIR/record_with_prompts.py" "${ROBOT_ARGS[@]}" \
  --teleop.type=so101_leader "--teleop.port=$LEADER_PORT" --teleop.id=zihao_leader_arm \
  "--dataset.repo_id=$DATASET_ID" "--dataset.root=$DATASET_ROOT" \
  "--dataset.single_task=$TASK_TEXT" --dataset.no_stamp=true \
  --dataset.fps=25 "--dataset.num_episodes=${NUM_EPISODES:-5}" \
  "--dataset.episode_time_s=${EPISODE_SECONDS:-30}" "--dataset.reset_time_s=${RESET_SECONDS:-30}" \
  --dataset.push_to_hub=false --dataset.rgb_encoder.vcodec=h264 \
  --display_data=true "--resume=${RESUME:-false}" "$@"