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

UMM_ROOT="${UMM_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
SOURCE_IMAGES="${SOURCE_IMAGES:-/rczhang/rczhang/zhangrch/video_gen/gen_dataset/covt/images}"
FILES_PER_SHARD="${FILES_PER_SHARD:-12000}"
LIST_DIR="${UMM_ROOT}/data/covt/shard_lists"
SHARD_DIR="${UMM_ROOT}/data/covt/shards"

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
mkdir -p "${LIST_DIR}" "${SHARD_DIR}"

if [[ ! -d "${SOURCE_IMAGES}" ]]; then
  echo "Missing source image directory: ${SOURCE_IMAGES}" >&2
  exit 1
fi

if ! compgen -G "${LIST_DIR}/images-*.txt" >/dev/null; then
  find "${SOURCE_IMAGES}" -maxdepth 1 -type f -printf '%f\n' \
    | LC_ALL=C sort \
    | split -d -a 5 -l "${FILES_PER_SHARD}" --additional-suffix=.txt - "${LIST_DIR}/images-"
fi

for list_path in "${LIST_DIR}"/images-*.txt; do
  shard_name="$(basename "${list_path}" .txt)"
  shard_path="${SHARD_DIR}/${shard_name}.tar"
  checksum_path="${shard_path}.sha256"
  temporary_path="${shard_path}.partial"

  if [[ -s "${shard_path}" && -s "${checksum_path}" ]]; then
    echo "reuse ${shard_name}"
    continue
  fi

  echo "build ${shard_name}"
  tar -C "${SOURCE_IMAGES}" -cf "${temporary_path}" -T "${list_path}"
  mv "${temporary_path}" "${shard_path}"
  (
    cd "${SHARD_DIR}"
    sha256sum "$(basename "${shard_path}")" > "$(basename "${checksum_path}")"
  )
done

expected="$(find "${SOURCE_IMAGES}" -maxdepth 1 -type f | wc -l)"
listed="$(awk 'END {print NR}' "${LIST_DIR}"/images-*.txt)"
if [[ "${expected}" != "${listed}" ]]; then
  echo "Image list mismatch: source=${expected}, listed=${listed}" >&2
  exit 1
fi

echo "Built $(find "${SHARD_DIR}" -maxdepth 1 -name 'images-*.tar' | wc -l) image shards for ${listed} images"