#!/usr/bin/env bash
# Build the CPU, CUDA, ROCm, and XPU ptychodus container images with the
# current ptychodus version threaded in as a build arg and included in the
# image tag.
#
# Tag format:
#   ptychodus:<tag-safe-version>-<backend>
#
#   backend ∈ { cpu, cuda12.8, cuda13.0, cuda13.2, rocm7.2.4, xpu }
#
# The PEP 440 version goes into the image verbatim via the PTYCHODUS_VERSION
# build arg, but OCI tags only accept [A-Za-z0-9._-], so characters outside
# that set are replaced with `_` in the tag. In practice this is the `+` of a
# dev checkout's local segment: 1.5.2.dev14+g79157a2b2 tags as
# 1.5.2.dev14_g79157a2b2, while a clean release tag is unchanged.
#
# Dectris is intentionally excluded: containers/Dockerfile.dectris installs
# ptychodus from PyPI rather than the checkout, so a version-in-tag would
# misrepresent the image. Build it separately with a plain `podman build
# -f containers/Dockerfile.dectris -t ptychodus:dectris .`.
#
# Usage:
#   scripts/podman/build                                # cpu + all cuda + rocm + xpu
#   scripts/podman/build cpu cuda13.0                   # subset from the matrix below
#
# Env-var overrides:
#   PTYCHODUS_VERSION   PEP 440 version to bake in (default: `uv run --with
#                       setuptools-scm python -m setuptools_scm` in the
#                       checkout, matching what a normal pip install would
#                       derive from git).
#   CONTAINER_ENGINE    build engine (default: podman; e.g. docker)
#
# Lint with: shellcheck scripts/podman/build

set -euo pipefail
IFS=$'\n\t'

# --- Configuration -----------------------------------------------------------

# ROCm version tracks the OLCF Frontier line, pinned to a patch release that
# rocm/pytorch actually publishes -- plain 7.2.0 has no tag on Docker Hub.
# Keep in sync with containers/Dockerfile.rocm's ARG default.
ROCM_VERSION=7.2.4

# Pair each CUDA_VERSION with a PYTORCH_VERSION that is actually published
# on Docker Hub as pytorch/pytorch:${PYTORCH_VERSION}-cuda${CUDA_VERSION}-cudnn9-devel.
# The containers/Dockerfile.cuda default (2.12.0) is not published for
# CUDA 12.8, so that row overrides.
declare -A VARIANT_ARGS=(
    [cpu]="-f containers/Dockerfile.cpu"
    [cuda12.8]="-f containers/Dockerfile.cuda --build-arg CUDA_VERSION=12.8 --build-arg PYTORCH_VERSION=2.11.0"
    [cuda13.0]="-f containers/Dockerfile.cuda --build-arg CUDA_VERSION=13.0 --build-arg PYTORCH_VERSION=2.12.0"
    [cuda13.2]="-f containers/Dockerfile.cuda --build-arg CUDA_VERSION=13.2 --build-arg PYTORCH_VERSION=2.12.0"
    [rocm${ROCM_VERSION}]="-f containers/Dockerfile.rocm --build-arg ROCM_VERSION=${ROCM_VERSION}"
    [xpu]="-f containers/Dockerfile.xpu"
)

DEFAULT_VARIANTS=(cpu cuda12.8 cuda13.0 cuda13.2 "rocm${ROCM_VERSION}" xpu)

# --- Helpers -----------------------------------------------------------------

die() {
    printf 'build: %s\n' "$*" >&2
    exit 1
}

# --- Setup -------------------------------------------------------------------

# Resolve the checkout from this script's own location rather than $PWD.
# `git rev-parse --show-toplevel` against the caller's cwd succeeds in
# whatever repo the caller happens to be sitting in, so a build launched from
# an unrelated checkout would silently cd there and fail later on a missing
# containers/Dockerfile instead of reporting the real problem here.
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P) \
    || die "cannot resolve the directory containing this script"
repo_root=$(git -C "$script_dir" rev-parse --show-toplevel) \
    || die "$script_dir is not inside a git checkout; run this script from a ptychodus clone"
cd "$repo_root" || die "cannot enter the ptychodus checkout at $repo_root"

if [[ -n "${PTYCHODUS_VERSION:-}" ]]; then
    VERSION="$PTYCHODUS_VERSION"
else
    command -v uv >/dev/null 2>&1 \
        || die "uv not on PATH; set PTYCHODUS_VERSION explicitly (PEP 440) instead"
    VERSION=$(uv run --quiet --with setuptools-scm python -m setuptools_scm) \
        || die "setuptools_scm failed; set PTYCHODUS_VERSION explicitly (PEP 440) instead"
fi

# OCI tags allow only [A-Za-z0-9._-] after a leading alphanumeric/underscore,
# so a PEP 440 local segment ('+g79157a2b2') is rejected outright. Sanitize for
# the tag only; PTYCHODUS_VERSION keeps the real version.
TAG_VERSION="${VERSION//[^A-Za-z0-9._-]/_}"

ENGINE="${CONTAINER_ENGINE:-podman}"
command -v "$ENGINE" >/dev/null 2>&1 \
    || die "$ENGINE not on PATH (set CONTAINER_ENGINE=docker to use docker)"

# --- Resolve requested variants ----------------------------------------------

if (( $# == 0 )); then
    variants=("${DEFAULT_VARIANTS[@]}")
else
    variants=("$@")
    for v in "${variants[@]}"; do
        [[ -n "${VARIANT_ARGS[$v]+set}" ]] \
            || die "unknown variant '$v' (choose from: $(IFS=' '; echo "${!VARIANT_ARGS[*]}"))"
    done
fi

# --- Build -------------------------------------------------------------------

for backend in "${variants[@]}"; do
    tag="ptychodus:${TAG_VERSION}-${backend}"
    printf '\n=== %s (version=%s) ===\n' "$tag" "$VERSION"
    # VARIANT_ARGS entries are space-separated argv, but the strict-mode IFS
    # set above contains no space, so unquoted word-splitting would hand the
    # engine one giant argument. Split on spaces explicitly instead.
    IFS=' ' read -r -a variant_args <<<"${VARIANT_ARGS[$backend]}"
    "$ENGINE" build "${variant_args[@]}" \
        --build-arg "PTYCHODUS_VERSION=$VERSION" \
        -t "$tag" .
done
