CREATE SCHEMA IF NOT EXISTS documentdb_test_helpers;

SELECT MIN(datcollate), MIN(datctype), MIN(pg_encoding_to_char(encoding)), MIN(datlocprovider) FROM pg_database;
SELECT MAX(datcollate), MAX(datctype), MAX(pg_encoding_to_char(encoding)), MAX(datlocprovider) FROM pg_database;

-- binary version should return the installed version after recreating the extension
SELECT documentdb_api.binary_version() = (SELECT REPLACE(extversion, '-', '.') FROM pg_extension where extname = 'documentdb_core');

-- Wait for the background worker to be launched in the `regression` database
-- When the extension is loaded, this isn't created yet. 
CREATE OR REPLACE PROCEDURE documentdb_test_helpers.wait_for_background_worker()
AS $$
DECLARE 
  v_bg_worker_app_name text := NULL;
BEGIN
  LOOP
    SELECT application_name INTO v_bg_worker_app_name FROM pg_stat_activity WHERE application_name = 'documentdb_bg_worker_leader';
    IF v_bg_worker_app_name IS NOT NULL THEN
      RETURN;
    END IF;

    COMMIT; -- This is needed so that we grab a fresh snapshot of pg_stat_activity
    PERFORM pg_sleep_for('100 ms');
  END LOOP;
END
$$
LANGUAGE plpgsql;

CALL documentdb_test_helpers.wait_for_background_worker();

-- validate background worker is launched
SELECT application_name FROM pg_stat_activity WHERE application_name = 'documentdb_bg_worker_leader';


CREATE OR REPLACE FUNCTION documentdb_test_helpers.run_explain_and_trim(p_query text, p_ignore_heap_fetches boolean DEFAULT false, p_normalize_window boolean DEFAULT false)
RETURNS SETOF text
AS $$
DECLARE
  v_explain_row text;
BEGIN
  FOR v_explain_row IN EXECUTE p_query
  LOOP
    IF v_explain_row ~ '^\s+Disabled: true\s*$' THEN
      CONTINUE;
    ELSIF p_normalize_window AND v_explain_row ~ '^\s+Window: ' THEN
      CONTINUE;
    ELSIF v_explain_row ~ '^\s+Index Searches: [0-9]+\s*$' THEN
      CONTINUE;
    ELSIF p_ignore_heap_fetches AND v_explain_row ~ '^\s+Heap Fetches: [0-9]+\s*$' THEN
      SELECT regexp_replace(v_explain_row, 'Heap Fetches: [0-9]+', 'Heap Fetches: xxx') INTO v_explain_row;
    ELSIF v_explain_row ~ 'Parallel Index Scan using .+ on documents_[0-9]+ collection \(actual rows=[0-9\.]+ loops=[0-9]+\)' THEN
      SELECT regexp_replace(v_explain_row, 'Parallel Index Scan using (.+) on documents_([0-9]+) collection \(actual rows=[0-9\.]+ loops=([0-9]+)\)',
                                           'Parallel Index Scan using \1 on documents_\2 collection (actual rows=xyz loops=\3)') INTO v_explain_row;
    ELSIF v_explain_row ~ 'actual rows=[0-9]+\.00' THEN
      SELECT regexp_replace(v_explain_row, 'actual rows=([0-9]+)\.00', 'actual rows=\1') INTO v_explain_row;
    ELSIF v_explain_row ~ 'Sort Method: quicksort  Memory: [0-9]+kB' THEN
      SELECT regexp_replace(v_explain_row, 'Sort Method: quicksort  Memory: [0-9]+kB', 'Sort Method: quicksort  Memory: xxxkB') INTO v_explain_row;
    ELSIF v_explain_row ~ 'Sort Method: quicksort  Average Memory: [0-9]+kB  Peak Memory: [0-9]+kB' THEN
      SELECT regexp_replace(v_explain_row, 'Sort Method: quicksort  Average Memory: [0-9]+kB  Peak Memory: [0-9]+kB', 'Sort Method: quicksort  Average Memory: xxxkB  Peak Memory: xxxkB') INTO v_explain_row;
    ELSIF v_explain_row ~ 'Memory Usage: [0-9]+kB' THEN
      SELECT regexp_replace(v_explain_row, 'Memory Usage: [0-9]+kB', 'Memory Usage: xxxkB') INTO v_explain_row;
    END IF;
    -- Normalize PG18 window function references (OVER w1 -> OVER (?))
    IF p_normalize_window AND v_explain_row ~ ' OVER w[0-9]+' THEN
      SELECT regexp_replace(v_explain_row, ' OVER w[0-9]+', ' OVER (?)', 'g') INTO v_explain_row;
    END IF;
    RETURN NEXT v_explain_row;
  END LOOP;
END
$$
LANGUAGE plpgsql;

-- query documentdb_api_catalog.collection_indexes for given collection
CREATE OR REPLACE FUNCTION documentdb_test_helpers.get_collection_indexes(
    p_database_name text,
    p_collection_name text,
    OUT collection_id bigint,
    OUT index_id integer,
    OUT index_spec_as_bson documentdb_core.bson,
    OUT index_is_valid bool)
RETURNS SETOF RECORD
AS $$
BEGIN
  RETURN QUERY
  SELECT ci.collection_id, ci.index_id,
         documentdb_api_internal.index_spec_as_bson(ci.index_spec, for_get_indexes=>true),
         ci.index_is_valid
  FROM documentdb_api_catalog.collection_indexes AS ci
  WHERE ci.collection_id = (SELECT hc.collection_id FROM documentdb_api_catalog.collections AS hc
                            WHERE collection_name = p_collection_name AND
                                  database_name = p_database_name)
  ORDER BY ci.index_id;
END;
$$ LANGUAGE plpgsql;

-- query pg_index for the documents table backing given collection
CREATE OR REPLACE FUNCTION documentdb_test_helpers.get_data_table_indexes (
    p_database_name text,
    p_collection_name text)
RETURNS TABLE (LIKE pg_index)
AS $$
DECLARE
  v_collection_id bigint;
  v_data_table_name text;
BEGIN
  SELECT collection_id INTO v_collection_id
  FROM documentdb_api_catalog.collections
  WHERE collection_name = p_collection_name AND
        database_name = p_database_name;

  v_data_table_name := format('documentdb_data.documents_%s', v_collection_id);

  RETURN QUERY
  SELECT * FROM pg_index WHERE indrelid = v_data_table_name::regclass;
END;
$$ LANGUAGE plpgsql;

-- Returns the command (without "CONCURRENTLY" option) used to create given
-- index on a collection.
CREATE OR REPLACE FUNCTION documentdb_test_helpers.documentdb_index_get_pg_def(
    p_database_name text,
    p_collection_name text,
    p_index_name text)
RETURNS SETOF TEXT
AS
$$
BEGIN
    RETURN QUERY
    SELECT pi.indexdef
    FROM documentdb_api_catalog.collection_indexes hi,
         documentdb_api_catalog.collections hc,
         pg_indexes pi
    WHERE hc.database_name = p_database_name AND
          hc.collection_name = p_collection_name AND
          (hi.index_spec).index_name = p_index_name AND
          hi.collection_id = hc.collection_id AND
          pi.indexname = concat('documents_rum_index_', index_id::text) AND
          pi.schemaname = 'documentdb_data';
END;
$$
LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION documentdb_test_helpers.drop_primary_key(p_database_name text, p_collection_name text)
RETURNS void
AS $$
DECLARE
  v_collection_id bigint;
BEGIN
    SELECT collection_id INTO v_collection_id
    FROM documentdb_api_catalog.collections
    WHERE collection_name = p_collection_name AND
          database_name = p_database_name;

    DELETE FROM documentdb_api_catalog.collection_indexes
    WHERE (index_spec).index_key operator(documentdb_core.=) '{"_id": 1}'::documentdb_core.bson AND
          collection_id = v_collection_id;
	EXECUTE format('ALTER TABLE documentdb_data.documents_%s DROP CONSTRAINT collection_pk_%s', v_collection_id, v_collection_id);
END;
$$ LANGUAGE plpgsql;

-- This is a helper for create_indexes_background. It performs the submission of index requests in background and wait for their completion.
CREATE OR REPLACE PROCEDURE documentdb_test_helpers.create_indexes_background(IN p_database_name text,
                                                        IN p_index_spec documentdb_core.bson,
                                                        INOUT retVal documentdb_core.bson DEFAULT null,
                                                        INOUT ok boolean DEFAULT false)
AS $procedure$
DECLARE
  create_index_response record;
  check_build_index_status record;
  completed boolean := false;
  indexRequest text;
  attempt_count int := 0;
BEGIN
  SET search_path TO documentdb_core,documentdb_api;
  SELECT * INTO create_index_response FROM documentdb_api.create_indexes_background(p_database_name, p_index_spec);
  COMMIT;

  IF create_index_response.ok THEN
    SELECT create_index_response.requests->>'indexRequest' INTO indexRequest;
    IF indexRequest IS NOT NULL THEN
      LOOP
          SELECT * INTO check_build_index_status FROM documentdb_api_internal.check_build_index_status(create_index_response.requests);
          IF check_build_index_status.ok THEN
            completed := check_build_index_status.complete;
            IF completed THEN
              ok := create_index_response.ok;
              retVal := create_index_response.retval;
              RETURN;
            END IF;
          ELSE
            ok := check_build_index_status.ok;
            retVal := check_build_index_status.retval;
            RETURN;
          END IF;

          COMMIT; -- COMMIT so that CREATE INDEX CONCURRENTLY does not wait for documentdb_distributed_test_helpers.create_indexes_background
          PERFORM pg_sleep_for('100 ms');
          attempt_count := attempt_count + 1;

          -- don't wait longer than 10 seconds for the index build to complete. This avoids timeouts due to infinite waits.
          IF attempt_count > 100 THEN
            RAISE EXCEPTION 'Waited too long for index build to complete. Last response from check_build_index_status: %', check_build_index_status;
          END IF;
      END LOOP;
    ELSE
      ok := create_index_response.ok;
      retVal := create_index_response.retval;
      RETURN;
    END IF;
  ELSE
    ok := create_index_response.ok;
    retVal := create_index_response.retval;
  END IF;
END;
$procedure$
LANGUAGE plpgsql;

-- query pg_index for the documents table backing given collection
CREATE OR REPLACE FUNCTION documentdb_test_helpers.get_data_table_indexes (
    p_database_name text,
    p_collection_name text)
RETURNS TABLE (LIKE pg_index)
AS $$
DECLARE
  v_collection_id bigint;
  v_data_table_name text;
BEGIN
  SELECT collection_id INTO v_collection_id
  FROM documentdb_api_catalog.collections
  WHERE collection_name = p_collection_name AND
        database_name = p_database_name;

  v_data_table_name := format('documentdb_data.documents_%s', v_collection_id);

  RETURN QUERY
  SELECT * FROM pg_index WHERE indrelid = v_data_table_name::regclass;
END;
$$ LANGUAGE plpgsql;

-- count collection indexes grouping by "pg_index.indisprimary" attr
CREATE OR REPLACE FUNCTION documentdb_test_helpers.count_collection_indexes(
    p_database_name text,
    p_collection_name text)
RETURNS TABLE (
  index_type_is_primary boolean,
  index_type_count bigint
)
AS $$
BEGIN
  RETURN QUERY
  SELECT indisprimary, COUNT(*) FROM pg_index
  WHERE indrelid = (SELECT ('documentdb_data.documents_' || collection_id::text)::regclass
                    FROM documentdb_api_catalog.collections
                    WHERE database_name = p_database_name AND
                          collection_name = p_collection_name)
  GROUP BY indisprimary;
END;
$$ LANGUAGE plpgsql;

-- Wait until the database-wide xmin horizon advances past every transaction that
-- committed before this call so a subsequent VACUUM is guaranteed to see prior
-- DELETEs as fully dead and run ambulkdelete.
--
-- This is a procedure rather than a function because it must COMMIT between
-- iterations: otherwise the caller's own snapshot (and any xid this routine
-- allocates) would itself pin the horizon and the loop could never make
-- progress. We use pg_snapshot_xmax() as the target (an upper bound on every
-- already-committed xid) so we never have to allocate a new xid ourselves.
CREATE OR REPLACE PROCEDURE documentdb_test_helpers.wait_for_vacuum_horizon(p_timeout_ms int default 30000)
  LANGUAGE plpgsql
AS $$
DECLARE
  v_target xid8;
  v_deadline timestamptz := clock_timestamp() + (p_timeout_ms || ' ms')::interval;
  v_holders text;
  v_prepared_holders text;
  v_slot_holders text;
  v_horizon_ok boolean;
BEGIN
  v_target := pg_snapshot_xmax(pg_current_snapshot());
  COMMIT;
  LOOP
    -- VACUUM uses GetOldestNonRemovableTransactionId(), which (via
    -- ComputeXidHorizons) walks the ProcArray for both proc->xid AND
    -- proc->xmin, plus replication slot xmins, plus prepared xacts.
    --
    -- A read-only backend with proc->xid = InvalidTransactionId but
    -- proc->xmin set (e.g., a long-running SELECT in a background worker)
    -- pins VACUUM's horizon but is *invisible* to pg_snapshot_xmin(),
    -- which GetSnapshotData computes from proc->xid only. That asymmetry
    -- is the historical flake source -- the helper would return ✓ while
    -- VACUUM still saw an older OldestXmin and silently skipped
    -- ambulkdelete. Use pg_stat_activity.backend_xmin (which exposes
    -- proc->xmin) so we wait on the same holders VACUUM does. Also check
    -- backend_xid to defend against the narrow race PG itself defends
    -- against in ComputeXidHorizons (xmin = TransactionIdOlder(xmin, xid)):
    -- a writer whose proc->xmin is momentarily cleared between snapshots
    -- but whose proc->xid still pins the horizon.
    --
    -- We use age() for the xid comparisons (rather than ::text::bigint)
    -- to be wraparound-safe: backend_xmin / backend_xid / pg_prepared_xacts
    -- / pg_replication_slots all expose 32-bit xid which wraps every ~4B
    -- transactions. age() computes next_xid - xid in modulo arithmetic,
    -- so "older than v_target" is exactly age(x) > age(v_target::xid),
    -- which behaves correctly across wraparound boundaries.
    --
    -- ComputeXidHorizons's *data* horizon counts:
    --   * backends in the current database (proc->databaseId == MyDatabaseId)
    --   * backends with PROC_AFFECTS_ALL_HORIZONS set (physical walsenders
    --     doing hot-standby-feedback), which are not tied to any specific
    --     database and appear in pg_stat_activity with datname IS NULL
    -- so we scope the backend check to (current_database() OR NULL).
    SELECT NOT EXISTS (
             SELECT 1 FROM pg_stat_activity
              WHERE pid <> pg_backend_pid()
                AND (datname = current_database() OR datname IS NULL)
                AND (
                  (backend_xmin IS NOT NULL
                   AND age(backend_xmin) > age(v_target::text::xid))
                  OR
                  (backend_xid IS NOT NULL
                   AND age(backend_xid) > age(v_target::text::xid))
                )
           )
           AND NOT EXISTS (
             SELECT 1 FROM pg_prepared_xacts
              WHERE age(transaction) > age(v_target::text::xid)
                AND database = current_database()
           )
           AND NOT EXISTS (
             SELECT 1 FROM pg_replication_slots
              WHERE xmin IS NOT NULL
                AND age(xmin) > age(v_target::text::xid)
           )
      INTO v_horizon_ok;

    EXIT WHEN v_horizon_ok;

    IF clock_timestamp() > v_deadline THEN
      SELECT string_agg(format('pid=%s xid=%s xmin=%s state=%s backend_type=%s db=%s',
                               pid, backend_xid, backend_xmin, state, backend_type, datname), '; ')
        INTO v_holders
        FROM pg_stat_activity
       WHERE pid <> pg_backend_pid()
         AND (datname = current_database() OR datname IS NULL)
         AND (backend_xmin IS NOT NULL OR backend_xid IS NOT NULL);
      SELECT string_agg(format('gid=%s xid=%s db=%s',
                               gid, transaction, database), '; ')
        INTO v_prepared_holders
        FROM pg_prepared_xacts
       WHERE database = current_database();
      SELECT string_agg(format('slot=%s xmin=%s active=%s',
                               slot_name, xmin, active), '; ')
        INTO v_slot_holders
        FROM pg_replication_slots
       WHERE xmin IS NOT NULL;
      RAISE EXCEPTION 'wait_for_vacuum_horizon: xmin horizon did not advance to % within %ms (backends: %; prepared: %; slots: %)',
        v_target, p_timeout_ms,
        COALESCE(v_holders, '<none>'),
        COALESCE(v_prepared_holders, '<none>'),
        COALESCE(v_slot_holders, '<none>');
    END IF;
    PERFORM pg_sleep(0.05);
    COMMIT;
  END LOOP;
END;
$$;

CREATE OR REPLACE FUNCTION documentdb_test_helpers.change_index_jobs_status(active_status boolean)
RETURNS void
AS $$
DECLARE
    job_id integer;
BEGIN
    FOR job_id IN (SELECT jobid FROM cron.job WHERE jobname LIKE 'documentdb_index_%' order by jobid)
    LOOP
        UPDATE cron.job SET active = active_status WHERE jobid = job_id;
        RAISE NOTICE 'Processing job_id: %', job_id;
    END LOOP;
END;
$$
LANGUAGE plpgsql;
