#!/usr/bin/env bash

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
DATASET_ID="pipecat-ai/smart-turn-data-v3.2-train"
DATASET_REVISION="${DATASET_REVISION:-e564e2ac567f774d1880aa1db6ce97afb8c519b7}"
TARGET_DIR="${PROJECT_ROOT}/data/raw/smart-turn-data-v3.2-train"
EXPECTED_PARQUET_SHARDS=83

if [[ -f "${PROJECT_ROOT}/.env" ]]; then
  set -a
  # shellcheck disable=SC1091
  source "${PROJECT_ROOT}/.env"
  set +a
fi

export HF_TOKEN="${HF_TOKEN:-${hf_token:-}}"
if [[ -z "${HF_TOKEN}" ]]; then
  echo "Missing Hugging Face token. Set HF_TOKEN or hf_token in .env." >&2
  exit 1
fi
unset hf_token

export UV_CACHE_DIR="${UV_CACHE_DIR:-${PROJECT_ROOT}/.cache/uv}"
export UV_TOOL_DIR="${UV_TOOL_DIR:-${PROJECT_ROOT}/.cache/uv-tools}"
export UV_TOOL_BIN_DIR="${UV_TOOL_BIN_DIR:-${PROJECT_ROOT}/.cache/uv-bin}"
export HF_XET_HIGH_PERFORMANCE="${HF_XET_HIGH_PERFORMANCE:-1}"

mkdir -p "${TARGET_DIR}"

if [[ -n "${HF_CLI_BIN:-}" ]]; then
  if [[ ! -x "${HF_CLI_BIN}" ]]; then
    echo "HF_CLI_BIN is not executable: ${HF_CLI_BIN}" >&2
    exit 1
  fi
  hf_command=("${HF_CLI_BIN}")
elif command -v hf >/dev/null 2>&1; then
  hf_command=("$(command -v hf)")
else
  hf_command=(uvx --from 'huggingface_hub>=0.27,<2' hf)
fi

"${hf_command[@]}" download \
  "${DATASET_ID}" \
  --repo-type dataset \
  --revision "${DATASET_REVISION}" \
  --local-dir "${TARGET_DIR}" \
  "$@"

actual_shards="$(find "${TARGET_DIR}/data" -type f -name '*.parquet' 2>/dev/null | wc -l | tr -d ' ')"
if [[ "${actual_shards}" != "${EXPECTED_PARQUET_SHARDS}" ]]; then
  echo "Incomplete snapshot: expected ${EXPECTED_PARQUET_SHARDS} Parquet shards, found ${actual_shards}." >&2
  echo "Re-run this command to resume safely." >&2
  exit 1
fi

echo "Verified ${actual_shards} Parquet shards at revision ${DATASET_REVISION}."