#!/usr/bin/env bash
# prepare.sh: everything that needs the network happens here, OUTSIDE the probe.
#
#   1. fetch_wheels   pip download onnxruntime 1.30.0 / 1.29.0 / 1.28.0 (+ the same onnx and numpy pins)
#   2. bundle_strace  copy this host's strace, its libraries and its loader into $WORK/strace
#   3. build_venvs    create one venv per version INSIDE a --network none container, offline (pip --no-index)
#
# Environment (all optional except PROBE_PYTHON_DIR when the image has no python3):
#   PROBE_WORK        scratch dir for wheels, venvs, strace (default: ./work next to this script)
#   PROBE_PYTHON_DIR  a CPython 3.12 install on the host, mounted read-only at /probe/python in the container
#   PROBE_IMAGE       container image (default debian:12-slim)
set -euo pipefail
KIT=$(cd "$(dirname "$0")" && pwd)
WORK=${PROBE_WORK:-$KIT/work}
IMAGE=${PROBE_IMAGE:-debian:12-slim}
VERSIONS="1.30.0 1.29.0 1.28.0"
PINS="onnx==1.23.0 numpy==2.5.3"
HOST_PY=${PROBE_PYTHON_DIR:+$PROBE_PYTHON_DIR/bin/python3}
HOST_PY=${HOST_PY:-python3}
CPY=${PROBE_PYTHON_DIR:+/probe/python/bin/python3}
CPY=${CPY:-python3}

venv_name() { echo "v$(echo "$1" | tr -d .)"; }   # 1.30.0 -> v1300

fetch_wheels() {
  for v in $VERSIONS; do
    PIP_NO_CACHE_DIR=1 PIP_DISABLE_PIP_VERSION_CHECK=1 "$HOST_PY" -m pip download -q --only-binary=:all: --python-version 3.12 \
      --platform manylinux_2_28_x86_64 --platform manylinux_2_27_x86_64 \
      --platform manylinux_2_17_x86_64 --platform manylinux2014_x86_64 \
      -d "$WORK/wheels/ort-$v" "onnxruntime==$v" $PINS
  done
  (cd "$WORK/wheels" && sha256sum ./*/*.whl > SHA256SUMS)
}

bundle_strace() {
  mkdir -p "$WORK/strace"
  cp -L "$(command -v strace)" /lib64/ld-linux-x86-64.so.2 "$WORK/strace/"
  ldd "$(command -v strace)" | awk '/=>/ {print $3}' | xargs -I{} cp -L {} "$WORK/strace/"
}

build_venvs() {
  mkdir -p "$WORK/venvs"
  local script=""
  for v in $VERSIONS; do
    d=/probe/work/venvs/$(venv_name "$v")
    script+="$CPY -m venv $d && $d/bin/pip install -q --no-index --find-links /probe/work/wheels/ort-$v onnxruntime==$v $PINS && "
  done
  docker run --rm --network none --user "$(id -u):$(id -g)" --cap-drop ALL \
    --security-opt no-new-privileges --read-only --tmpfs /tmp:rw,mode=1777 \
    -e HOME=/tmp -e PIP_NO_CACHE_DIR=1 -e PIP_DISABLE_PIP_VERSION_CHECK=1 \
    -v "$WORK:/probe/work:rw" ${PROBE_PYTHON_DIR:+-v "$PROBE_PYTHON_DIR:/probe/python:ro"} \
    "$IMAGE" sh -c "${script}true"
}

# no argument: all three steps in order; or name one step, e.g. `prepare.sh build_venvs`
if [ $# -eq 0 ]; then fetch_wheels; bundle_strace; build_venvs; else "$@"; fi