#!/bin/bash
# batch_convert.sh
# 在已运行的 ROS 容器内执行，批量转换尚未转换的数据集
#
# 用法（容器内）:
#   cd /mnt && bash batch_convert.sh
#
# 它会自动检测:
#   - 哪些 .bag 文件尚未转换（通过检查 dataset/<name>/metadata.json）
#   - 跳过已转换的数据集

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
OUTPUT_DIR="${SCRIPT_DIR}/dataset"
BAG_DIR="${SCRIPT_DIR}/bags"

# 已转换的数据集列表（通过检查 metadata.json）
echo "=========================================="
echo "批量转换 UZH FPV 数据集"
echo "=========================================="

# 收集所有 bag 文件
BAG_FILES=()
while IFS= read -r -d '' f; do
    BAG_FILES+=("$f")
done < <(find "$BAG_DIR" -maxdepth 2 -name "*_davis_with_gt.bag" -print0)

if [ ${#BAG_FILES[@]} -eq 0 ]; then
    echo "❌ 未找到任何 .bag 文件"
    exit 1
fi

echo "找到 ${#BAG_FILES[@]} 个 bag 文件"

# 逐个检查并转换
CONVERTED=0
SKIPPED=0
FAILED=0

for bag_path in "${BAG_FILES[@]}"; do
    bag_name="$(basename "$bag_path")"
    dataset_name="${bag_name%_davis_with_gt.bag}"

    # 检查是否已转换
    metadata_file="${OUTPUT_DIR}/${dataset_name}/metadata.json"
    if [ -f "$metadata_file" ]; then
        echo "  ⏭️  跳过 ${dataset_name} (已转换)"
        SKIPPED=$((SKIPPED + 1))
        continue
    fi

    echo ""
    echo "  🔄 转换: ${dataset_name}"

    # 检查依赖
    if ! python3 -c "import webdataset" 2>/dev/null; then
        echo "  ⚠️   正在安装 webdataset..."
        pip3 install webdataset tqdm scipy 2>/dev/null || {
            echo "  ❌ pip install 失败，跳过 ${dataset_name}"
            FAILED=$((FAILED + 1))
            continue
        }
    fi

    # 执行转换
    if python3 "${SCRIPT_DIR}/rosbag2wds.py" \
        --bag "$bag_path" \
        --output "$OUTPUT_DIR" \
        --name "$dataset_name" \
        --shard_size 2000 \
        --width 320 --height 240; then
        echo "  ✅ ${dataset_name} 转换完成"
        CONVERTED=$((CONVERTED + 1))
    else
        echo "  ❌ ${dataset_name} 转换失败"
        FAILED=$((FAILED + 1))
    fi
done

echo ""
echo "=========================================="
echo "批量转换结束"
echo "  已转换: ${CONVERTED}"
echo "  已跳过: ${SKIPPED}"
echo "  失败:   ${FAILED}"
echo "=========================================="