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

PG_VERSION="${PG_VERSION:-pg17}"
PG_FEATURE="${PG_FEATURE:-pg17}"
PG_CONFIG="${PG_CONFIG:-/opt/homebrew/opt/postgresql@17/bin/pg_config}"
PGHOST="${PGHOST:-localhost}"
PGPORT="${PGPORT:-28817}"
PGRX_HOME="${PGRX_HOME:-${HOME}/.pgrx}"
PGDATA="${PGDATA:-${PGRX_HOME}/data-${PG_VERSION#pg}}"
PG_CTL="${PG_CTL:-$(dirname "${PG_CONFIG}")/pg_ctl}"
DBNAME="${DBNAME:-pgcontext_hnsw_restart_check}"

if [[ ! "${DBNAME}" =~ ^[A-Za-z_][A-Za-z0-9_]*$ ]]; then
    echo "DBNAME must be a simple SQL identifier" >&2
    exit 2
fi

restart_after_stop=false
cleanup() {
    if [[ "${restart_after_stop}" == "true" ]]; then
        cargo pgrx start "${PG_VERSION}" >/dev/null || true
    fi
}
trap cleanup EXIT

cargo pgrx start "${PG_VERSION}"
cargo pgrx install -p context-pg --features "${PG_FEATURE}" --pg-config "${PG_CONFIG}"

psql -h "${PGHOST}" -p "${PGPORT}" -d postgres -v ON_ERROR_STOP=1 \
    -c "DROP DATABASE IF EXISTS ${DBNAME}" \
    -c "CREATE DATABASE ${DBNAME}"

psql -h "${PGHOST}" -p "${PGPORT}" -d "${DBNAME}" -v ON_ERROR_STOP=1 <<'SQL'
CREATE EXTENSION pgcontext;

CREATE TABLE hnsw_restart_items (
    id bigint GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
    embedding vector NOT NULL
);

INSERT INTO hnsw_restart_items (embedding)
VALUES ('[1,1,1]'::vector), ('[2,2,2]'::vector), ('[3,3,3]'::vector);

CREATE INDEX hnsw_restart_items_embedding_idx
    ON hnsw_restart_items USING pgcontext_hnsw (embedding);

INSERT INTO hnsw_restart_items (embedding)
VALUES ('[4,4,4]'::vector);
SQL

restart_after_stop=true
"${PG_CTL}" -D "${PGDATA}" stop -m immediate
cargo pgrx start "${PG_VERSION}"
restart_after_stop=false

psql -h "${PGHOST}" -p "${PGPORT}" -d "${DBNAME}" -v ON_ERROR_STOP=1 <<'SQL'
DO $$
DECLARE
    visible_rows bigint;
    index_bytes bigint;
BEGIN
    SELECT count(*) INTO visible_rows FROM hnsw_restart_items;
    IF visible_rows <> 4 THEN
        RAISE EXCEPTION 'unexpected visible row count after HNSW restart: %', visible_rows;
    END IF;

    SELECT pg_relation_size('hnsw_restart_items_embedding_idx'::regclass)
    INTO index_bytes;
    IF index_bytes < current_setting('block_size')::bigint THEN
        RAISE EXCEPTION 'HNSW index relation is too small after restart: %', index_bytes;
    END IF;
END
$$;

INSERT INTO hnsw_restart_items (embedding)
VALUES ('[5,5,5]'::vector);

VACUUM hnsw_restart_items;

SELECT true AS restart_preserved_hnsw_index;
SQL