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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$SCRIPT_DIR"
WHEELHOUSE_DIR="$PROJECT_ROOT/wheelhouse"

mkdir -p "$WHEELHOUSE_DIR"

python3 -m venv "$PROJECT_ROOT/.wheelhouse-venv"
source "$PROJECT_ROOT/.wheelhouse-venv/bin/activate"

python -m pip install --upgrade pip setuptools wheel

pip download -d "$WHEELHOUSE_DIR" pip setuptools wheel
pip download -d "$WHEELHOUSE_DIR" "numpy<2"
pip download -d "$WHEELHOUSE_DIR" \
  --index-url https://download.pytorch.org/whl/cu126 \
  --extra-index-url https://pypi.nvidia.com \
  torch torchvision
pip download -d "$WHEELHOUSE_DIR" \
  --extra-index-url https://pypi.nvidia.com \
  "tensorflow[and-cuda]==2.16.1" "tf-keras==2.16.0"
pip download -d "$WHEELHOUSE_DIR" \
  --find-links "$WHEELHOUSE_DIR" \
  -c "$PROJECT_ROOT/constraints-wheelhouse.txt" \
  -r "$PROJECT_ROOT/requirements-native.txt"

deactivate

echo "wheelhouse_ready=$WHEELHOUSE_DIR"