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

: "${HF_NAMESPACE:?HF_NAMESPACE must be set, e.g. username}"
: "${ARTIFACT_REPO:?ARTIFACT_REPO must be set, e.g. username/assaylens-artifacts}"

JOB_SCRIPT="${1:?Usage: scripts/run_hf_job.sh scripts/01_profile_eve.py [flavor] [script args...]}"
FLAVOR="${2:-cpu-upgrade}"
if [[ "$#" -ge 2 ]]; then
  shift 2
else
  shift 1
fi
SCRIPT_ARGS="$*"
TIMEOUT="${TIMEOUT:-4h}"
CODE_REPO="${ASSAYLENS_CODE_REPO:-${HF_NAMESPACE}/assaylens-code-snapshot}"
case "$JOB_SCRIPT" in
  *06_train_assaylens.py)
    JOB_DEPS="huggingface_hub>=0.36.0 datasets>=3.0.0 pyarrow>=15.0.0 pandas>=2.2.0 numpy>=1.26.0 scikit-learn>=1.4.0 scipy>=1.12.0 joblib>=1.3.0 torch>=2.3.0"
    ;;
  *08_next_assay_policy.py)
    JOB_DEPS="huggingface_hub>=0.36.0 datasets>=3.0.0 pyarrow>=15.0.0 pandas>=2.2.0 numpy>=1.26.0 scikit-learn>=1.4.0 scipy>=1.12.0 joblib>=1.3.0 rdkit>=2023.9.0 openai>=1.109.0"
    ;;
  *11_gradio_demo.py)
    JOB_DEPS="huggingface_hub>=0.36.0 pandas>=2.2.0 numpy>=1.26.0 scikit-learn>=1.4.0 joblib>=1.3.0 rdkit>=2023.9.0 openai>=1.109.0 gradio>=5.0.0"
    ;;
  *03_build_features.py|*04_make_splits.py|*05_train_baselines.py|*07_evaluate.py)
    JOB_DEPS="huggingface_hub>=0.36.0 datasets>=3.0.0 pyarrow>=15.0.0 pandas>=2.2.0 numpy>=1.26.0 scikit-learn>=1.4.0 scipy>=1.12.0 joblib>=1.3.0 rdkit>=2023.9.0"
    ;;
  *10_upload_artifacts.py)
    JOB_DEPS="huggingface_hub>=0.36.0"
    ;;
  *)
    JOB_DEPS="huggingface_hub>=0.36.0 datasets>=3.0.0 pyarrow>=15.0.0 pandas>=2.2.0 numpy>=1.26.0 pyyaml>=6.0.0"
    ;;
esac
SNAPSHOT_DIR="$(mktemp -d)"
cleanup() {
  rm -rf "$SNAPSHOT_DIR"
}
trap cleanup EXIT

tar \
  --exclude=.git \
  --exclude=.venv \
  --exclude=.pytest_cache \
  --exclude='__pycache__' \
  --exclude='*.pyc' \
  --exclude='*.Zone.Identifier' \
  --exclude='*:Zone.Identifier' \
  --exclude=./data \
  --exclude=./outputs \
  --exclude=./coverage \
  --exclude=./reports \
  -cf - . | tar -xf - -C "$SNAPSHOT_DIR"

hf repos create "$CODE_REPO" --type dataset --private --exist-ok >/dev/null
hf upload "$CODE_REPO" "$SNAPSHOT_DIR" \
  --type dataset \
  --delete "*" \
  --commit-message "Upload AssayLens code snapshot" >/dev/null

hf jobs run \
  --flavor "$FLAVOR" \
  --timeout "$TIMEOUT" \
  --secrets HF_TOKEN \
  --env ARTIFACT_REPO="$ARTIFACT_REPO" \
  --env HF_NAMESPACE="$HF_NAMESPACE" \
  --env HF_HOME="/tmp/hf" \
  --env HF_HUB_CACHE="/tmp/hf/hub" \
  --env HF_DATASETS_CACHE="/tmp/hf/datasets" \
  --env ASSAYLENS_CODE_REPO="$CODE_REPO" \
  --env ASSAYLENS_JOB_DEPS="$JOB_DEPS" \
  --env ASSAYLENS_AUTO_UPLOAD="1" \
  python:3.11 \
  -- \
  bash -lc "
    set -euo pipefail
    python -m pip install -U pip &&
    python -m pip install \$ASSAYLENS_JOB_DEPS &&
    mkdir -p /workspace/assaylens &&
    hf download \"\$ASSAYLENS_CODE_REPO\" --repo-type dataset --local-dir /workspace/assaylens &&
    cd /workspace/assaylens &&
    python -m pip install -e . --no-deps &&
    python ${JOB_SCRIPT} ${SCRIPT_ARGS}
  "