#!/usr/bin/env bash
# Fetch raw research data. Idempotent.
set -euo pipefail
RAW=${1:-/home/user/tinychess/data/raw}
mkdir -p "$RAW"

PGN_URL="https://database.lichess.org/standard/lichess_db_standard_rated_2014-07.pgn.zst"
if [ ! -f "$RAW/games.pgn.zst" ]; then
  echo "[fetch] lichess games ..."
  curl -sL -o "$RAW/games.pgn.zst.part" "$PGN_URL"
  mv "$RAW/games.pgn.zst.part" "$RAW/games.pgn.zst"
fi
echo "[fetch] games: $(du -h "$RAW/games.pgn.zst" | cut -f1)"

for i in 0 1 2; do
  f="$RAW/puzzles-$i.parquet"
  if [ ! -f "$f" ]; then
    echo "[fetch] puzzles shard $i ..."
    curl -sL -o "$f.part" \
      "https://huggingface.co/datasets/Lichess/chess-puzzles/resolve/main/data/train-0000$i-of-00003.parquet"
    mv "$f.part" "$f"
  fi
done
echo "[fetch] done"
ls -lh "$RAW"