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

usage() {
  cat >&2 <<'USAGE'
Usage: $0 [--pip-version <version>] [--clean] [--local]

Options:
  --pip-version <version>   额外下载指定版本的 pip wheel（例如 25.2）。
  --clean                   清理 offline_wheels/*.whl 后重新下载。
  --local                   使用本地 python 执行下载（默认通过 docker python:3.11-slim）。
USAGE
}

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
REQUIREMENTS_FILE="$PROJECT_ROOT/requirements.txt"
WHEEL_DIR="$PROJECT_ROOT/offline_wheels"
PIP_VERSION=""
CLEAN=0
USE_LOCAL=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    --pip-version)
      [[ $# -ge 2 ]] || { usage; exit 1; }
      PIP_VERSION="$2"
      shift 2
      ;;
    --clean)
      CLEAN=1
      shift
      ;;
    --local)
      USE_LOCAL=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage
      exit 1
      ;;
  esac
 done

if [[ ! -f "$REQUIREMENTS_FILE" ]]; then
  echo "[ERROR] requirements.txt not found at $REQUIREMENTS_FILE" >&2
  exit 1
fi

mkdir -p "$WHEEL_DIR"

if [[ "$CLEAN" -eq 1 ]]; then
  echo "[INFO] Cleaning existing wheels in $WHEEL_DIR"
  find "$WHEEL_DIR" -maxdepth 1 -type f -name '*.whl' -delete
fi

run_with_python() {
  local cmd=("python" "-m" "pip" "$@")
  eval "${cmd[@]}"
}

if [[ "$USE_LOCAL" -eq 1 ]]; then
  PYTHON_BIN=${PYTHON_BIN:-python3}
  if ! command -v "$PYTHON_BIN" >/dev/null 2>&1; then
    echo "[ERROR] $PYTHON_BIN not found" >&2
    exit 1
  fi
  echo "[INFO] Using local python ($PYTHON_BIN) to download wheels"
  "$PYTHON_BIN" -m pip download -r "$REQUIREMENTS_FILE" -d "$WHEEL_DIR"
  if [[ -n "$PIP_VERSION" ]]; then
    "$PYTHON_BIN" -m pip download "pip==${PIP_VERSION}" -d "$WHEEL_DIR"
  fi
else
  if ! command -v docker >/dev/null 2>&1; then
    echo "[ERROR] docker not found; rerun with --local or安装 docker" >&2
    exit 1
  fi
  echo "[INFO] Using docker image python:3.11-slim 下载 wheel"
  docker run --rm \
    -v "$WHEEL_DIR":/wheels \
    -v "$REQUIREMENTS_FILE":/tmp/requirements.txt \
    python:3.11-slim \
    bash -c "set -euo pipefail && python -m pip install --upgrade pip && python -m pip download -r /tmp/requirements.txt -d /wheels"
  if [[ -n "$PIP_VERSION" ]]; then
    docker run --rm \
      -v "$WHEEL_DIR":/wheels \
      python:3.11-slim \
      bash -c "set -euo pipefail && python -m pip download pip==${PIP_VERSION} -d /wheels"
  fi
fi

echo "[INFO] Offline wheels prepared at $WHEEL_DIR"