#!/bin/sh
set -e

REPO="${ATLAS_DATA_REPO:-michxu/nvidia-nemotron-v3-atlas-data}"
DATASET_FILE="${ATLAS_DATASET_FILE:-dataset_search.parquet}"
PCA_FILE="${ATLAS_PCA_FILE:-pca_model.json}"
DETAILS_FILE="${ATLAS_DETAILS_FILE:-dataset_details.parquet}"
ACCESS_TOKEN=""
ACCESS_TOKEN_SOURCE="none"
if [ -n "${michxu_HF_TOKEN:-}" ]; then
  ACCESS_TOKEN="$michxu_HF_TOKEN"
  ACCESS_TOKEN_SOURCE="michxu_HF_TOKEN"
elif [ -n "${MICHXU_HF_TOKEN:-}" ]; then
  ACCESS_TOKEN="$MICHXU_HF_TOKEN"
  ACCESS_TOKEN_SOURCE="MICHXU_HF_TOKEN"
elif [ -n "${HF_TOKEN:-}" ]; then
  ACCESS_TOKEN="$HF_TOKEN"
  ACCESS_TOKEN_SOURCE="HF_TOKEN"
fi
DATA_DIR="/usr/share/nginx/html/data"
mkdir -p "$DATA_DIR"

download_file() {
  remote_file="$1"
  output_file="$2"
  url="https://huggingface.co/datasets/${REPO}/resolve/main/${remote_file}"

  echo "Downloading ${remote_file} from ${REPO} ..."
  if [ -n "$ACCESS_TOKEN" ]; then
    curl -L -f --retry 3 -H "Authorization: Bearer ${ACCESS_TOKEN}" "$url" -o "$output_file" || {
      status="$?"
      echo "ERROR: Failed to download ${remote_file} from ${REPO} with access token from ${ACCESS_TOKEN_SOURCE}." >&2
      echo "       Check that the token can read the dataset repo and that the file exists on main." >&2
      echo "       URL: ${url}" >&2
      exit "$status"
    }
  else
    curl -L -f --retry 3 "$url" -o "$output_file" || {
      status="$?"
      echo "ERROR: Failed to download ${remote_file} from ${REPO} without HF_TOKEN." >&2
      echo "       If the dataset is private, add a Space secret named HF_TOKEN with read access." >&2
      echo "       URL: ${url}" >&2
      exit "$status"
    }
  fi
}

download_file "$DATASET_FILE" "$DATA_DIR/dataset.parquet"
echo "  dataset: $(du -h $DATA_DIR/dataset.parquet | cut -f1)"

download_file "$PCA_FILE" "$DATA_DIR/pca_model.json"
echo "  pca_model: $(du -h $DATA_DIR/pca_model.json | cut -f1)"

download_file "$DETAILS_FILE" "$DATA_DIR/details.parquet"
echo "  details: $(du -h $DATA_DIR/details.parquet | cut -f1)"

echo "Starting nginx ..."
exec nginx -g 'daemon off;'