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

ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
PYTHON_BIN=python3.12
VENV_DIR="$ROOT/.venv"
DRY_RUN=0
DEPENDENCIES_ONLY=0
LOCK_FILE="$ROOT/requirements/bootstrap.lock"

usage() {
  printf '%s\n' \
    "Usage: scripts/bootstrap-rocm.sh [--python PATH] [--venv PATH] [--dependencies-only] [--dry-run]" \
    "" \
    "Creates an isolated Python 3.12 environment with AMD's verified ROCm 7.2.1 wheels." \
    "It never modifies the system Python or the host ROCm installation."
}

while (($#)); do
  case "$1" in
    --python)
      PYTHON_BIN=${2:?--python requires a path}
      shift 2
      ;;
    --venv)
      VENV_DIR=${2:?--venv requires a path}
      shift 2
      ;;
    --dependencies-only)
      DEPENDENCIES_ONLY=1
      shift
      ;;
    -n|--dry-run)
      DRY_RUN=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      printf 'error: unknown argument: %s\n' "$1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

ROCM_RELEASE=7.2.1
CACHE_BASE=${XDG_CACHE_HOME:-${HOME}/.cache}/unlimited-ocr-rdna4
WHEEL_DIR="$CACHE_BASE/wheels/rocm-$ROCM_RELEASE"
TASK_TMP="$CACHE_BASE/tmp"

TORCH_FILE='torch-2.9.1+rocm7.2.1.lw.gitff65f5bc-cp312-cp312-linux_x86_64.whl'
TORCH_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/torch-2.9.1%2Brocm7.2.1.lw.gitff65f5bc-cp312-cp312-linux_x86_64.whl'
TORCH_SHA='fb45ace0a27e9f0d0e3c4c6efd8932162743f8376f2aa4752a4d31ef5a1bd3d7'

VISION_FILE='torchvision-0.24.0+rocm7.2.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl'
VISION_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/torchvision-0.24.0%2Brocm7.2.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl'
VISION_SHA='d5fca8cda173235a3b7434baeebe04c3ebffec3c6fc191e79aa8aa300633f2c9'

TRITON_FILE='triton-3.5.1+rocm7.2.1.gita272dfa8-cp312-cp312-linux_x86_64.whl'
TRITON_URL='https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/triton-3.5.1%2Brocm7.2.1.gita272dfa8-cp312-cp312-linux_x86_64.whl'
TRITON_SHA='07787af1d28c273852f897bfeaa7bca29f2fa4a13ca0f28f535832b240ce7016'

if ((DRY_RUN)); then
  printf 'python=%s\nvenv=%s\nrocm_release=%s\nwheel_cache=%s\ndependency_lock=%s\n' \
    "$PYTHON_BIN" "$VENV_DIR" "$ROCM_RELEASE" "$WHEEL_DIR" "$LOCK_FILE"
  printf 'dependencies_only=%s\n' "$DEPENDENCIES_ONLY"
  printf 'download=%s\ndownload=%s\ndownload=%s\n' "$TORCH_URL" "$VISION_URL" "$TRITON_URL"
  printf 'next=%s/bin/unlimited-ocr-rdna4 doctor --device 0\n' "$VENV_DIR"
  exit 0
fi

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64) ;;
  *)
    printf 'error: the verified bootstrap supports Linux x86_64 only\n' >&2
    exit 3
    ;;
esac

if [[ ! -r /opt/rocm/.info/version ]]; then
  printf 'error: /opt/rocm/.info/version is missing; install a supported host ROCm stack first\n' >&2
  exit 3
fi
HOST_ROCM=$(tr -d '\r\n' < /opt/rocm/.info/version)
if [[ $HOST_ROCM != "$ROCM_RELEASE" ]]; then
  printf 'error: this bootstrap is verified for host ROCm 7.2.1, found %s\n' "$HOST_ROCM" >&2
  exit 3
fi

PYTHON_VERSION=$("$PYTHON_BIN" -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
if [[ $PYTHON_VERSION != 3.12 ]]; then
  printf 'error: AMD wheel set requires Python 3.12, found %s\n' "$PYTHON_VERSION" >&2
  exit 3
fi
command -v curl >/dev/null || { printf 'error: curl is required\n' >&2; exit 3; }
command -v sha256sum >/dev/null || { printf 'error: sha256sum is required\n' >&2; exit 3; }
command -v flock >/dev/null || { printf 'error: flock is required\n' >&2; exit 3; }
[[ -f $LOCK_FILE && ! -L $LOCK_FILE ]] || { printf 'error: dependency lock is missing or symlinked: %s\n' "$LOCK_FILE" >&2; exit 3; }

mkdir -p "$WHEEL_DIR" "$TASK_TMP"
export TMPDIR="$TASK_TMP"
export PIP_CACHE_DIR="$CACHE_BASE/pip"
exec 9>"$WHEEL_DIR/.bootstrap.lock"
flock 9

download_wheel() {
  local filename=$1
  local url=$2
  local expected=$3
  local destination="$WHEEL_DIR/$filename"
  local partial="$destination.part"

  if [[ -L $destination || -L $partial ]]; then
    printf 'error: refusing symlinked wheel cache entry: %s\n' "$destination" >&2
    exit 4
  fi
  if [[ -f $destination ]]; then
    if printf '%s  %s\n' "$expected" "$destination" | sha256sum --check --status; then
      printf 'wheel=reused path=%s\n' "$destination"
      return
    fi
    printf 'wheel=discard-corrupt path=%s\n' "$destination" >&2
    rm -f -- "$destination"
  fi

  for attempt in 1 2; do
    printf 'downloading=%s attempt=%s\n' "$filename" "$attempt"
    if ! curl --fail --location --proto '=https' --tlsv1.2 --continue-at - --output "$partial" "$url"; then
      printf 'wheel=discard-failed-partial path=%s\n' "$partial" >&2
      rm -f -- "$partial"
      continue
    fi
    if printf '%s  %s\n' "$expected" "$partial" | sha256sum --check --status; then
      mv "$partial" "$destination"
      return
    fi
    printf 'wheel=discard-corrupt-partial path=%s\n' "$partial" >&2
    rm -f -- "$partial"
  done
  printf 'error: downloaded wheel checksum mismatch after clean retry: %s\n' "$filename" >&2
  exit 4
}

download_wheel "$TORCH_FILE" "$TORCH_URL" "$TORCH_SHA"
download_wheel "$VISION_FILE" "$VISION_URL" "$VISION_SHA"
download_wheel "$TRITON_FILE" "$TRITON_URL" "$TRITON_SHA"

validate_venv() {
  if [[ -L $VENV_DIR || ! -d $VENV_DIR || -L $VENV_DIR/pyvenv.cfg || ! -f $VENV_DIR/pyvenv.cfg ]]; then
    printf 'error: existing --venv path is not a regular Python virtual environment: %s\n' "$VENV_DIR" >&2
    exit 3
  fi
  if [[ ! -x $VENV_DIR/bin/python ]]; then
    printf 'error: virtual environment has no executable bin/python: %s\n' "$VENV_DIR" >&2
    exit 3
  fi
  "$VENV_DIR/bin/python" - "$VENV_DIR" <<'PY'
import os
import sys

expected = os.path.realpath(sys.argv[1])
if os.path.realpath(sys.prefix) != expected or sys.prefix == sys.base_prefix:
    raise SystemExit("virtual environment prefix validation failed")
if sys.version_info[:2] != (3, 12):
    raise SystemExit(f"virtual environment requires Python 3.12, found {sys.version_info.major}.{sys.version_info.minor}")
PY
}

if [[ -L $VENV_DIR ]]; then
  printf 'error: --venv path must not be a symlink: %s\n' "$VENV_DIR" >&2
  exit 3
fi
if [[ -e $VENV_DIR && ! -d $VENV_DIR ]]; then
  printf 'error: --venv path exists and is not a directory: %s\n' "$VENV_DIR" >&2
  exit 3
fi
if [[ ! -e $VENV_DIR ]]; then
  "$PYTHON_BIN" -m venv "$VENV_DIR"
fi
validate_venv
"$VENV_DIR/bin/python" -m pip install --require-hashes --only-binary=:all: -r "$LOCK_FILE"
"$VENV_DIR/bin/python" -m pip install \
  --no-deps \
  "$WHEEL_DIR/$TORCH_FILE" \
  "$WHEEL_DIR/$VISION_FILE" \
  "$WHEEL_DIR/$TRITON_FILE"
if ((DEPENDENCIES_ONLY)); then
  "$VENV_DIR/bin/python" -m pip check
  printf 'bootstrap=PASS mode=dependencies-only venv=%s\n' "$VENV_DIR"
  exit 0
fi
BUILD_DIR=$(mktemp -d "$TASK_TMP/package-build.XXXXXX")
trap 'rm -rf -- "$BUILD_DIR"' EXIT
"$VENV_DIR/bin/python" -m pip wheel --no-build-isolation --no-deps --wheel-dir "$BUILD_DIR" "$ROOT"
PACKAGE_WHEEL=$(find "$BUILD_DIR" -maxdepth 1 -type f -name 'unlimited_ocr_rdna4-*.whl' -print -quit)
[[ -n $PACKAGE_WHEEL ]] || { printf 'error: local package wheel was not produced\n' >&2; exit 4; }
"$VENV_DIR/bin/python" -m pip install --force-reinstall --no-deps "$PACKAGE_WHEEL"
"$VENV_DIR/bin/python" -m pip check
"$VENV_DIR/bin/unlimited-ocr-rdna4" doctor --device "${UNLIMITED_OCR_DEVICE:-0}"
printf 'bootstrap=PASS venv=%s\n' "$VENV_DIR"