#!/usr/bin/env bash
# Transparent podman launcher for ptychodus.
# Detects host GPU family, picks the matching ptychodus image, forwards X11,
# bind-mounts $HOME / $PWD / beamline data roots, and runs rootless so files
# written to the host are owned by the invoking user.
#
# Lint with: shellcheck scripts/podman/ptychodus
#
# Env-var overrides:
#   PTYCHODUS_IMAGE   skip GPU detection; use this image tag verbatim
#   PTYCHODUS_VERSION pin the ptychodus version, keeping GPU detection; matched
#                     exactly against the version part of the tag
#   PTYCHODUS_QUIET   set to any non-empty value to suppress the stderr notice

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

# --- Configuration (admins edit here) ----------------------------------------

# Backend suffix that pairs with each detected GPU family. The actual image
# tag is `ptychodus:<version>-<suffix>`; the newest matching versioned tag
# in local podman storage is picked at runtime so a rebuild "just works"
# without editing this file.
#
# The nvidia row is only a FALLBACK, used when the driver's CUDA ceiling
# cannot be read. Normally the CUDA variant is sensed at runtime: the
# candidates come from the `ptychodus:*-cuda<X.Y>` tags actually present in
# podman storage, and the newest one the installed driver supports wins. See
# choose_cuda_suffix below.
declare -A BACKEND_FOR=(
    [nvidia]=cuda13.0
    [rocm]=rocm7.2.4
    [xpu]=xpu
    [cpu]=cpu
)

BEAMLINE_MOUNTS=(/local /gdata)

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

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

notice() {
    [[ -n "${PTYCHODUS_QUIET:-}" ]] && return 0
    printf 'ptychodus: %s\n' "$*" >&2
}

detect_gpu() {
    if command -v nvidia-smi >/dev/null 2>&1 \
       && [[ -n "$(nvidia-smi -L 2>/dev/null)" ]]; then
        printf nvidia
        return
    fi

    if [[ -e /dev/kfd ]]; then
        shopt -s nullglob
        local render_nodes=(/dev/dri/renderD*)
        shopt -u nullglob
        if (( ${#render_nodes[@]} > 0 )); then
            printf rocm
            return
        fi
    fi

    if command -v xpu-smi >/dev/null 2>&1 \
       && xpu-smi discovery >/dev/null 2>&1; then
        printf xpu
        return
    fi
    shopt -s nullglob
    local vendor_files=(/sys/class/drm/card*/device/vendor)
    shopt -u nullglob
    local vf
    for vf in "${vendor_files[@]}"; do
        if [[ -r "$vf" ]] && grep -q '0x8086' "$vf"; then
            printf xpu
            return
        fi
    done

    printf cpu
}

# Print the highest CUDA runtime the installed NVIDIA driver supports, as
# `<major>.<minor>` (empty if it cannot be determined).
#
# This is a property of the DRIVER, not of any CUDA toolkit on the host -- a
# beamline workstation running containers has no nvcc, and needs none. The
# ceiling matters because a CUDA major version bump is not covered by minor
# version compatibility: a cuda13.x image on an r570 driver (CUDA 12.8) leaves
# torch.cuda.is_available() False, which surfaces much later as a mystifying
# CPU-speed reconstruction rather than as a launch error.
#
# `nvidia-smi -q` is the machine-readable source ("CUDA Version : 12.8"); the
# banner regex is the fallback for drivers whose -q output omits the field.
detect_cuda_version() {
    local version
    version=$(nvidia-smi -q 2>/dev/null \
        | sed -n 's/^ *CUDA Version *: *\([0-9][0-9.]*\).*/\1/p' \
        | head -n1)
    if [[ -z "$version" ]]; then
        version=$(nvidia-smi 2>/dev/null \
            | sed -n 's/.*CUDA Version: *\([0-9][0-9.]*\).*/\1/p' \
            | head -n1)
    fi
    printf '%s' "$version"
}

# True when version $1 is no newer than version $2.
version_le() {
    [[ "$1" == "$2" ]] && return 0
    [[ "$(printf '%s\n%s\n' "$1" "$2" | sort -V | head -n1)" == "$1" ]]
}

add_gpu_flags() {
    local -n _args=$1
    local family=$2
    case "$family" in
        nvidia)
            _args+=( --device nvidia.com/gpu=all )
            ;;
        rocm)
            _args+=( --device /dev/kfd --device /dev/dri
                     --group-add video --group-add render )
            ;;
        xpu)
            _args+=( --device /dev/dri
                     --group-add video --group-add render )
            ;;
        cpu)
            ;;
        *)
            ;;
    esac
}

add_x11_flags() {
    local -n _args=$1
    _args+=( --env DISPLAY -v /tmp/.X11-unix:/tmp/.X11-unix )
    if [[ -n "${XAUTHORITY:-}" && -r "${XAUTHORITY:-}" ]]; then
        _args+=( --env XAUTHORITY -v "$XAUTHORITY:$XAUTHORITY:ro" )
    fi
}

# --- Pre-scan argv for batch mode (no argv mutation) -------------------------

is_batch=0
for arg in "$@"; do
    case "$arg" in
        -b|--batch|--batch=*) is_batch=1 ;;
    esac
done

# --- Prerequisite checks -----------------------------------------------------

# Ahead of the image choice below rather than after it, because choosing
# shells out to `podman images` and would otherwise fail first, on a bare
# exit status instead of on this message.
command -v podman >/dev/null 2>&1 \
    || die "podman not installed; see docs/source/getting_started.md"

# --- Choose image ------------------------------------------------------------

# Print every ptychodus image in local podman storage, one per line, as four
# TAB-separated fields: `<sort-key> <version> <backend> <reference>`. The
# helpers below are all filters over this table, and $TAG_TABLE holds one
# snapshot of it, because `podman images` is the slowest thing this script
# runs.
#
# podman reports a locally built image's repository fully qualified, as
# `localhost/ptychodus`, so the name has to be compared against the last path
# component rather than the whole field -- matching `$1 == "ptychodus"` finds
# nothing at all. The full reference is carried along, since that is what
# `podman run` wants.
#
# The sort key is the version with a `~` inserted before a PEP 440 `dev` or
# `rc` marker, which is not a typo: GNU `sort -V` ranks `~` below the empty
# string, so `1.5.2~dev14` sorts below `1.5.2` the way PEP 440 says a dev build
# precedes its release. Without it a stale `1.5.2.dev14_g79157a2b2` image
# shadows the real 1.5.2 forever. `a` and `b` markers are deliberately left
# out: a setuptools_scm local segment is a hex hash (`_g79157a2b2`), where
# `a2` can occur but `dev` and `rc` cannot.
ptychodus_tags() {
    # The tag is taken as the field after the LAST colon so that a reference
    # carrying a registry port parses; splitting at the first colon reads the
    # port as the tag. Version and backend split at the last `-`, which no
    # backend name and no normalized PEP 440 version contains; a tag with no
    # `-` at all is skipped rather than mis-parsed.
    podman images --format '{{.Repository}}:{{.Tag}}' 2>/dev/null \
        | awk -F: '{ tag = $NF; ref = $0; sub(/:[^:]*$/, "", ref)
                     repo = ref; sub(/.*\//, "", repo) }
                   repo == "ptychodus" && match(tag, /-[^-]*$/) {
                       version = substr(tag, 1, RSTART - 1)
                       key = version; sub(/\.?(dev|rc)[0-9]/, "~&", key)
                       sub(/~\./, "~", key)
                       printf "%s\t%s\t%s\t%s\n",
                              key, version, substr(tag, RSTART + 1), $0 }'
}

# Print the newest ptychodus image with backend suffix $1 (empty if none),
# honouring a $version_pin when one is set.
#
# The suffix is compared literally rather than as a regex: every backend name
# carrying a version has dots in it, and as a pattern `cuda13.0` would also
# match a `cuda1320` tag.
newest_image_for_suffix() {
    local suffix=$1
    awk -F'\t' -v s="$suffix" -v pin="$version_pin" \
        '$3 == s && (pin == "" || $2 == pin)' <<<"$TAG_TABLE" \
        | sort -V -k1,1 \
        | tail -n1 \
        | cut -f4
}

# Print the CUDA versions that have a ptychodus image in local storage, newest
# first, one `<major>.<minor>` per line. A $version_pin restricts the
# candidates to that version, so the chosen CUDA variant is one the pinned
# version was actually built for.
#
# Reading the candidates out of podman storage rather than restating
# scripts/podman/build's VARIANT_ARGS matrix here is deliberate: a new CUDA row
# over there becomes selectable the moment it is built, and a row that was
# never built is never chosen.
available_cuda_versions() {
    awk -F'\t' -v pin="$version_pin" \
        '(pin == "" || $2 == pin) && $3 ~ /^cuda[0-9]+\.[0-9]+$/ { print substr($3, 5) }' \
        <<<"$TAG_TABLE" \
        | sort -Vru
}

# Print the versions built for backend suffix $1, and the backend suffixes
# built for version $1, respectively. Both ignore $version_pin: they exist to
# tell an operator what a failed pinned lookup could have matched.
versions_for_suffix() {
    awk -F'\t' -v s="$1" '$3 == s { print $2 }' <<<"$TAG_TABLE" | sort -V -u
}

backends_for_version() {
    awk -F'\t' -v v="$1" '$2 == v { print $3 }' <<<"$TAG_TABLE" | sort -u
}

# Render a newline-separated list on one line for an error message.
list_or_none() {
    local items=$1
    if [[ -n "$items" ]]; then
        printf '%s' "${items//$'\n'/ }"
    else
        printf '(none)'
    fi
}

# Set $1 to the backend suffix for an NVIDIA host: the newest built CUDA
# variant the driver can actually run. $2 is the driver's CUDA ceiling,
# possibly empty.
#
# This writes through a nameref rather than printing, because the
# no-compatible-variant path below exits -- and an exit inside the `$(...)`
# this would otherwise be called from only kills the subshell, leaving the
# caller to carry on with an empty suffix.
choose_cuda_suffix() {
    local -n _suffix=$1
    local host_cuda=$2

    if [[ -z "$host_cuda" ]]; then
        notice "cannot read the driver CUDA version from nvidia-smi; falling back to ${BACKEND_FOR[nvidia]}"
        _suffix="${BACKEND_FOR[nvidia]}"
        return
    fi

    local -a built
    mapfile -t built < <(available_cuda_versions)

    # Nothing built yet: hand back the fallback so the "no image found" branch
    # below reports it with an actionable build command.
    if (( ${#built[@]} == 0 )); then
        _suffix="${BACKEND_FOR[nvidia]}"
        return
    fi

    local version
    for version in "${built[@]}"; do
        if version_le "$version" "$host_cuda"; then
            _suffix="cuda$version"
            return
        fi
    done

    # Every built variant is newer than the driver supports. Fail loudly
    # instead of quietly dropping to the cpu image -- on a beamline a
    # CPU-speed reconstruction reads as a hang, not as a misconfiguration.
    {
        printf 'ptychodus: host NVIDIA driver supports CUDA <= %s, but every ptychodus\n' "$host_cuda"
        printf '  CUDA image in podman storage is newer:'
        printf ' cuda%s' "${built[@]}"
        printf '\n'
        if [[ -n "$version_pin" ]]; then
            printf '  (candidates restricted to the pinned version %s)\n' "$version_pin"
        fi
        printf '  Build a variant no newer than CUDA %s from the repository root,\n' "$host_cuda"
        printf '  choosing one from the matrix at the top of scripts/podman/build:\n'
        printf '    scripts/podman/build cuda%s\n' "$host_cuda"
        printf '  Or force an image with PTYCHODUS_IMAGE=<tag>.\n'
    } >&2
    exit 1
}

# OCI tags accept only [A-Za-z0-9._-], and scripts/podman/build substitutes `_`
# for everything else when it tags an image. Applying the same rule to the pin
# lets a real PEP 440 version be pasted in: 1.5.2.dev14+g79157a2b2 finds the
# image tagged 1.5.2.dev14_g79157a2b2.
version_pin=''
if [[ -n "${PTYCHODUS_VERSION:-}" ]]; then
    version_pin="${PTYCHODUS_VERSION//[^A-Za-z0-9._-]/_}"
    if [[ "$version_pin" != "$PTYCHODUS_VERSION" ]]; then
        notice "PTYCHODUS_VERSION=$PTYCHODUS_VERSION is not tag-safe; matching $version_pin"
    fi
fi

if [[ -n "${PTYCHODUS_IMAGE:-}" ]]; then
    IMAGE="$PTYCHODUS_IMAGE"
    family="override"
    if [[ -n "$version_pin" ]]; then
        notice "PTYCHODUS_IMAGE is set; ignoring PTYCHODUS_VERSION=$PTYCHODUS_VERSION"
    fi
    notice "using image $IMAGE (override: PTYCHODUS_IMAGE)"
else
    TAG_TABLE=$(ptychodus_tags)
    family=$(detect_gpu)
    detected="$family"
    if [[ "$family" == nvidia ]]; then
        host_cuda=$(detect_cuda_version)
        choose_cuda_suffix suffix "$host_cuda"
        if [[ -n "$host_cuda" ]]; then
            detected="$family, driver CUDA $host_cuda"
        fi
    else
        suffix="${BACKEND_FOR[$family]:-}"
    fi
    [[ -n "$suffix" ]] || die "no backend mapped for family '$family'"
    [[ -n "$version_pin" ]] && detected="$detected, pinned to $version_pin"
    IMAGE=$(newest_image_for_suffix "$suffix")
    if [[ -z "$IMAGE" ]]; then
        # Fall back to the tag that was looked for, so the "not found" branch
        # below can report it with an actionable build command.
        IMAGE="ptychodus:${version_pin:+$version_pin-}$suffix"
    fi
    notice "using image $IMAGE (detected: $detected)"
fi

# --- Verify the chosen image exists ------------------------------------------

if ! podman image exists "$IMAGE"; then
    {
        if [[ -n "${suffix:-}" && -n "$version_pin" ]]; then
            # A pinned lookup fails for one of two reasons -- the version was
            # never built, or it was never built for this backend -- and the
            # operator cannot tell which without seeing both lists.
            printf 'ptychodus: no ptychodus:%s-%s image found in rootless podman storage.\n' \
                "$version_pin" "$suffix"
            printf '  Versions built for %s: %s\n' \
                "$suffix" "$(list_or_none "$(versions_for_suffix "$suffix")")"
            printf '  Backends built for %s: %s\n' \
                "$version_pin" "$(list_or_none "$(backends_for_version "$version_pin")")"
            # Not `PTYCHODUS_VERSION=... scripts/podman/build`: over there the
            # same variable only relabels the tag, so suggesting it would
            # invite tagging current code as an old release.
            printf '  Build it from a %s checkout with:\n' "$version_pin"
            printf '    scripts/podman/build %s\n' "$suffix"
        elif [[ -n "${suffix:-}" ]]; then
            printf 'ptychodus: no ptychodus:*-%s image found in rootless podman storage.\n' "$suffix"
            printf '  Build it from the repository root with:\n'
            printf '    scripts/podman/build %s\n' "$suffix"
        else
            printf 'ptychodus: image "%s" not found in rootless podman storage.\n' "$IMAGE"
            printf '  Build the desired variant with scripts/podman/build from the repository root.\n'
        fi
    } >&2
    exit 1
fi

# --- Assemble podman flags ---------------------------------------------------

args=(
    run --rm -i
    --shm-size=2g
    --ulimit host
    --userns=keep-id
    --security-opt label=type:container_runtime_t
    --network host
    --workdir "$PWD"
    -v "$HOME:$HOME"
    --env HOME
)

[[ -t 0 && -t 1 ]] && args+=( -t )

# A symlinked home -- /home/beams/$USER -> /home/beams0/$USER on APS
# workstations -- is reachable inside the container only at the literal $HOME
# spelling: podman resolves the bind source but creates the destination
# verbatim, so the real path is simply absent. Absolute paths that a previous
# host-side run canonicalized (settings.ini entries, Qt file-dialog history,
# -i/-o arguments typed as realpaths) all point at that absent spelling, so
# mount it too.
HOME_REAL=$(realpath -- "$HOME" 2>/dev/null) || HOME_REAL="$HOME"
if [[ "$HOME_REAL" != "$HOME" ]]; then
    args+=( -v "$HOME_REAL:$HOME_REAL" )
fi

# Containment tests below consider both spellings: with only the literal one,
# a $PWD or beamline root under the real home is mounted a second time,
# nested inside a mount that already covers it.
is_under_home() {
    local path=$1
    [[ "$path" == "$HOME" || "$path" == "$HOME"/* ]] && return 0
    [[ "$path" == "$HOME_REAL" || "$path" == "$HOME_REAL"/* ]]
}

if ! is_under_home "$PWD"; then
    args+=( -v "$PWD:$PWD" )
fi

for mount in "${BEAMLINE_MOUNTS[@]}"; do
    [[ -e "$mount" ]] || continue
    is_under_home "$mount" && continue
    [[ "$PWD" == "$mount"* ]] && continue
    args+=( -v "$mount:$mount" )
done

add_gpu_flags args "$family"

if (( is_batch == 0 )) && [[ -n "${DISPLAY:-}" ]]; then
    add_x11_flags args
fi

# --- Launch ------------------------------------------------------------------

exec podman "${args[@]}" "$IMAGE" python3 -m ptychodus "$@"
