#!/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_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.
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
}

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

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

# Print the newest ptychodus:*-<suffix> tag in local storage (empty if none).
newest_image_for_suffix() {
    local suffix=$1
    podman images --format '{{.Repository}}:{{.Tag}}' 2>/dev/null \
        | awk -F: -v s="$suffix" '$1=="ptychodus" && $2 ~ ("-"s"$") {print}' \
        | sort -V \
        | tail -n1
}

if [[ -n "${PTYCHODUS_IMAGE:-}" ]]; then
    IMAGE="$PTYCHODUS_IMAGE"
    family="override"
    notice "using image $IMAGE (override: PTYCHODUS_IMAGE)"
else
    family=$(detect_gpu)
    suffix="${BACKEND_FOR[$family]:-}"
    [[ -n "$suffix" ]] || die "no backend mapped for family '$family'"
    IMAGE=$(newest_image_for_suffix "$suffix")
    if [[ -z "$IMAGE" ]]; then
        # Fall back to the unversioned tag so the "not found" branch below
        # can print an actionable build command.
        IMAGE="ptychodus:$suffix"
    fi
    notice "using image $IMAGE (detected: $family)"
fi

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

command -v podman >/dev/null 2>&1 \
    || die "podman not installed; see docs/source/getting_started.md"

if ! podman image exists "$IMAGE"; then
    {
        if [[ -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
    --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 )

if [[ "$PWD" != "$HOME"* ]]; then
    args+=( -v "$PWD:$PWD" )
fi

for mount in "${BEAMLINE_MOUNTS[@]}"; do
    [[ -e "$mount" ]] || continue
    [[ "$mount" == "$HOME"* ]] && 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 "$@"
