#!/bin/bash

source /opt/TrainPerfBench/check_existence.sh
source /opt/TrainPerfBench/download_utils.sh

BASE_URL="https://taco-public-1251001002.cos.ap-shanghai.myqcloud.com/datasets/qwen3"
DATASET_DIR=${1:-"/data/datasets"}

prepare_redpajamaQwen_dataset()
{
  local target_path="$1"

  # 1. 检查数据集是否已完整存在
  if check_dir "${target_path}" '^[d-]' "4" && check_file "${target_path}/redpajama-sample_text_document.idx" "d07331a7940c28994ebbb2073c33f49f"; then
    echo "✅ 数据集 ${target_path} 已完整存在，跳过所有准备步骤。"
    return 0
  fi

  # 2. 如果目录存在但不完整，删除然后重新下载
  if check_dir "${target_path}"; then
    echo "⚠️ 数据集目录 ${target_path} 存在但不完整，删除目录并重新准备数据。"
    rm -rf "${target_path}"
  fi

  echo "创建数据集目录: ${target_path}"
  mkdir -p "${target_path}"

  # 3. 下载数据集
  download_target="$(dirname "${target_path}")/redpajama-qwen.tar"
  if check_file "${download_target}" "86e13dcd058f92959ce605d24349f7f6"; then
    echo "数据集压缩包已存在，跳过下载步骤。"
  else
    rm -f "${download_target}"
    safe_download "${BASE_URL}/redpajama-qwen.tar" "${download_target}.downloading"
    if [[ $? -ne 0 ]]; then
      echo "❌ 数据集下载失败，请检查网络连接或重试。"
      rm -f "${download_target}.downloading"
      return 1
    fi
    mv "${download_target}.downloading" "${download_target}"
  fi

  # 4. 解压数据集
  echo "解压数据集到: ${target_path}"
  tar -xf "${download_target}" --strip-components=1 -C "${target_path}"
  if [[ $? -ne 0 ]]; then
    echo "❌ 数据集解压失败，请检查压缩包完整性或重试。"
    rm -rf "${target_path}/*"
    return 1
  else
    mkdir "${target_path}/redpajama-sample_text_document"
    echo "✅ 数据集准备完成: ${target_path}"
    rm -f "${download_target}"
  fi
}

prepare_redpajamaQwen_dataset "${DATASET_DIR}/redpajama-qwen"
