CREATE SCHEMA IF NOT EXISTS documentdb_distributed_test_helpers;

SELECT citus_set_coordinator_host('localhost', current_setting('port')::integer);
SELECT citus_set_node_property('localhost', current_setting('port')::integer, 'shouldhaveshards', true);

/* see the comment written for its definition at create_indexes.c */
CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.generate_create_index_arg(
    p_collection_name text,
    p_index_name text,
    p_index_key documentdb_core.bson)
RETURNS documentdb_core.bson LANGUAGE C STRICT AS 'pg_documentdb', $$generate_create_index_arg$$;

-- Returns the command (without "CONCURRENTLY" option) used to create given
-- documentdb index on given collection.
CREATE FUNCTION documentdb_distributed_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 mi,
         documentdb_api_catalog.collections mc,
         pg_indexes pi
    WHERE mc.database_name = p_database_name AND
          mc.collection_name = p_collection_name AND
          (mi.index_spec).index_name = p_index_name AND
          mi.collection_id = mc.collection_id AND
          pi.indexname = concat('documents_rum_index_', index_id::text) AND
          pi.schemaname = 'documentdb_data';
END;
$$
LANGUAGE plpgsql;


-- query documentdb_api_catalog.collection_indexes for given collection
CREATE OR REPLACE FUNCTION documentdb_distributed_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 mi.collection_id, mi.index_id,
         documentdb_api_internal.index_spec_as_bson(mi.index_spec, for_get_indexes=>true),
         mi.index_is_valid
  FROM documentdb_api_catalog.collection_indexes AS mi
  WHERE mi.collection_id = (SELECT mc.collection_id FROM documentdb_api_catalog.collections AS mc
                            WHERE collection_name = p_collection_name AND
                                  database_name = p_database_name);
END;
$$ LANGUAGE plpgsql;

-- query pg_index for the documents table backing given collection
CREATE OR REPLACE FUNCTION documentdb_distributed_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_distributed_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;

-- function to mask variable plan id from the explain output of a distributed subplan
CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.mask_plan_id_from_distributed_subplan(explain_command text, out query_plan text)
RETURNS SETOF TEXT AS $$
BEGIN
  FOR query_plan IN EXECUTE explain_command LOOP
    IF query_plan ILIKE '%Distributed Subplan %_%'
    THEN
      RETURN QUERY SELECT REGEXP_REPLACE(query_plan,'[[:digit:]]+','X', 'g');
    ELSE
      RETURN next;
    END IF;
  END LOOP;
  RETURN;
END; $$ language plpgsql;


CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.run_explain_and_trim(
    p_query text,
    p_ignore_heap_fetches boolean DEFAULT false,
    p_ignore_distributed_runtime_details 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 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 p_ignore_distributed_runtime_details AND v_explain_row ~ 'Intermediate Data Size: [0-9]+ bytes' THEN
      SELECT regexp_replace(v_explain_row, 'Intermediate Data Size: [0-9]+ bytes', 'Intermediate Data Size: xxx bytes')
      INTO v_explain_row;
    ELSIF p_ignore_distributed_runtime_details AND v_explain_row ~ 'Tuple data received from (node|nodes): [0-9]+ bytes' THEN
      SELECT regexp_replace(v_explain_row,
                           'Tuple data received from (node|nodes): [0-9]+ bytes',
                           'Tuple data received from \1: xxx bytes')
      INTO v_explain_row;
    ELSIF p_ignore_distributed_runtime_details AND 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;
    ELSIF p_ignore_distributed_runtime_details AND v_explain_row ~ 'Seq Scan on documentdb_data\.documents_[0-9]+_[0-9]+ collection \(actual rows=[0-9\.]+ loops=[0-9]+\)' THEN
      SELECT regexp_replace(v_explain_row,
                           'Seq Scan on documentdb_data\.documents_[0-9]+_[0-9]+ collection \(actual rows=[0-9\.]+ loops=([0-9]+)\)',
                           'Seq Scan on documentdb_data.documents_x_x collection (actual rows=xyz loops=\1)')
      INTO v_explain_row;
    ELSIF p_ignore_distributed_runtime_details AND v_explain_row ~ 'documents_[0-9]+_[0-9]+' THEN
      SELECT regexp_replace(v_explain_row, 'documents_[0-9]+_[0-9]+', 'documents_x_x', 'g') 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 ~ '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;
    END IF;

    IF 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;
    END IF;

    RETURN NEXT v_explain_row;
  END LOOP;
END
$$
LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION documentdb_distributed_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 database_name = p_database_name AND collection_name = p_collection_name;
    DELETE FROM documentdb_api_catalog.collection_indexes
    WHERE (index_spec).index_key operator(documentdb_core.=) '{"_id": 1}' 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;

-- Function to avoid flakiness of a SQL query typically on a sharded multi-node collection. 
-- One way to fix such falkiness is to add an order by clause to inject determinism, but 
-- many queryies like cursors don't support order by. This test function bridges that gap 
-- by storing the result of such queries in a TEMP table and then ordering the entries in the 
-- temp table. One caveat is that the sql query in the argument is expacted to have exact two
-- columns object_id, and document. This seems to be sufficient for now for our use cases.
-- If the caller wants to project multiple columns, thaey can be concatenated as aliased as 'document'
CREATE OR REPLACE FUNCTION execute_and_sort(p_sql TEXT)
RETURNS TABLE (document text) AS $$
BEGIN
    EXECUTE 'CREATE TEMP TABLE temp_dynamic_results ON COMMIT DROP AS ' || p_sql;
    RETURN QUERY EXECUTE 'SELECT document FROM temp_dynamic_results ORDER BY object_id';
    EXECUTE 'DROP TABLE temp_dynamic_results';
END;
$$ LANGUAGE plpgsql;

-- This method mimics how the 2d index extract the geometries from `p_document` from `p_keyPath`
-- This function expects the geospatial data in form of legacy coordinate pairs (longitude, latitude).
-- returns the 2d flat geometry in form of public.geometry.
--
-- This function does strict validation of the values at path for geometry formats and
-- checks for valid points and multipoints input format and throws
-- error if not valid and only is applicable for creating the geospatial index and control
-- insert behaviors for invalid geodetic data points.
--
-- example scenario with the reference implementation:
-- - db.coll.createIndex({loc: "2dsphere"});
-- 
-- - db.insert({loc: [10, 'text']}); => This throws error
-- 
-- - db.insert({non-loc: [10, 'text']}) => This is normal insert as no 2d index
CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.bson_extract_geometry(
    p_document documentdb_core.bson,
    p_keyPath text)
 RETURNS public.geometry
 LANGUAGE c
 IMMUTABLE PARALLEL SAFE STRICT 
AS 'pg_documentdb', $function$bson_extract_geometry$function$;


-- This method mimics how the runtime extract the geometries from `p_document` from `p_keyPath`
-- This is similar to bson_extract_geometry function but
-- it performs a `weak` validation and doesn't throw error in case where the `bson_extract_geometry` function may throw error
-- e.g. scenarios with the reference implementation:
-- - db.coll.insert({loc: [[10, 20], [30, 40], ["invalid"]]}); (without 2d index on 'loc')
-- - db.coll.find({loc: {$geoWithin: { $box: [[30, 30], [40, 40]] }}})
--
-- The above find should match the object if any of the point (in multikey point case) matches the
-- geospatial query.
CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.bson_extract_geometry_runtime(
    p_document documentdb_core.bson,
    p_keyPath text)
 RETURNS public.geometry
 LANGUAGE c
 IMMUTABLE PARALLEL SAFE STRICT
AS 'pg_documentdb', $function$bson_extract_geometry_runtime$function$;

-- 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_distributed_test_helpers.create_indexes_background(IN p_database_name text, 
                                                        IN p_index_spec documentdb_core.bson,
                                                        IN p_log_index_queue boolean DEFAULT false,
                                                        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;
  index_cmd_stored text;
  attempt_count int := 0;
  max_attempts int := 600;
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);
  IF p_log_index_queue THEN
    SELECT string_agg(index_cmd, ',') into index_cmd_stored FROM documentdb_api_catalog.documentdb_index_queue;
    RAISE INFO 'Index Queue Commands: %', index_cmd_stored;
  END IF;
  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

          -- Background index builds can take longer on slower CI runners.
          -- Allow up to 60 seconds before treating the build as hung.
          IF attempt_count >= max_attempts THEN
            SELECT string_agg(index_cmd || index_cmd_status || comment || attempt || update_time, ',') into index_cmd_stored FROM documentdb_api_catalog.documentdb_index_queue;
            RAISE INFO 'Index Queue Commands: %', index_cmd_stored;
            SELECT string_agg(index_id || ':' || collection_id || ':' || index_is_valid, ',') into index_cmd_stored FROM documentdb_api_catalog.collection_indexes;
            RAISE INFO 'Collection Indexes: %', index_cmd_stored;
            RAISE EXCEPTION 'Waited too long for index build to complete. Last response from check_build_index_status: %', check_build_index_status;
          END IF;
          PERFORM pg_sleep_for('100 ms');
          attempt_count := attempt_count + 1;
      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;


-- This is a helper function to evaluate expressions for testing purposes.
-- This is used by backend tests to validate functionality of comparisons.
CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.evaluate_query_expression(expression documentdb_core.bson, value documentdb_core.bson)
 RETURNS bool
 LANGUAGE c
 IMMUTABLE STRICT
AS '$libdir/pg_documentdb.so', $function$command_evaluate_query_expression$function$;

CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.evaluate_expression_get_first_match(expression documentdb_core.bson, value documentdb_core.bson)
 RETURNS documentdb_core.bson
 LANGUAGE c
 IMMUTABLE STRICT
AS '$libdir/pg_documentdb.so', $function$command_evaluate_expression_get_first_match$function$;

-- 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_distributed_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_distributed_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 a single table in the 'db' database so that existing tests don't change behavior (yet)
set documentdb.enableNativeColocation to off;
SELECT documentdb_api.create_collection('db', 'firstCollection');
set documentdb.enableNativeColocation to on;
SELECT documentdb_api.create_collection('db', 'secondCollection');

CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.gin_bson_get_single_path_generated_terms(
    document documentdb_core.bson,
    path text,
    isWildcard bool,
    generateNotFoundTerm bool default false,
    addMetadata bool default false,
    indexTermSizeLimit int default -1,
    enableReducedWildcardTerms bool default false)
 RETURNS SETOF documentdb_core.bson
 LANGUAGE c
 IMMUTABLE PARALLEL SAFE STRICT ROWS 100
AS 'pg_documentdb', $$gin_bson_get_single_path_generated_terms$$;

CREATE OR REPLACE FUNCTION documentdb_distributed_test_helpers.gin_bson_get_wildcard_project_generated_terms(
    document documentdb_core.bson,
    pathSpec text,
    isExclusion bool,
    includeId bool,
    addMetadata bool default false,
    indexTermSizeLimit int default -1)
 RETURNS SETOF documentdb_core.bson
 LANGUAGE c
 IMMUTABLE PARALLEL SAFE STRICT ROWS 100
AS 'pg_documentdb', $$gin_bson_get_wildcard_project_generated_terms$$;

CREATE FUNCTION documentdb_distributed_test_helpers.get_feature_counter_pretty(p_reset_counter bool)
RETURNS SETOF json
AS
$$
BEGIN
    RETURN QUERY
    SELECT row_to_json(result) FROM ( 
        SELECT coalesce(json_agg(json_build_object(feature_name, usage_count)), '[]'::json) AS "Feature_usage" 
        FROM documentdb_api_internal.command_feature_counter_stats(p_reset_counter)
    ) result;
END;
$$
LANGUAGE plpgsql;

-- Waits until the global xmin horizon advances past the current transaction,
-- ensuring that a subsequent VACUUM can see recently-deleted tuples as removable.
CREATE OR REPLACE PROCEDURE documentdb_distributed_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_distributed_test_helpers.change_index_jobs_status(active_status boolean)
RETURNS void
AS $$
DECLARE
    job_id integer;
    index_row_stored text;
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;
        SELECT row_to_json(ROW(jobid,schedule,command,active,jobname))::text into index_row_stored FROM cron.job WHERE jobid = job_id;
        RAISE INFO 'Processing job_id: % state: %', job_id, index_row_stored;
    END LOOP;
END;
$$
LANGUAGE plpgsql;
