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

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

if [ -n "${OPENPI_PYTHON:-}" ]; then
    PY="$OPENPI_PYTHON"
elif [ -x "$ROOT/.venv/bin/python" ]; then
    PY="$ROOT/.venv/bin/python"
else
    PY="$(command -v python3 || true)"
fi

if [ -z "$PY" ] || [ ! -x "$PY" ]; then
    echo "ERROR: no usable Python interpreter found"
    exit 1
fi

CHECKPOINT="${POLICY_CHECKPOINT:-$ROOT/checkpoints/29999}"
HOST="${POLICY_HOST:-0.0.0.0}"
PORT="${POLICY_PORT:-8000}"
DEFAULT_PROMPT="${POLICY_DEFAULT_PROMPT:-}"

export PYTHONDONTWRITEBYTECODE=1
export PYTHONPATH="$ROOT:$ROOT/src:$ROOT/packages/openpi-client/src"
export OPENPI_REAL_PIPER_X_ASSETS="$CHECKPOINT/assets"
export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}"

if [ -n "${POLICY_GPU:-}" ]; then
    export CUDA_VISIBLE_DEVICES="$POLICY_GPU"
fi

ARGS=(
    --checkpoint "$CHECKPOINT"
    --host "$HOST"
    --port "$PORT"
)

if [ -n "$DEFAULT_PROMPT" ]; then
    ARGS+=(--default-prompt "$DEFAULT_PROMPT")
fi

echo "Repository: $ROOT"
echo "Python: $PY"
echo "Checkpoint: $CHECKPOINT"
echo "Server: ws://$HOST:$PORT"
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unchanged}"

cd "$ROOT"
exec "$PY" -u deployment/serve_piperx.py "${ARGS[@]}" "$@"