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

PY_MODULES_DIR="$(dirname "$0")/../py_modules"
REQUIREMENTS="$(dirname "$0")/../requirements.txt"
TMPDIR=$(mktemp -d)
trap 'rm -rf "$TMPDIR"' EXIT

# Detect Python version (default: 3.13)
PYTHON_VERSION="${PYTHON_VERSION:-3.11}"
ABI="cp${PYTHON_VERSION//./}"
PLATFORMS="manylinux2014_x86_64,manylinux_2_17_x86_64,linux_x86_64"

echo "==> Bundling Python dependencies for Python ${PYTHON_VERSION}"
echo "==> Requirements: ${REQUIREMENTS}"
echo "==> Output: ${PY_MODULES_DIR}"

# Clean py_modules
rm -rf "${PY_MODULES_DIR:?}"
mkdir -p "${PY_MODULES_DIR}"

# Download all dependencies (with transitive deps)
echo "==> Downloading wheels..."
pip download \
    --python-version "${PYTHON_VERSION}" \
    --abi "${ABI}" \
    --only-binary=:all: \
    --platform manylinux2014_x86_64 \
    --platform manylinux_2_17_x86_64 \
    --platform linux_x86_64 \
    -r "${REQUIREMENTS}" \
    -d "${TMPDIR}"

echo "==> Downloaded $(ls "${TMPDIR}"/*.whl 2>/dev/null | wc -l) wheels"

# Extract into py_modules
echo "==> Extracting wheels..."
for whl in "${TMPDIR}"/*.whl; do
    unzip -qo "$whl" -d "${PY_MODULES_DIR}"
done

# Clean up dist-info and RECORD
echo "==> Cleaning up..."
find "${PY_MODULES_DIR}" -name "*.dist-info" -type d -exec rm -rf {} + 2>/dev/null || true
find "${PY_MODULES_DIR}" -name "*.egg-info" -type d -exec rm -rf {} + 2>/dev/null || true
find "${PY_MODULES_DIR}" -name "RECORD" -delete 2>/dev/null || true
find "${PY_MODULES_DIR}" -name "*.pyc" -delete 2>/dev/null || true

# Show results
echo "==> Bundled packages:"
ls -d "${PY_MODULES_DIR}"/*/ 2>/dev/null | xargs -I{} basename {} | sort
echo ""
SO_COUNT=$(find "${PY_MODULES_DIR}" -name "*.so" | wc -l)
echo "==> ${SO_COUNT} native extensions (.so files)"
echo "==> Done!"