// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use pgrx::pg_sys::panic::ErrorReport;
use pgrx::{
    function_name, pg_sys, GucContext, GucFlags, GucRegistry, GucSetting, PgLogLevel,
    PgSqlErrorCode,
};
use std::ffi::CStr;
use std::num::NonZeroUsize;
use tantivy::aggregation::DEFAULT_BUCKET_LIMIT;

use crate::postgres::options::MAX_MUTABLE_SEGMENT_ROWS;

/// Allows the user to toggle the use of our "ParadeDB Base Scan".
static ENABLE_CUSTOM_SCAN: GucSetting<bool> = GucSetting::<bool>::new(true);

/// Allows the user to toggle the use of our "ParadeDB Aggregate Scan".
static ENABLE_AGGREGATE_CUSTOM_SCAN: GucSetting<bool> = GucSetting::<bool>::new(false);

/// Allows the user to toggle the use of our "ParadeDB Join Scan".
static ENABLE_JOIN_CUSTOM_SCAN: GucSetting<bool> = GucSetting::<bool>::new(false);

/// Allows the user to toggle the use of the custom scan without use of the `@@@` operator. The
/// default is `false`.
static ENABLE_CUSTOM_SCAN_WITHOUT_OPERATOR: GucSetting<bool> = GucSetting::<bool>::new(false);

/// Allows the user to toggle the use of custom scan for queries that include non-indexed fields.
/// When enabled, queries with non-indexed predicates will use HeapExpr for heap filtering.
static ENABLE_FILTER_PUSHDOWN: GucSetting<bool> = GucSetting::<bool>::new(true);

/// Allows the user to enable or disable the FastFieldsExecState executor. Default is `true`.
static ENABLE_FAST_FIELD_EXEC: GucSetting<bool> = GucSetting::<bool>::new(true);

/// Allows the user to enable or disable the ColumnarExecState executor. Default is `true`.
static ENABLE_COLUMNAR_EXEC: GucSetting<bool> = GucSetting::<bool>::new(true);

/// Allows the user to enable or disable sorted execution for ColumnarExecState.
/// When disabled, sorted paths will not be created even if the index has sort_by.
static ENABLE_COLUMNAR_SORT: GucSetting<bool> = GucSetting::<bool>::new(true);

/// In a Top K query, the limit is multiplied by this factor to determine the chunk size.
static LIMIT_FETCH_MULTIPLIER: GucSetting<f64> = GucSetting::<f64>::new(1.0);

/// The scale factor for the chunk size in a Top K query.
static TOPK_RETRY_SCALE_FACTOR: GucSetting<i32> = GucSetting::<i32>::new(2);

/// The maximum chunk size for a Top K query.
static MAX_TOPK_CHUNK_SIZE: GucSetting<i32> = GucSetting::<i32>::new(100_000);

/// The maximum number of buckets that can be returned by a TermsAggregation
static MAX_TERM_AGG_BUCKETS: GucSetting<i32> = GucSetting::<i32>::new(DEFAULT_BUCKET_LIMIT as i32);

/// The maximum response size in bytes for a window aggregate.
static MAX_WINDOW_AGGREGATE_RESPONSE_BYTES: GucSetting<i32> = GucSetting::<i32>::new(1_048_576);

/// For testing, ensures the same handling of null aggregates as Postgres
static ADD_DOC_COUNT_TO_AGGS: GucSetting<bool> = GucSetting::<bool>::new(false);

/// The number of fast-field columns below-which the ColumnarExecState will be used, rather
/// than the NormalExecState. The Columnar execution mode fetches data as column-oriented, whereas
/// the Normal mode fetches data as row-oriented.
///
/// Each fetch from a fast-field column costs one or two disk seeks, whereas a fetch of a row
/// generally costs one. But with a wide enough row, fetching multiple columns might still result
/// in better cache performance than fetching a row.
static COLUMNAR_EXEC_COLUMN_THRESHOLD: GucSetting<i32> = GucSetting::<i32>::new(3);
static COLUMNAR_EXEC_COLUMN_THRESHOLD_NAME: &CStr = c"paradedb.columnar_exec_column_threshold";

/// The `PER_TUPLE_COST` is an arbitrary value that needs to be really high.  In fact, we default
/// to one hundred million.
///
/// The reason for this is we really do not want Postgres to choose a plan where the `@@@` operator
/// is used in a sequential scan, filter, or recheck condition... unless of course there's no
/// other way to solve the query.
///
/// This value is a multiplier that Postgres applies to the estimated row count any given `@@@`
/// query clause will return.  In our case, higher is better.
///
/// Our IAM impl has its own costing functions that don't use this GUC and provide sensible estimates
/// for the overall IndexScan.  That plus this help to persuade Postgres to use our IAM whenever
/// it logically can.
static PER_TUPLE_COST: GucSetting<f64> = GucSetting::<f64>::new(100_000_000.0);

static GLOBAL_TARGET_SEGMENT_COUNT: GucSetting<i32> = GucSetting::<i32>::new(0);
static GLOBAL_ENABLE_BACKGROUND_MERGING: GucSetting<bool> = GucSetting::<bool>::new(true);
static GLOBAL_MUTABLE_SEGMENT_ROWS: GucSetting<i32> = GucSetting::<i32>::new(-1);
static EXPLAIN_RECURSIVE_ESTIMATES: GucSetting<bool> = GucSetting::<bool>::new(false);

/// Validate Top K scan eligibility for LIMIT queries
static CHECK_TOPK_SCAN: GucSetting<bool> = GucSetting::<bool>::new(true);

/// When true, queries with expensive scorer construction (fuzzy, regex, range)
/// use a cheap heuristic for selectivity estimation instead of building a full Tantivy scorer.
static ENABLE_HEURISTIC_SELECTIVITY: GucSetting<bool> = GucSetting::<bool>::new(true);

/// Minimum number of rows per parallel worker.
/// Controls how many workers are spawned based on estimated row count.
/// Based on benchmarks, the crossover point where parallel becomes beneficial
/// is around 300K rows total. Setting to 300K ensures tables under ~300K rows
/// won't use parallel, and larger tables scale workers appropriately.
static MIN_ROWS_PER_WORKER: GucSetting<i32> = GucSetting::<i32>::new(300000);

/// Override the scanner batch size when dynamic filters are pushed down.
/// 0 means disabled (use the scanner's default). When > 0, the scanner's batch
/// size is capped to this value during filter pushdown so that Top K can tighten
/// its threshold between batches.
static DYNAMIC_FILTER_BATCH_SIZE: GucSetting<i32> = GucSetting::<i32>::new(0);

/// Allows the user to enable or disable the SegmentedTopK optimization.
/// When enabled, Top K queries on deferred (late-materialized) string/bytes columns
/// use per-segment ordinal pruning to reduce dictionary decoding.
static ENABLE_SEGMENTED_TOPK: GucSetting<bool> = GucSetting::<bool>::new(true);

pub fn init() {
    // Note that Postgres is very specific about the naming convention of variables.
    // They must be namespaced... we use 'paradedb.<variable>' below.

    GucRegistry::define_bool_guc(
        c"paradedb.enable_custom_scan",
        c"Enable ParadeDB's custom scan",
        c"Enable ParadeDB's custom scan, which replaces table scans with ParadeDB index scans in cases where beneficial",
        &ENABLE_CUSTOM_SCAN,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_aggregate_custom_scan",
        c"Enable ParadeDB's custom aggregate scan",
        c"Enable ParadeDB's custom aggregate scan, which replaces row-based aggregates with column-based aggregates where beneficial",
        &ENABLE_AGGREGATE_CUSTOM_SCAN,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_join_custom_scan",
        c"Enable ParadeDB's experimental join custom scan",
        c"Enable ParadeDB's experimental join custom scan. Default is false.",
        &ENABLE_JOIN_CUSTOM_SCAN,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_custom_scan_without_operator",
        c"Enable ParadeDB's custom scan to run without the `@@@` operator",
        c"Enable ParadeDB's custom scan to run even when the `@@@` operator has not been used in a query, as long as the entire WHERE clause is able to be pushed down",
        &ENABLE_CUSTOM_SCAN_WITHOUT_OPERATOR,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_filter_pushdown",
        c"Enable ParadeDB's custom scan for queries with non-indexed fields",
        c"Enable ParadeDB's custom scan to handle queries that include non-indexed field predicates using HeapExpr for heap filtering. When disabled, such queries will fall back to standard PostgreSQL execution",
        &ENABLE_FILTER_PUSHDOWN,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_fast_field_exec",
        c"Enable StringFastFieldsExecState and NumericFastFieldsExecState executor",
        c"Enable the StringFastFieldsExecState and NumericFastFieldsExecState executors for handling one string fast field or multiple numeric fast fields",
        &ENABLE_FAST_FIELD_EXEC,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_columnar_exec",
        c"Enable ColumnarExecState executor",
        c"Enable the ColumnarExecState executor for handling multiple string fast fields or mixed string/numeric fast fields",
        &ENABLE_COLUMNAR_EXEC,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_columnar_sort",
        c"Enable sorted execution for ColumnarExecState",
        c"Enable sorted execution for ColumnarExecState when the index has sort_by and the query ORDER BY matches the prefix. Disabling this forces unsorted execution.",
                &ENABLE_COLUMNAR_SORT,
                GucContext::Userset,
                GucFlags::default(),
            );

    GucRegistry::define_int_guc(
                COLUMNAR_EXEC_COLUMN_THRESHOLD_NAME,        c"Threshold of fetched columns below which ColumnarExecState will be used.",
        c"The number of fast-field columns below-which the ColumnarExecState will be used, rather \
         than the NormalExecState. The Columnar execution mode fetches data as column-oriented, whereas \
         the Normal mode fetches data as row-oriented.",
        &COLUMNAR_EXEC_COLUMN_THRESHOLD,
        0,
        i32::MAX,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_float_guc(
        c"paradedb.per_tuple_cost",
        c"Arbitrary multiplier for the cost of retrieving a tuple from a USING bm25 index outside of an IndexScan",
        c"Default is 100,000,000.0.  It is very expensive to use a USING bm25 index in the wrong query plan",
        &PER_TUPLE_COST,
        0.0,
        f64::MAX,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_float_guc(
        c"paradedb.limit_fetch_multiplier",
        c"Multiplier for the limit in a Top K query",
        c"The limit is multiplied by this factor to determine the chunk size. A higher value reduces the probability of a re-query for a Top K query but increases query times.",
        &LIMIT_FETCH_MULTIPLIER,
        1.0,
        100.0,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.max_topk_chunk_size",
        c"Maximum chunk size for a Top K query",
        c"A higher value reduces the probability of a re-query for a Top K query but increases the memory usage",
        &MAX_TOPK_CHUNK_SIZE,
        1,
        1_000_000,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.topk_retry_scale_factor",
        c"Scale factor for the chunk size in a Top K query",
        c"The chunk size is multiplied by this factor on subsequent retries. A higher value reduces the probability of a re-query for a Top K query but increases query times.",
        &TOPK_RETRY_SCALE_FACTOR,
        1,
        100,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.max_term_agg_buckets",
        c"Maximum number of buckets/groups that can be returned by a terms aggregation",
        c"Mostly used for testing. If this number is exceeded, that means the result could be truncated and the query will be cancelled.",
        &MAX_TERM_AGG_BUCKETS,
        1,
        DEFAULT_BUCKET_LIMIT as i32,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.max_window_aggregate_response_bytes",
        c"Maximum response size in bytes for a window aggregate.",
        c"The maximum response size in bytes for a window aggregate during a parallel scan. If this is exceeded, the query will be cancelled.",
        &MAX_WINDOW_AGGREGATE_RESPONSE_BYTES,
        1024,
        i32::MAX,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.global_target_segment_count",
        c"a global target segment count override",
        c"Setting this to a non-zero value ignores the `target_segment_count` property on all indexes in favor of this value",
        &GLOBAL_TARGET_SEGMENT_COUNT,
        0,
        8192,
        GucContext::Sighup,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.global_enable_background_merging",
        c"Enable background merging",
        c"Enable background merging",
        &GLOBAL_ENABLE_BACKGROUND_MERGING,
        GucContext::Sighup,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.explain_recursive_estimates",
        c"Enable recursive estimates in EXPLAIN VERBOSE",
        c"Shows estimated document counts for nested query components. Expensive operation, use for debugging only.",
        &EXPLAIN_RECURSIVE_ESTIMATES,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.global_mutable_segment_rows",
        c"a global mutable segment rows override",
        c"Setting this to a non-negative value ignores the `mutable_segment_rows` property on all indexes in favor of this value",
        &GLOBAL_MUTABLE_SEGMENT_ROWS,
        -1,
        MAX_MUTABLE_SEGMENT_ROWS as i32,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.add_doc_count_to_aggs",
        c"for testing, ensures the same handling of null aggregates as Postgres",
        c"Meant for internal testing usage",
        &ADD_DOC_COUNT_TO_AGGS,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.check_topk_scan",
        c"Validate Top K scan eligibility for LIMIT queries",
        c"When enabled, logs a warning if a query with LIMIT cannot use Top K scan. \
          This helps detect performance issues during development where queries expected \
          to use Top K optimization fall back to slower execution methods.",
        &CHECK_TOPK_SCAN,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_heuristic_selectivity",
        c"Use heuristic selectivity for expensive query types",
        c"When enabled, fuzzy, regex, and range queries use a cheap heuristic for planner selectivity estimation instead of constructing a full Tantivy scorer. Default is true.",
        &ENABLE_HEURISTIC_SELECTIVITY,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.min_rows_per_worker",
        c"Minimum rows per parallel worker",
        c"Controls how many parallel workers are used based on estimated row count. \
          Workers are limited so each processes at least this many rows. \
          Set to 0 to disable this check and use segment-based parallelism only.",
        &MIN_ROWS_PER_WORKER,
        0,
        i32::MAX,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_bool_guc(
        c"paradedb.enable_segmented_topk",
        c"Enable SegmentedTopK optimization for Top K queries on deferred columns",
        c"When enabled, ORDER BY on a late-materialized string/bytes column with LIMIT \
          uses per-segment ordinal pruning to reduce dictionary decoding. \
          All input is collected before emitting (EmissionType::Final) so only \
          the exact Top K rows per segment are sent to dictionary decoding. \
          Per-segment thresholds are published progressively to the scanner \
          for early row pruning during collection.",
        &ENABLE_SEGMENTED_TOPK,
        GucContext::Userset,
        GucFlags::default(),
    );

    GucRegistry::define_int_guc(
        c"paradedb.dynamic_filter_batch_size",
        c"Scanner batch size override for dynamic filter pushdown",
        c"When > 0, caps the scanner batch size during dynamic filter pushdown so that \
          Top K can tighten its threshold between batches. 0 disables the override.",
        &DYNAMIC_FILTER_BATCH_SIZE,
        0,
        128_000,
        GucContext::Userset,
        GucFlags::default(),
    );
}

pub fn enable_custom_scan() -> bool {
    ENABLE_CUSTOM_SCAN.get()
}

pub fn enable_aggregate_custom_scan() -> bool {
    ENABLE_AGGREGATE_CUSTOM_SCAN.get()
}

pub fn enable_join_custom_scan() -> bool {
    ENABLE_JOIN_CUSTOM_SCAN.get()
}

pub fn enable_custom_scan_without_operator() -> bool {
    ENABLE_CUSTOM_SCAN_WITHOUT_OPERATOR.get()
}

pub fn enable_filter_pushdown() -> bool {
    ENABLE_FILTER_PUSHDOWN.get()
}

pub fn is_fast_field_exec_enabled() -> bool {
    ENABLE_FAST_FIELD_EXEC.get()
}

pub fn is_columnar_exec_enabled() -> bool {
    ENABLE_COLUMNAR_EXEC.get()
}

pub fn is_columnar_sort_enabled() -> bool {
    ENABLE_COLUMNAR_SORT.get()
}

pub fn columnar_exec_column_threshold() -> usize {
    COLUMNAR_EXEC_COLUMN_THRESHOLD
        .get()
        .try_into()
        .unwrap_or_else(|e| {
            panic!(
                "{} must be positive. {e}",
                COLUMNAR_EXEC_COLUMN_THRESHOLD_NAME.to_str().unwrap()
            );
        })
}

pub fn per_tuple_cost() -> f64 {
    PER_TUPLE_COST.get()
}

pub fn max_topk_chunk_size() -> i32 {
    MAX_TOPK_CHUNK_SIZE.get()
}

pub fn global_target_segment_count() -> i32 {
    GLOBAL_TARGET_SEGMENT_COUNT.get()
}

pub fn global_enable_background_merging() -> bool {
    GLOBAL_ENABLE_BACKGROUND_MERGING.get()
}

// NB:  These limits come from [`tantivy::index_writer::MEMORY_BUDGET_NUM_BYTES_MAX`], which is not publicly exposed
mod limits {
    const MARGIN_IN_BYTES: usize = 1_000_000;
    // Size of the margin for the `memory_arena`. A segment is closed when the remaining memory
    // in the `memory_arena` goes below MARGIN_IN_BYTES.
    pub const MEMORY_BUDGET_NUM_BYTES_MIN: usize = 15 * MARGIN_IN_BYTES;

    // We impose the memory per thread to be no greater than 4GB as that's tantivy's limit
    pub const MEMORY_BUDGET_NUM_BYTES_MAX: usize = (4 * 1024 * 1024 * 1024) - MARGIN_IN_BYTES;
}

pub fn adjust_maintenance_work_mem(nlaunched: usize) -> NonZeroUsize {
    let nlaunched = nlaunched.max(1);
    let mwm_as_bytes = unsafe { pg_sys::maintenance_work_mem as usize } * 1024;
    let per_worker_budget = mwm_as_bytes / nlaunched;

    if per_worker_budget < limits::MEMORY_BUDGET_NUM_BYTES_MIN {
        ErrorReport::new(
            PgSqlErrorCode::ERRCODE_INSUFFICIENT_RESOURCES,
            "`maintenance_work_mem` is not high enough to give each parallel worker 15MB",
            function_name!(),
        )
        .set_detail(format!("this query asked for {nlaunched} workers, so `maintenance_work_mem` must be at least {nlaunched} * 15MB"))
        .set_hint("`SET maintenance_work_mem = <number>`")
        .report(PgLogLevel::ERROR);
    } else {
        pgrx::debug1!(
            "adjust_maintenance_work_mem: per_worker_budget: {per_worker_budget}, minimum: {}",
            limits::MEMORY_BUDGET_NUM_BYTES_MIN
        );
    }

    // clamp the per_thread_budget to the min/max values
    let per_worker_budget = per_worker_budget.clamp(
        limits::MEMORY_BUDGET_NUM_BYTES_MIN,
        limits::MEMORY_BUDGET_NUM_BYTES_MAX - 1,
    );

    NonZeroUsize::new(per_worker_budget * nlaunched).unwrap()
}

pub fn adjust_work_mem() -> NonZeroUsize {
    let wm_as_bytes = unsafe { pg_sys::work_mem as usize * 1024 };
    let wm_as_bytes = wm_as_bytes.clamp(
        limits::MEMORY_BUDGET_NUM_BYTES_MIN,
        limits::MEMORY_BUDGET_NUM_BYTES_MAX - 1,
    );

    NonZeroUsize::new(wm_as_bytes).unwrap()
}

pub fn limit_fetch_multiplier() -> f64 {
    LIMIT_FETCH_MULTIPLIER.get()
}

pub fn max_term_agg_buckets() -> i32 {
    let v = MAX_TERM_AGG_BUCKETS.get();
    if v <= 0 {
        pgrx::error!("paradedb.max_term_agg_buckets must be a positive integer");
    }
    v
}

pub fn max_window_aggregate_response_bytes() -> usize {
    MAX_WINDOW_AGGREGATE_RESPONSE_BYTES.get() as usize
}

pub fn topk_retry_scale_factor() -> i32 {
    TOPK_RETRY_SCALE_FACTOR.get()
}

pub fn global_mutable_segment_rows() -> Option<usize> {
    let value = GLOBAL_MUTABLE_SEGMENT_ROWS.get();
    if value >= 0 {
        Some(value as usize)
    } else {
        None
    }
}

pub fn explain_recursive_estimates() -> bool {
    EXPLAIN_RECURSIVE_ESTIMATES.get()
}

pub fn check_topk_scan() -> bool {
    CHECK_TOPK_SCAN.get()
}

pub fn enable_heuristic_selectivity() -> bool {
    ENABLE_HEURISTIC_SELECTIVITY.get()
}

pub fn min_rows_per_worker() -> i32 {
    MIN_ROWS_PER_WORKER.get()
}

pub fn add_doc_count_to_aggs() -> bool {
    ADD_DOC_COUNT_TO_AGGS.get()
}

pub fn dynamic_filter_batch_size() -> i32 {
    DYNAMIC_FILTER_BATCH_SIZE.get()
}

pub fn enable_segmented_topk() -> bool {
    ENABLE_SEGMENTED_TOPK.get()
}

#[cfg(any(test, feature = "pg_test"))]
#[pgrx::pg_schema]
mod tests {
    use super::*;
    use crate::postgres::rel::PgSearchRelation;
    use pgrx::prelude::*;

    macro_rules! assert_approx_eq {
        ($a:expr, $b:expr, $percent:expr) => {{
            let a = $a;
            let b = $b;
            let diff = if a > b { a - b } else { b - a };
            let max_val = a.max(b);
            let max_diff = ((max_val as f64) * ($percent as f64 / 100.0)).ceil() as usize;

            assert!(
                diff <= max_diff,
                "assertion failed: `a = {}`, `b = {}` differ by more than {}% (allowed: {}, actual: {})",
                a,
                b,
                $percent,
                max_diff,
                diff
            );
        }};
    }

    #[pg_test]
    fn test_adjust_work_mem() {
        Spi::run("SET work_mem = '4MB';").unwrap();
        assert_approx_eq!(adjust_work_mem().get(), 15 * 1_000_000, 1.0);

        Spi::run("SET work_mem = '1GB';").unwrap();
        assert_approx_eq!(adjust_work_mem().get(), 1024 * 1024 * 1024, 1.0);
    }

    #[pg_test]
    fn test_adjust_maintenance_work_mem() {
        Spi::run("SET maintenance_work_mem = '16MB';").unwrap();
        assert_approx_eq!(adjust_maintenance_work_mem(0).get(), 16 * 1024 * 1024, 1.0);
        assert_approx_eq!(adjust_maintenance_work_mem(1).get(), 16 * 1024 * 1024, 1.0);
        assert!(std::panic::catch_unwind(|| adjust_maintenance_work_mem(2)).is_err());
        assert!(std::panic::catch_unwind(|| adjust_maintenance_work_mem(10)).is_err());

        Spi::run("SET maintenance_work_mem = '1GB';").unwrap();
        assert_approx_eq!(
            adjust_maintenance_work_mem(0).get(),
            1024 * 1024 * 1024,
            1.0
        );
        assert_approx_eq!(
            adjust_maintenance_work_mem(1).get(),
            1024 * 1024 * 1024,
            1.0
        );
        assert_approx_eq!(
            adjust_maintenance_work_mem(2).get(),
            1024 * 1024 * 1024,
            1.0
        );
        assert_approx_eq!(
            adjust_maintenance_work_mem(10).get(),
            1024 * 1024 * 1024,
            1.0
        );
        assert!(std::panic::catch_unwind(|| adjust_maintenance_work_mem(128)).is_err());
    }

    #[pg_test]
    fn test_global_mutable_segment_rows() {
        // valid options
        Spi::run("SET paradedb.global_mutable_segment_rows = 1000;")
            .expect("mutable_segment_rows should be 1000");
        assert_eq!(global_mutable_segment_rows(), Some(1000));

        Spi::run("SET paradedb.global_mutable_segment_rows = 0;")
            .expect("mutable_segment_rows should be 0");
        assert_eq!(global_mutable_segment_rows(), Some(0));

        Spi::run("SET paradedb.global_mutable_segment_rows = -1;").unwrap();
        assert_eq!(global_mutable_segment_rows(), None);

        // invalid options
        assert!(std::panic::catch_unwind(|| Spi::run(
            "SET paradedb.global_mutable_segment_rows = 10001;"
        ))
        .is_err());
        assert!(std::panic::catch_unwind(|| Spi::run(
            "SET paradedb.global_mutable_segment_rows = -2;"
        ))
        .is_err());

        // global override
        Spi::run("CREATE TABLE t (id SERIAL, data TEXT);").unwrap();
        Spi::run("INSERT INTO t (data) VALUES ('test');").unwrap();
        Spi::run("CREATE INDEX t_idx ON t USING bm25(id, data) WITH (key_field = 'id', mutable_segment_rows = 500)").unwrap();
        let relation_oid: pg_sys::Oid =
            Spi::get_one("SELECT oid FROM pg_class WHERE relname = 't_idx' AND relkind = 'i';")
                .expect("spi should succeed")
                .unwrap();
        let indexrel = PgSearchRelation::open(relation_oid);
        let options = indexrel.options();

        Spi::run("SET paradedb.global_mutable_segment_rows = -1;").unwrap();
        assert_eq!(
            options.mutable_segment_rows(),
            Some(NonZeroUsize::new(500).unwrap())
        );

        Spi::run("SET paradedb.global_mutable_segment_rows = 0;").unwrap();
        assert_eq!(options.mutable_segment_rows(), None);

        Spi::run("SET paradedb.global_mutable_segment_rows = 1000;").unwrap();
        assert_eq!(
            options.mutable_segment_rows(),
            Some(NonZeroUsize::new(1000).unwrap())
        );
    }
}
