#!/usr/bin/env bash
set -euo pipefail

#################################
# 全局环境（静默无关 warning）
#################################
export NO_ALBUMENTATIONS_UPDATE=1

#################################
# 配置区（按你的要求，不改并发）
#################################
PARQUET_DIR="./parquet"
OUT_DIR="./images"

MIN_FREE_GB=500

# 你指定的并发（不改）
PROCESSES=32
THREADS=64

# 每个 part 的最大运行时间（自己调）
# 例如：6h、8h、12h
TIME_LIMIT="6h"

#################################
# 初始化
#################################
mkdir -p "$OUT_DIR"

echo "📂 Parquet dir : $PARQUET_DIR"
echo "📂 Output dir  : $OUT_DIR"
echo "🧵 processes=$PROCESSES threads=$THREADS"
echo "⏱️  per-part time limit=$TIME_LIMIT"
echo "💾 min free space=${MIN_FREE_GB}GB"
echo

#################################
# 主循环
#################################
for p in ${PARQUET_DIR}/part_*.parquet; do
    part_name=$(basename "$p" .parquet)
    out_path="${OUT_DIR}/${part_name}"

    echo "===================================================="
    echo "➡️  ENTER ${part_name}"

    # 1️⃣ 校验 schema（防止误喂）
    python - << EOF
import pyarrow.parquet as pq
pf = pq.ParquetFile("${p}")
if "source_url" not in pf.schema.names:
    raise SystemExit("❌ source_url missing in ${p}")
EOF

    # 2️⃣ 检查磁盘剩余空间
    free_gb=$(df -BG "$OUT_DIR" | tail -1 | awk '{print $4}' | sed 's/G//')
    if [ "$free_gb" -lt "$MIN_FREE_GB" ]; then
        echo "🛑 Disk free space ${free_gb}GB < ${MIN_FREE_GB}GB"
        echo "🛑 Stop at ${part_name}"
        exit 0
    fi

    # 3️⃣ 已完成则跳过（可续跑）
    if [ -d "$out_path" ] && [ "$(ls -A "$out_path" 2>/dev/null)" ]; then
        echo "⏭️  Skip ${part_name} (already exists)"
        continue
    fi

    mkdir -p "$out_path"

    echo "🚀 Start ${part_name} (free ${free_gb}GB)"

    #################################
    # 4️⃣ img2dataset（强制兜底）
    #################################
    if ! timeout --signal=SIGTERM "$TIME_LIMIT" img2dataset \
      --url_list "$p" \
      --input_format parquet \
      --url_col source_url \
      --output_folder "$out_path" \
      --output_format webdataset \
      --processes_count "$PROCESSES" \
      --thread_count "$THREADS" \
      --image_size 0 \
      --resize_mode no \
      --disable_all_reencoding True \
      --skip_reencode True \
      --resize_only_if_bigger False \
      --retries 1 \
      --ignore_ssl_certificate \
      --save_additional_columns '[
        "uuid",
        "kingdom",
        "phylum",
        "class",
        "order",
        "family",
        "genus",
        "species",
        "scientific_name",
        "common",
        "data_source",
        "publisher",
        "basis_of_record",
        "img_type",
        "source_id",
        "shard_filename",
        "shard_file_path",
        "base_dataset_file_path"
      ]' \
 
    then
        echo "⚠️  ${part_name} exited (error or timeout), continue next"
    else
        echo "✅ EXIT ${part_name}"
    fi

    echo "➡️  NEXT PART"
    echo
done

echo "🎉 All parquet parts processed (with high-concurrency mode)"
