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

export NO_ALBUMENTATIONS_UPDATE=1

BASE_URL="http://104.238.221.200:8080/TreeOfLife-200M/dataset/parquet"

START="${1:-0}"
END="${2:-1999}"

PARQUET_DIR="./parquet"
OUT_DIR="./images"

MIN_FREE_GB=100
PROCESSES=8
THREADS=16

#################################
# Ctrl+C：无条件强杀
#################################
cleanup() {
    echo
    echo "🛑 Ctrl+C detected"
    echo "🔪 Force killing img2dataset & python workers"

    pkill -9 -f img2dataset 2>/dev/null || true
    pkill -9 -f multiprocessing 2>/dev/null || true
    pkill -9 -f python 2>/dev/null || true

    exit 130
}

trap cleanup INT TERM

mkdir -p "$PARQUET_DIR" "$OUT_DIR"

for i in $(seq -f "%04g" "$START" "$END"); do
    part="part_${i}"
    parquet="${PARQUET_DIR}/${part}.parquet"
    out_path="${OUT_DIR}/${part}"
    url="${BASE_URL}/${part}.parquet"

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

    if [ ! -f "$parquet" ]; then
        wget -c "$url" -O "$parquet" || {
            echo "❌ Download failed, skip"
            rm -f "$parquet"
            continue
        }
    fi

    free_gb=$(df -BG "$OUT_DIR" | awk 'NR==2 {gsub("G","",$4); print $4}')
    if [ "$free_gb" -lt "$MIN_FREE_GB" ]; then
        echo "🛑 Disk space low, stop"
        exit 0
    fi

    if [ -d "$out_path" ] && [ "$(ls -A "$out_path" 2>/dev/null)" ]; then
        echo "⏭️  Skip ${part}"
        continue
    fi

    mkdir -p "$out_path"
    echo "🚀 Start img2dataset ${part}"

    img2dataset \
        --url_list "$parquet" \
        --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"
        ]'

    echo "✅ EXIT ${part}"
done
