#!/usr/bin/env bash
# pyspark — launch a Python REPL with a `spark` global wired to our
# Spark Connect server. Mirrors Apache Spark's bin/pyspark.
#
# Usage:
#   pyspark                       # interactive
#   pyspark script.py             # run a script with `spark` available
#   pyspark --port 15002          # pin the connect port (default: random)
#
# Requirements: pyspark>=4.0.0 installed (`pip install pyspark[connect]`).

set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"

# Which profile to build if no server binary exists yet. Release is the
# right default for the demo/serving path: the server's per-request cost
# is planning-dominated (see docs/archive/decisions/connect-concurrency-bench/
# B1-server-planning-cost-2026-07-01.md), and a debug build runs that
# planning several times slower — a laggy first impression for `pyspark`.
# Override with PYSPARK_BUILD_PROFILE=debug for a faster one-time compile
# during local server development.
BUILD_PROFILE="${PYSPARK_BUILD_PROFILE:-release}"
case "$BUILD_PROFILE" in
    release|debug) ;;
    *)
        echo "==> PYSPARK_BUILD_PROFILE must be 'release' or 'debug' (got '$BUILD_PROFILE'); using 'release'" >&2
        BUILD_PROFILE="release"
        ;;
esac

# Prefer an already-built binary of EITHER profile (don't force a rebuild
# if the user already has one): release first (faster at runtime), then
# debug. Only when neither exists do we build — and then we honor
# BUILD_PROFILE (release by default).
SERVER_BIN="$ROOT/target/release/spark-connect-server"
if [[ ! -x "$SERVER_BIN" ]]; then
    SERVER_BIN="$ROOT/target/debug/spark-connect-server"
fi
if [[ ! -x "$SERVER_BIN" ]]; then
    if [[ "$BUILD_PROFILE" == "release" ]]; then
        echo "==> building spark-connect-server (release, one-time — this is slower to" >&2
        echo "    compile but much faster to run; set PYSPARK_BUILD_PROFILE=debug to" >&2
        echo "    trade runtime speed for a quicker one-time build)" >&2
        (cd "$ROOT" && cargo build -q --release -p spark-connect --bin spark-connect-server) >&2
        SERVER_BIN="$ROOT/target/release/spark-connect-server"
    else
        echo "==> building spark-connect-server (debug, one-time — PYSPARK_BUILD_PROFILE=debug)" >&2
        (cd "$ROOT" && cargo build -q -p spark-connect --bin spark-connect-server) >&2
        SERVER_BIN="$ROOT/target/debug/spark-connect-server"
    fi
fi

PORT="${SPARK_CONNECT_PORT:-0}"   # 0 => pick a free port
if [[ "$PORT" == "0" ]]; then
    PORT=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1",0)); print(s.getsockname()[1]); s.close()')
fi

DIALECT="${SPARK_DIALECT:-spark4}"

echo "==> pyspark: launching spark-connect-server on :$PORT (dialect=$DIALECT)" >&2
SPARK_CONNECT_PORT="$PORT" "$SERVER_BIN" >/tmp/spark-connect-pyspark.log 2>&1 &
SERVER_PID=$!
trap "kill $SERVER_PID 2>/dev/null || true" EXIT INT TERM

# Poll until the server is listening (max 5s).
for _ in $(seq 1 50); do
    if python3 -c "import socket; s=socket.socket(); s.connect(('127.0.0.1', $PORT)); s.close()" 2>/dev/null; then
        break
    fi
    sleep 0.1
done

STARTUP_PY=$(mktemp /tmp/pyspark-startup-XXXX.py)
cat > "$STARTUP_PY" <<EOF
import os, sys
from pyspark.sql.connect.session import SparkSession
spark = SparkSession.builder.remote("sc://127.0.0.1:$PORT").getOrCreate()
if "$DIALECT" == "spark4":
    spark.conf.set("spark.sql.ansi.enabled", "true")
print(f"spark-rust-jaaicode pyspark — Spark Connect at sc://127.0.0.1:$PORT")
print(f"  dialect: $DIALECT")
print(f"  Use the variable: spark")
print()
EOF
trap "rm -f $STARTUP_PY; kill $SERVER_PID 2>/dev/null || true" EXIT INT TERM

if [[ $# -eq 0 ]]; then
    # PYTHONSTARTUP loads on every interactive python (REPL or stdin).
    # Use `python3 -i` so it triggers even when stdin is a pipe.
    if [[ -t 0 ]]; then
        PYTHONSTARTUP="$STARTUP_PY" python3
    else
        PYTHONSTARTUP="$STARTUP_PY" python3 -i </dev/stdin
    fi
else
    # Run a script: prepend startup, pass remaining args.
    SCRIPT="$1"; shift
    python3 -c "exec(open('$STARTUP_PY').read()); exec(open('$SCRIPT').read())" "$@"
fi
