// 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/>.

#![allow(clippy::unnecessary_cast)] // helps with integer casting differences between postgres versions
pub mod exec_methods;
pub mod parallel;
mod privdat;
pub mod projections;
mod scan_state;

use std::ffi::CStr;
use std::ptr::addr_of_mut;
use std::sync::atomic::Ordering;

use crate::api::operator::estimate_selectivity;
use crate::api::window_aggregate::window_agg_oid;
use crate::api::{HashMap, HashSet, OrderByFeature, OrderByInfo, Varno};
use crate::gucs;
use crate::index::fast_fields_helper::WhichFastField;
use crate::index::mvcc::MvccSatisfies;
use crate::index::reader::index::{SearchIndexReader, MAX_TOPK_FEATURES};
use crate::postgres::customscan::basescan::exec_methods::{
    fast_fields, normal::NormalScanExecState, ExecState,
};
use crate::postgres::customscan::basescan::privdat::PrivateData;
use crate::postgres::customscan::basescan::projections::score::uses_scores;
use crate::postgres::customscan::basescan::projections::snippet::{
    snippet_funcoids, snippet_positions_funcoids, snippets_funcoids, uses_snippets, SnippetType,
};
use crate::postgres::customscan::basescan::projections::window_agg::{
    deserialize_window_agg_placeholders, resolve_window_aggregate_filters_at_plan_time,
    WindowAggregateInfo,
};
use crate::postgres::customscan::basescan::scan_state::BaseScanState;
use crate::postgres::customscan::builders::custom_path::{
    restrict_info, CustomPathBuilder, ExecMethodType, Flags, RestrictInfoType,
};
use crate::postgres::customscan::builders::custom_scan::CustomScanBuilder;
use crate::postgres::customscan::builders::custom_state::{
    CustomScanStateBuilder, CustomScanStateWrapper,
};
use crate::postgres::customscan::explainer::Explainer;
use crate::postgres::customscan::orderby::{
    extract_pathkey_styles_with_sortability_check, find_sort_by_pathkey, PathKeyInfo,
    UnusableReason,
};
use crate::postgres::customscan::parallel::{compute_nworkers, list_segment_ids, RowEstimate};
use crate::postgres::customscan::projections::{
    inject_placeholders, maybe_needs_const_projections, pullout_funcexprs,
};
use crate::postgres::customscan::qual_inspect::{
    extract_join_predicates, extract_quals, is_subplan, optimize_quals_with_heap_expr,
    PlannerContext, Qual, QualExtractState,
};
use crate::postgres::customscan::score_funcoids;
use crate::postgres::customscan::solve_expr::SolvePostgresExpressions;
use crate::postgres::customscan::{
    self, range_table, CustomScan, CustomScanState, RelPathlistHookArgs,
};
use crate::postgres::heap::{HeapFetchState, VisibilityChecker};
use crate::postgres::rel::PgSearchRelation;
use crate::postgres::rel_get_bm25_index;
use crate::postgres::storage::metadata::MetaPage;
use crate::postgres::utils::{filter_implied_predicates, is_unnest_func};
use crate::query::pdb_query::pdb;
use crate::query::SearchQueryInput;
use crate::schema::SearchIndexSchema;
use crate::{nodecast, DEFAULT_STARTUP_COST, PARAMETERIZED_SELECTIVITY, UNKNOWN_SELECTIVITY};
use crate::{FULL_RELATION_SELECTIVITY, UNASSIGNED_SELECTIVITY};

use crate::postgres::customscan::limit_offset::extract_const_i64;
use pgrx::{pg_sys, FromDatum, IntoDatum, PgList, PgMemoryContexts};
use tantivy::snippet::SnippetGenerator;
use tantivy::Index;

#[derive(Default)]
pub struct BaseScan;

impl BaseScan {
    /// (Re-)initializes the search reader for the current execution context.
    ///
    /// This function handles three distinct execution scenarios:
    ///
    /// 1. **Leader Execution** (`ParallelWorkerNumber == -1`):
    ///    The scan is running in the main backend process. It uses `MvccSatisfies::Snapshot`
    ///    to see all segments visible to the current transaction's snapshot.
    ///
    /// 2. **Parallel-Aware Worker** (Worker with `parallel_state`):
    ///    The scan is part of a `parallel_aware` path (Partial Scan). Workers coordinate
    ///    via shared memory (DSM) to divide segments. It uses `MvccSatisfies::ParallelWorker`
    ///    to ensure it only queries segments explicitly identified and pinned by the leader.
    ///
    /// 3. **Replicated Worker** (Worker with NO `parallel_state`):
    ///    The scan is `parallel_safe` but NOT `parallel_aware`. This happens when a serial
    ///    scan runs inside a worker context (e.g., on the inner side of a Parallel Hash Join).
    ///    In this case, every worker executes the full scan independently using its own
    ///    transaction snapshot (`MvccSatisfies::Snapshot`).
    pub(crate) fn init_search_reader(state: &mut CustomScanStateWrapper<Self>) {
        let planstate = state.planstate();
        let expr_context = state.runtime_context;
        state
            .custom_state_mut()
            .prepare_query_for_execution(planstate, expr_context);

        // Open the index
        let indexrel = state
            .custom_state()
            .indexrel
            .as_ref()
            .expect("custom_state.indexrel should already be open");

        let search_query_input = state.custom_state().search_query_input();
        let need_scores = state.custom_state().need_scores();

        let search_reader = SearchIndexReader::open_with_context(
            indexrel,
            search_query_input.clone(),
            need_scores,
            unsafe {
                if pg_sys::ParallelWorkerNumber == -1 {
                    // the leader only sees snapshot-visible segments
                    MvccSatisfies::Snapshot
                } else if let Some(parallel_state) = state.custom_state().parallel_state {
                    // the workers have their own rules, which is literally every segment
                    // this is because the workers pick a specific segment to query that
                    // is known to be held open/pinned by the leader but might not pass a ::Snapshot
                    // visibility test due to concurrent merges/garbage collects
                    MvccSatisfies::ParallelWorker(list_segment_ids(parallel_state))
                } else {
                    // We are in a worker, but this is not a parallel-aware scan (e.g. we are running
                    // a serial scan inside a parallel worker, like in a Parallel Nested Loop Join).
                    // In this case, we behave like a normal snapshot scan.
                    MvccSatisfies::Snapshot
                }
            },
            std::ptr::NonNull::new(expr_context),
            std::ptr::NonNull::new(planstate),
        )
        .expect("should be able to open the search index reader");
        state.custom_state_mut().search_reader = Some(search_reader);

        let csstate = addr_of_mut!(state.csstate);
        state.custom_state_mut().init_exec_method(csstate);

        if state.custom_state().need_snippets() {
            let mut snippet_generators: HashMap<SnippetType, Option<SnippetGenerator>> = state
                .custom_state_mut()
                .snippet_generators
                .drain()
                .collect();

            // Pre-compute enhanced queries for snippet generation if we have join predicates
            let enhanced_query_for_snippets =
                if let Some(ref join_predicate) = state.custom_state().join_predicates {
                    // Combine base query with join predicate for snippet generation
                    let base_query = state.custom_state().search_query_input();
                    Some(SearchQueryInput::Boolean {
                        must: vec![base_query.clone()],
                        should: vec![join_predicate.clone()],
                        must_not: vec![],
                    })
                } else {
                    None
                };

            for (snippet_type, generator) in &mut snippet_generators {
                // Use enhanced query if available, otherwise use base query
                let query_to_use = enhanced_query_for_snippets
                    .as_ref()
                    .unwrap_or_else(|| state.custom_state().search_query_input());

                let mut new_generator = state
                    .custom_state()
                    .search_reader
                    .as_ref()
                    .unwrap()
                    .snippet_generator(
                        snippet_type.field().root(),
                        query_to_use,
                        std::ptr::NonNull::new(expr_context),
                    );

                snippet_type.configure_generator(&mut new_generator.1);

                *generator = Some(new_generator.1);
            }

            state.custom_state_mut().snippet_generators = snippet_generators;
        }

        unsafe {
            inject_pdb_placeholders(state);
        }
    }

    #[allow(clippy::too_many_arguments)]
    unsafe fn extract_all_possible_quals(
        builder: &mut CustomPathBuilder<BaseScan>,
        root: *mut pg_sys::PlannerInfo,
        rti: pg_sys::Index,
        restrict_info: PgList<pg_sys::RestrictInfo>,
        ri_type: RestrictInfoType,
        indexrel: &PgSearchRelation,
        uses_score_or_snippet: bool,
        attempt_pushdown: bool,
    ) -> Option<Qual> {
        let mut state = QualExtractState::default();
        let context = PlannerContext::from_planner(root);

        // Filter out predicates that are implied by the partial index predicate.
        // If a partial index has predicate P (e.g., "deleted_at IS NULL"), and the query
        // also has predicate P, we don't need to create a heap filter for P since the
        // partial index already guarantees it.
        let filtered_restrict_info = filter_implied_predicates(indexrel.rd_indpred, &restrict_info);

        let mut quals = extract_quals(
            &context,
            rti,
            filtered_restrict_info.as_ptr().cast(),
            ri_type,
            indexrel,
            false, // Base relation quals should not convert external to all
            &mut state,
            attempt_pushdown,
        );

        // If full extraction failed (e.g., baserestrictinfo contains SubPlan from
        // RLS policies alongside our @@@ operator), try partial extraction: extract
        // each restrict_info item individually, skipping ones we can't handle.
        // Only use partial extraction when ALL skipped clauses are SubPlans
        // (which will be evaluated via plan.qual). If any non-SubPlan clause
        // is skipped, fall back to let PostgreSQL handle the query normally.
        //
        // TODO: We do something similar in `collect_join_sources_base_rel`,
        // is unification possible?
        if quals.is_none() {
            let mut partial_quals = Vec::new();
            let mut partial_state = QualExtractState::default();
            let mut all_skipped_are_subplans = true;
            for ri in filtered_restrict_info.iter_ptr() {
                if let Some(qual) = extract_quals(
                    &context,
                    rti,
                    ri.cast(),
                    ri_type,
                    indexrel,
                    false,
                    &mut partial_state,
                    attempt_pushdown,
                ) {
                    partial_quals.push(qual);
                } else if !is_subplan(ri.cast()) {
                    all_skipped_are_subplans = false;
                }
            }
            if !partial_quals.is_empty()
                && partial_state.uses_our_operator
                && all_skipped_are_subplans
            {
                state = partial_state;
                quals = if partial_quals.len() == 1 {
                    partial_quals.pop()
                } else {
                    Some(Qual::And(partial_quals))
                };
            }
        }

        // If we couldn't push down quals, try to push down quals from the join
        // This is only done if we have a join predicate, and only if we have used our operator
        let quals = if quals.is_none() {
            let joinri: PgList<pg_sys::RestrictInfo> =
                PgList::from_pg(builder.args().rel().joininfo);
            let mut quals = extract_quals(
                &context,
                rti,
                joinri.as_ptr().cast(),
                RestrictInfoType::Join,
                indexrel,
                true, // Join quals should convert external to all
                &mut state,
                attempt_pushdown,
            );

            let quals = Self::handle_heap_expr_optimization(&state, &mut quals);

            // If we have found something to push down in the join, then we can use the join quals
            // Note: these Join quals won't help in filtering down the data (as they contain
            // external vars, e.g. `b.category_name @@@ "technology"` in
            // `a.name @@@ "abc" OR b.category_name @@@ "technology"`), and we cannot evaluate
            // boolean expressions that contain external vars. That's why, when handling the Join
            // quals, we'd endup scanning the whole tantivy index.
            // However, the Join quals help with scoring and snippet generation, as the documents
            // that match partially the Join quals will be scored and snippets generated. That is
            // why it only makes sense to use the Join quals if we have used our operator and
            // also used pdb.score or pdb.snippet functions in the query.
            if state.uses_our_operator && uses_score_or_snippet {
                quals
            } else {
                None
            }
        } else {
            Self::handle_heap_expr_optimization(&state, &mut quals)
        };

        // Finally, decide whether we can actually use the extracted quals.
        // We allow custom scan if:
        // 1. The query uses @@@ operator, OR
        // 2. enable_custom_scan_without_operator is true, OR
        // 3. The query has window aggregates (pdb.agg()) that we must handle
        let has_window_aggs = query_has_window_agg_functions(root);
        if state.uses_our_operator || gucs::enable_custom_scan_without_operator() || has_window_aggs
        {
            quals
        } else {
            None
        }
    }

    unsafe fn handle_heap_expr_optimization(
        state: &QualExtractState,
        quals: &mut Option<Qual>,
    ) -> Option<Qual> {
        if state.uses_heap_expr && !state.uses_our_operator {
            return None;
        }

        // Apply HeapExpr optimization to the base relation quals
        if let Some(ref mut q) = quals {
            optimize_quals_with_heap_expr(q);
        }

        quals.clone()
    }
}

/// Check if the query's target list contains window_agg() function calls
///
/// This is called AFTER window function replacement in BaseScan's create_custom_path.
/// It looks for FuncExpr nodes with window_agg() OID, NOT WindowFunc nodes.
///
/// This is different from query_has_window_functions() in hook.rs which looks for WindowFunc
/// nodes BEFORE replacement in the planner hook.
///
/// Used to determine if we should create a custom path even without @@@ operator.
///
/// Also validates that pdb.agg() is not present - if it is, that means the planner hook
/// didn't replace it (e.g., not a Top K query), and we should reject it.
unsafe fn query_has_window_agg_functions(root: *mut pg_sys::PlannerInfo) -> bool {
    if root.is_null() || (*root).parse.is_null() {
        return false;
    }

    let parse = (*root).parse;
    let window_agg_func_oid = window_agg_oid();
    let paradedb_agg_func_oid = crate::api::agg_funcoid();

    // If functions don't exist yet (e.g., during extension creation), skip check
    if window_agg_func_oid == pg_sys::InvalidOid {
        return false;
    }

    let window_agg_func_oid = window_agg_func_oid.to_u32();
    let paradedb_agg_func_oid = paradedb_agg_func_oid.to_u32();

    // Check target list for window_agg() or pdb.agg() function calls
    if !(*parse).targetList.is_null() {
        let target_list = PgList::<pg_sys::TargetEntry>::from_pg((*parse).targetList);
        for te in target_list.iter_ptr() {
            if !(*te).expr.is_null() {
                // Check if this is a FuncExpr with window_agg or pdb.agg OID
                if let Some(func_expr) = nodecast!(FuncExpr, T_FuncExpr, (*te).expr) {
                    let func_oid = (*func_expr).funcid.to_u32();
                    if func_oid == window_agg_func_oid {
                        return true;
                    } else if func_oid == paradedb_agg_func_oid {
                        // pdb.agg() should have been replaced by planner hook
                        // If it's still here, it means it wasn't a valid Top K query
                        pgrx::error!(
                            "pdb.agg() can only be used as a window function in Top K queries \
                             (queries with ORDER BY and LIMIT). For GROUP BY aggregates, use standard \
                             SQL aggregates like COUNT(*), SUM(), etc. \
                             Hint: Try using '@@@ pdb.all()' with ORDER BY and LIMIT, \
                             or see https://github.com/paradedb/paradedb/issues for more information."
                        );
                    }
                }
            }
        }
    }

    false
}

/// Check if the query's target list contains only set-returning functions (e.g. `unnest()`) which
/// are safe for use with our LIMIT optimization, and if so, then return the LIMIT from the parse.
///
/// Only set returning functions which produce at least as many rows as they consume are safe.
unsafe fn maybe_limit_from_parse(root: *mut pg_sys::PlannerInfo) -> Option<f64> {
    if root.is_null() || (*root).parse.is_null() || (*(*root).parse).targetList.is_null() {
        return None;
    }
    let parse = *(*root).parse;
    // non-Const LIMIT is not a thing we can handle here
    let limit = extract_const_i64(parse.limitCount).map(|v| v as f64)?;

    let offset = if parse.limitOffset.is_null() {
        0.0
    } else if let Some(offset_const) = nodecast!(Const, T_Const, parse.limitOffset) {
        i64::from_datum((*offset_const).constvalue, (*offset_const).constisnull)
            .map(|v| v as f64)?
    } else {
        // non-Const OFFSET is not a thing we can handle here
        return None;
    };

    let mut found_limit_safe_srf = false;
    let target_list = PgList::<pg_sys::TargetEntry>::from_pg(parse.targetList);
    for te in target_list.iter_ptr() {
        if !(*te).expr.is_null() && pg_sys::expression_returns_set((*te).expr.cast()) {
            // It's a set-returning function, is it one we approve of?
            if let Some(func_expr) = nodecast!(FuncExpr, T_FuncExpr, (*te).expr) {
                if is_unnest_func((*func_expr).funcid) {
                    found_limit_safe_srf = true;
                } else {
                    // We don't recognize this SRF, and can't vouch for it.
                    return None;
                }
            }
        }
    }

    // A LIMIT was applied in the parse for our node, but is being handled elsewhere
    // due to a set-returning-function that we know produces at least as many tuples as
    // it is given.
    if found_limit_safe_srf {
        Some(limit + offset)
    } else {
        None
    }
}

impl CustomScan for BaseScan {
    const NAME: &'static CStr = c"ParadeDB Base Scan";

    type Args = RelPathlistHookArgs;
    type State = BaseScanState;
    type PrivateData = PrivateData;

    fn exec_methods() -> pg_sys::CustomExecMethods {
        pg_sys::CustomExecMethods {
            CustomName: Self::NAME.as_ptr(),
            BeginCustomScan: Some(customscan::exec::begin_custom_scan::<Self>),
            ExecCustomScan: Some(customscan::exec::exec_custom_scan::<Self>),
            EndCustomScan: Some(customscan::exec::end_custom_scan::<Self>),
            ReScanCustomScan: Some(customscan::exec::rescan_custom_scan::<Self>),
            MarkPosCustomScan: None,
            RestrPosCustomScan: None,
            EstimateDSMCustomScan: Some(customscan::dsm::estimate_dsm_custom_scan::<Self>),
            InitializeDSMCustomScan: Some(customscan::dsm::initialize_dsm_custom_scan::<Self>),
            ReInitializeDSMCustomScan: Some(customscan::dsm::reinitialize_dsm_custom_scan::<Self>),
            InitializeWorkerCustomScan: Some(
                customscan::dsm::initialize_worker_custom_scan::<Self>,
            ),
            ShutdownCustomScan: Some(customscan::exec::shutdown_custom_scan::<Self>),
            ExplainCustomScan: Some(customscan::exec::explain_custom_scan::<Self>),
        }
    }

    fn create_custom_path(mut builder: CustomPathBuilder<Self>) -> Vec<pg_sys::CustomPath> {
        let paths = (|| unsafe {
            let (restrict_info, ri_type) = restrict_info(builder.args().rel());
            let rel = builder.args().rel;

            // Check if the query has window aggregates (pdb.agg() or window_agg())
            let has_window_aggs = query_has_window_agg_functions(builder.args().root);

            if matches!(ri_type, RestrictInfoType::None) && !has_window_aggs {
                // this relation has no restrictions (WHERE clause predicates) and no window aggregates,
                // so there's no need for us to do anything
                return None;
            }

            let rti = builder.args().rti;
            let (table, bm25_index) = {
                let rte = builder.args().rte();

                // we only support plain relation and join rte's
                if rte.rtekind != pg_sys::RTEKind::RTE_RELATION
                    && rte.rtekind != pg_sys::RTEKind::RTE_JOIN
                {
                    return None;
                }

                // and we only work on plain relations
                let relkind = pg_sys::get_rel_relkind(rte.relid) as u8;
                if relkind != pg_sys::RELKIND_RELATION && relkind != pg_sys::RELKIND_MATVIEW {
                    return None;
                }

                // and that relation must have a `USING bm25` index
                let (table, bm25_index) = rel_get_bm25_index(rte.relid)?;

                (table, bm25_index)
            };

            let root = builder.args().root;

            // quick look at the target list to see if we might need to do our const projections
            let target_list = (*(*builder.args().root).parse).targetList;
            let maybe_needs_const_projections = maybe_needs_const_projections(target_list.cast());

            //
            // look for quals we can support.  we do this first so that we can get out early if this
            // isn't a query we can support.
            //
            // Opening the Directory and Index down below is expensive, so if we can avoid it,
            // especially for non-SELECT (ie, UPDATE) statements, that's good
            //
            let is_select =
                (*(*builder.args().root).parse).commandType == pg_sys::CmdType::CMD_SELECT;
            let quals = Self::extract_all_possible_quals(
                &mut builder,
                root,
                rti,
                PgList::from_pg(restrict_info.as_ptr()),
                ri_type,
                &bm25_index,
                maybe_needs_const_projections,
                is_select,
            );

            // If we have window aggregates but no quals, we must still create the custom path
            // because pdb.agg() can only be executed by our custom scan
            let quals = if let Some(q) = quals {
                q
            } else if has_window_aggs {
                // We have window aggregates but couldn't extract quals.
                // This can happen in two cases:
                // 1. No WHERE clause at all -> safe to use Qual::All
                // 2. WHERE clause exists but couldn't be extracted:
                //    a. filter_pushdown enabled -> HeapExpr was created during extraction, safe to use Qual::All
                //    b. filter_pushdown disabled -> unsafe, reject the query
                let has_where_clause = !(*root).parse.is_null()
                    && !(*(*root).parse).jointree.is_null()
                    && !(*(*(*root).parse).jointree).quals.is_null();

                if has_where_clause && !crate::gucs::enable_filter_pushdown() {
                    // There's a WHERE clause but we couldn't extract quals and filter_pushdown is disabled.
                    // This means qual extraction failed without creating HeapExpr, so we cannot handle
                    // this query safely - the WHERE clause would be silently ignored.
                    return None;
                }

                // Safe to use Qual::All because:
                // - Either there's no WHERE clause (nothing to filter), OR
                // - filter_pushdown is enabled, meaning HeapExpr was created during qual extraction
                //   and will be evaluated by PostgreSQL's executor after we return results
                Qual::All
            } else {
                // No quals and no window aggregates - we can't help
                return None;
            };

            // Check if this is a partial index and if the query is compatible with it
            if !bm25_index.rd_indpred.is_null() {
                // This is a partial index - we need to check if the query predicates
                // imply the partial index predicate using PostgreSQL's predicate_implied_by.
                //
                // For example:
                // - Partial index: WHERE deleted_at IS NULL
                // - Query: WHERE deleted_at IS NULL AND category_id = 5
                // - predicate_implied_by(index_pred, query_clauses) returns true
                //
                // But:
                // - Partial index: WHERE category = 'Electronics'
                // - Query: WHERE description @@@ 'Apple' AND rating >= 4
                // - predicate_implied_by returns false (query doesn't imply category = 'Electronics')

                // Extract the restriction clauses as a list of Expr nodes
                let mut clause_list: *mut pg_sys::List = std::ptr::null_mut();
                for ri in restrict_info.iter_ptr() {
                    clause_list =
                        pg_sys::lappend(clause_list, (*ri).clause as *mut std::ffi::c_void);
                }

                // Check if query clauses imply the partial index predicate
                let is_compatible =
                    pg_sys::predicate_implied_by(bm25_index.rd_indpred, clause_list, false);

                if !is_compatible {
                    // The query predicates don't imply the partial index predicate,
                    // so we cannot safely use this partial index
                    return None;
                }
            }

            //
            // ===================
            // If we make it this far, we're going to submit a path... it better be a good one!
            // ====================
            //

            // TODO: `impl Default for PrivateData` requires that many fields are in invalid
            // states. Should consider having a separate builder for PrivateData.
            let mut custom_private = PrivateData::default();

            let segment_count = {
                let directory = MvccSatisfies::LargestSegment.directory(&bm25_index);
                let segment_count = directory.total_segment_count(); // return value only valid after the index has been opened
                Index::open(directory).expect("custom_scan: should be able to open index");
                segment_count.load(Ordering::Relaxed)
            };
            let schema = bm25_index
                .schema()
                .expect("custom_scan: should have a schema");
            let topk_pathkey_info = pullup_topk_pathkeys(rti, &schema, root);

            #[cfg(feature = "pg15")]
            let baserels = (*builder.args().root).all_baserels;
            #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))]
            let baserels = (*builder.args().root).all_query_rels;

            // Detect if we are in a join context (more than 1 base relation in the query)
            // If so, we want to be aggressive with parallelism to enable Parallel Hash Join
            let is_join_context = pg_sys::bms_num_members(baserels) > 1;

            let limit = if (*builder.args().root).limit_tuples > -1.0 {
                // Check if this is a single relation or a partitioned table setup
                let rel_is_single_or_partitioned = pg_sys::bms_equal((*rel).relids, baserels)
                    || range_table::is_partitioned_table_setup(
                        builder.args().root,
                        (*rel).relids,
                        baserels,
                    );

                // Check for LEFT JOIN LATERAL where left side drives the query
                let is_left_driven_lateral = is_left_join_lateral(builder.args().root, rel)
                    && where_clause_only_references_left(builder.args().root, rti);

                if rel_is_single_or_partitioned || is_left_driven_lateral {
                    // We can use the limit for estimates if:
                    // a) we have a limit, and
                    // b) we're either:
                    //    * querying a single relation OR
                    //    * querying partitions of a partitioned table OR
                    //    * we're in a LEFT JOIN LATERAL where the left side drives the query
                    Some((*builder.args().root).limit_tuples)
                } else {
                    None
                }
            } else {
                maybe_limit_from_parse(builder.args().root)
            };

            // Get all columns referenced by this RTE throughout the entire query
            let referenced_columns = collect_maybe_fast_field_referenced_columns(rti, rel);

            // Save the count of referenced columns for decision-making
            custom_private.set_referenced_columns_count(referenced_columns.len());

            let is_maybe_topk = limit.is_some() && topk_pathkey_info.is_usable();

            // When collecting which_fast_fields, analyze the entire set of referenced columns,
            // not just those in the target list. To avoid execution-time surprises, the "planned"
            // fast fields must be a superset of the fast fields which are extracted from the
            // execution-time target list: see `assign_exec_method` for more info.
            custom_private.set_planned_which_fast_fields(
                exec_methods::fast_fields::collect_fast_fields(
                    target_list,
                    &referenced_columns,
                    rti,
                    &table,
                    &bm25_index,
                    false,
                )
                .into_iter()
                .collect(),
            );

            let query = SearchQueryInput::from(&quals);
            let norm_selec = if restrict_info.len() == 1 {
                (*restrict_info.get_ptr(0).unwrap()).norm_selec
            } else {
                UNASSIGNED_SELECTIVITY
            };

            let selectivity = if norm_selec != UNASSIGNED_SELECTIVITY {
                // we can use the norm_selec that already happened
                norm_selec
            } else if quals.contains_external_var() {
                // if the query has external vars (references to other relations which decide whether the rows in this
                // relation are visible) then we end up returning *everything* from _this_ relation
                FULL_RELATION_SELECTIVITY
            } else if quals.contains_exprs() {
                // if the query has expressions then it's parameterized and we have to guess something
                PARAMETERIZED_SELECTIVITY
            } else {
                // ask the index
                estimate_selectivity(&bm25_index, query.clone()).unwrap_or(UNKNOWN_SELECTIVITY)
            };

            custom_private.set_heaprelid(table.oid());
            custom_private.set_indexrelid(bm25_index.oid());
            custom_private.set_range_table_index(rti);
            custom_private.set_query(query);
            custom_private.set_limit(limit);
            custom_private.set_segment_count(segment_count);

            // Determine whether we might be able to sort.
            if is_maybe_topk && topk_pathkey_info.pathkeys().is_some() {
                custom_private.set_maybe_orderby_info(topk_pathkey_info.pathkeys());
            }

            // Choose the exec method type, and make claims about whether it is sorted.
            let limit_is_explicit =
                limit.is_some() && !is_minmax_implicit_limit(builder.args().root);

            // Check if the index has a sort_by configuration and the query has a matching pathkey.
            // If so, create an additional sorted path that declares the pathkey.
            // This allows Postgres to use merge joins when joining with other sorted inputs.
            let sort_by_pathkey = if gucs::is_columnar_sort_enabled() {
                let sort_by_fields = bm25_index.options().sort_by();
                sort_by_fields
                    .first()
                    .and_then(|sort_by| find_sort_by_pathkey(root, rti, sort_by, &table))
            } else {
                None
            };

            // calculate the total number of rows that might match the query, and the number of
            // rows that we expect that scan to return: these may be different in the case of a
            // `limit`.
            //
            // Use RowEstimate enum to distinguish between known and unknown row counts.
            // Unknown is used when the table hasn't been ANALYZEd.
            // TODO: Convert to use RowEstimate::from_reltuples.
            let row_estimate = match table.reltuples() {
                Some(reltuples) if reltuples > 0.0 => {
                    let estimated = (reltuples as f64 * selectivity).max(1.0) as u64;
                    RowEstimate::Known(estimated)
                }
                _ => RowEstimate::Unknown,
            };
            let base_result_rows = match row_estimate {
                RowEstimate::Known(rows) => (rows as f64).min(limit.unwrap_or(f64::MAX)).max(1.0),
                RowEstimate::Unknown => {
                    // For unknown row counts, use 1.0 as a conservative estimate for costing
                    limit.unwrap_or(1.0).max(1.0)
                }
            };

            let exec_method_types = choose_exec_method(
                &custom_private,
                &topk_pathkey_info,
                limit_is_explicit,
                table.name(),
                sort_by_pathkey.is_some(),
            );

            //
            // finally, we have enough information to set the cost and estimation information, and
            // to decide on parallelism
            //

            let startup_cost = DEFAULT_STARTUP_COST;
            let mut custom_paths = Vec::new();

            // For each execution method variant (e.g. sorted vs unsorted), we build a separate
            // CustomPath. This allows the Postgres planner to choose the most efficient
            // implementation based on costs and downstream requirements like ordering.
            for method in exec_method_types {
                let per_tuple_cost = match &method {
                    // returning fields from fast fields
                    ExecMethodType::Columnar { .. } => pg_sys::cpu_index_tuple_cost,
                    // requires heap access to return fields
                    _ => pg_sys::cpu_tuple_cost,
                };

                let mut path_builder = CustomPathBuilder::<Self>::new(
                    builder.args().root,
                    builder.args().rel,
                    *builder.args(),
                );

                // we must use this path if we need to do const projections for scores or snippets
                path_builder = path_builder.set_force_path(
                    maybe_needs_const_projections
                        || matches!(method, ExecMethodType::TopK { .. })
                        || quals.contains_all(),
                );

                let mut method_private = custom_private.clone();
                method_private.set_exec_method_type(method.clone());

                let is_sorted = method.declares_sorted_output();
                method_private.set_use_sorted_path(is_sorted);

                // Our BaseScan is always parallel-safe (can run in a worker),
                // even if it's not parallel-aware (splitting segments).
                let nworkers = if (*builder.args().rel).consider_parallel {
                    path_builder = path_builder.set_parallel_safe(true);

                    compute_nworkers(
                        is_sorted,
                        limit,
                        row_estimate,
                        segment_count,
                        quals.contains_external_var(),
                        quals.contains_correlated_param(builder.args().root),
                        is_join_context,
                    )
                } else {
                    0
                };

                let mut method_result_rows = base_result_rows;

                if nworkers > 0 {
                    path_builder = path_builder.set_parallel(nworkers);

                    // if we're likely to do a parallel scan, divide the result_rows by the number of workers
                    // we're likely to use.  this lets Postgres make better decisions based on what
                    // an individual parallel scan is actually going to return
                    let processes = std::cmp::max(
                        1,
                        nworkers
                            + if pg_sys::parallel_leader_participation {
                                1
                            } else {
                                0
                            },
                    );
                    method_result_rows /= processes as f64;
                }

                let mut total_cost = startup_cost + (method_result_rows * per_tuple_cost);

                if is_sorted && method.supports_sorted_index_merge() {
                    total_cost *= 1.01;
                }

                path_builder = path_builder.set_rows(method_result_rows);
                path_builder = path_builder.set_startup_cost(startup_cost);
                path_builder = path_builder.set_total_cost(total_cost);

                // indicate that we'll be doing projection ourselves
                path_builder = path_builder.set_flag(Flags::Projection);

                // If Top K, add pathkeys to builder
                if matches!(
                    method,
                    ExecMethodType::TopK {
                        orderby_info: Some(..),
                        ..
                    }
                ) {
                    path_builder = path_builder.set_pathkeys((*builder.args().root).query_pathkeys);
                } else if is_sorted {
                    // For sorted columnar execution, add the sort pathkey
                    if let Some(ref pathkey_style) = sort_by_pathkey {
                        path_builder = path_builder.add_path_key(pathkey_style);
                    }
                }

                custom_paths.push(path_builder.build(method_private));
            }

            Some(custom_paths)
        })();
        paths.unwrap_or_default()
    }

    fn plan_custom_path(mut builder: CustomScanBuilder<Self>) -> pg_sys::CustomScan {
        unsafe {
            let mut tlist = PgList::<pg_sys::TargetEntry>::from_pg(builder.args().tlist.as_ptr());

            // Store the length of the target list
            builder
                .custom_private_mut()
                .set_target_list_len(Some(tlist.len()));

            // Extract window_agg(json) calls from processed_tlist using expression tree walker
            // Similar to how uses_scores/uses_snippets work - walk the tree to find our placeholders
            // Note: This updates target_entry_index to match the processed_tlist positions
            let processed_tlist = (*builder.args().root).processed_tlist;

            let mut window_aggregates = deserialize_window_agg_placeholders(processed_tlist);

            if !window_aggregates.is_empty() {
                // Convert PostgresExpression filters to SearchQueryInput now that we have root
                // Note: root was not available in the planner hook, so we needed to delay this until now.
                let private_data = builder.custom_private();
                if let Some(heaprelid) = private_data.heaprelid() {
                    if let Some((_, bm25_index)) = rel_get_bm25_index(heaprelid) {
                        let root = builder.args().root;
                        let rti = private_data
                            .range_table_index()
                            .expect("range table index should be set");

                        resolve_window_aggregate_filters_at_plan_time(
                            &mut window_aggregates,
                            &bm25_index,
                            root,
                            rti,
                        );

                        // Validate that all fields in window aggregates exist in the index schema
                        // and are supported for aggregate pushdown (not NUMERIC)
                        if let Ok(schema) = crate::schema::SearchIndexSchema::open(&bm25_index) {
                            for window_agg in &window_aggregates {
                                for agg_type in window_agg.targetlist.aggregates() {
                                    if let Err(e) = agg_type.validate_fields(&schema) {
                                        pgrx::error!("{}", e);
                                    }
                                }
                            }
                        }
                    }
                }

                builder
                    .custom_private_mut()
                    .set_window_aggregates(window_aggregates);
            }

            let private_data = builder.custom_private();
            let rti = private_data
                .range_table_index()
                .expect("range table index should have been set")
                .try_into()
                .expect("range table index should not be negative");
            let processed_tlist = PgList::<pg_sys::TargetEntry>::from_pg(processed_tlist);

            let mut attname_lookup = HashMap::default();
            let funcoids: Vec<pg_sys::Oid> = score_funcoids()
                .iter()
                .copied()
                .chain(snippet_funcoids().iter().copied())
                .chain(snippets_funcoids().iter().copied())
                .chain(snippet_positions_funcoids().iter().copied())
                .collect();
            for te in processed_tlist.iter_ptr() {
                let func_vars_at_level =
                    pullout_funcexprs(te.cast(), &funcoids, rti, builder.args().root);

                for (funcexpr, var, attname) in func_vars_at_level {
                    // if we have a tlist, then we need to add the specific function that uses
                    // a Var at our level to that tlist.
                    //
                    // if we don't have a tlist (it's empty), then that means Postgres will later
                    // give us everything we need

                    if !tlist.is_empty() {
                        let te = pg_sys::copyObjectImpl(te.cast()).cast::<pg_sys::TargetEntry>();
                        (*te).resno = (tlist.len() + 1) as _;
                        (*te).expr = funcexpr.cast();

                        tlist.push(te);
                    }

                    // track a triplet of (varno, varattno, attname) as 3 individual
                    // entries in the `attname_lookup` List
                    attname_lookup.insert((rti as Varno, (*var).varattno), attname);
                }
            }

            // Extract join-level snippet predicates for this relation
            // Get values we need before the mutable borrow

            // Extract the indexrelid early to avoid borrow checker issues later
            let indexrelid = private_data.indexrelid().expect("indexrelid should be set");
            let indexrel = PgSearchRelation::with_lock(indexrelid, pg_sys::AccessShareLock as _);

            let join_predicates = extract_join_predicates(
                &PlannerContext::from_planner(builder.args().root),
                rti as pg_sys::Index,
                &indexrel,
                true,
            );

            builder
                .custom_private_mut()
                .set_join_predicates(join_predicates);

            builder
                .custom_private_mut()
                .set_var_attname_lookup(attname_lookup);

            builder
                .custom_private_mut()
                .set_ambulkdelete_epoch(MetaPage::open(&indexrel).ambulkdelete_epoch());

            // Collect subplans that our Custom Scan doesn't handle internally and set them as plan.qual.
            // PostgreSQL's ExecInitCustomScan will call ExecInitQual on plan.qual,
            // which properly initializes SubPlans. We then evaluate these in exec_custom_scan.
            let clauses = PgList::<pg_sys::Node>::from_pg(builder.args().clauses);
            let mut subplan_quals = PgList::<pg_sys::Node>::new();
            for clause in clauses.iter_ptr() {
                if is_subplan(clause) {
                    // strip RestrictInfo wrapper, plan.qual needs bare expressions
                    let bare_clause = if (*clause).type_ == pg_sys::NodeTag::T_RestrictInfo {
                        let ri = clause as *mut pg_sys::RestrictInfo;
                        (*ri).clause.cast()
                    } else {
                        clause
                    };
                    subplan_quals.push(bare_clause);
                }
            }

            // SubPlan quals require per-tuple heap access for ExecQual, which would
            // negate the benefit of columnar. Fall back to Normal so we don't pay for both
            // batch processing AND per-tuple heap fetches.
            if !subplan_quals.is_empty()
                && matches!(
                    builder.custom_private().exec_method_type(),
                    ExecMethodType::Columnar { .. }
                )
            {
                builder
                    .custom_private_mut()
                    .set_exec_method_type(ExecMethodType::Normal);
            }

            let mut scan = builder.build();
            if !subplan_quals.is_empty() {
                scan.scan.plan.qual = subplan_quals.into_pg();
            }
            scan
        }
    }

    fn create_custom_scan_state(
        mut builder: CustomScanStateBuilder<Self, Self::PrivateData>,
    ) -> *mut CustomScanStateWrapper<Self> {
        unsafe {
            builder.custom_state().heaprelid = builder
                .custom_private()
                .heaprelid()
                .expect("heaprelid should have a value");
            builder.custom_state().indexrelid = builder
                .custom_private()
                .indexrelid()
                .expect("indexrelid should have a value");

            builder
                .custom_state()
                .open_relations(pg_sys::AccessShareLock as _);

            builder.custom_state().execution_rti =
                (*builder.args().cscan).scan.scanrelid as pg_sys::Index;

            builder.custom_state().exec_method_type =
                builder.custom_private().exec_method_type().clone();

            builder.custom_state().targetlist_len = builder.target_list().len();

            builder.custom_state().segment_count = builder.custom_private().segment_count();
            builder.custom_state().var_attname_lookup = builder
                .custom_private()
                .var_attname_lookup()
                .as_ref()
                .cloned()
                .expect("should have an attribute name lookup");

            let score_funcoids = score_funcoids();
            let snippet_funcoids = snippet_funcoids();
            let snippets_funcoids = snippets_funcoids();
            let snippet_positions_funcoids = snippet_positions_funcoids();

            builder.custom_state().score_funcoids = score_funcoids;
            builder.custom_state().snippet_funcoids = snippet_funcoids;
            builder.custom_state().snippets_funcoids = snippets_funcoids;
            builder.custom_state().snippet_positions_funcoids = snippet_positions_funcoids;
            builder.custom_state().need_scores = uses_scores(
                builder.target_list().as_ptr().cast(),
                score_funcoids,
                builder.custom_state().execution_rti,
            );

            // Store join snippet predicates in the scan state
            builder.custom_state().join_predicates =
                builder.custom_private().join_predicates().clone();

            // Store window aggregates in the scan state
            let window_aggs = builder.custom_private().window_aggregates().clone();
            builder.custom_state().window_aggregates = window_aggs;

            // store our query into our custom state too
            let base_query = builder
                .custom_private()
                .query()
                .clone()
                .expect("should have a SearchQueryInput");
            builder
                .custom_state()
                .set_base_search_query_input(base_query);

            if builder.custom_state().need_scores {
                let state = builder.custom_state();
                // Pre-compute enhanced score query if we have join predicates that could affect scoring
                let mut enhanced_score_query = None;
                if let Some(ref join_predicate) = state.join_predicates {
                    // Check the ORIGINAL base query for this relation, not the modified search_query_input
                    // which may contain simplified join predicates from other relations
                    let original_base_query = state.base_search_query_input();

                    // Only enhance scoring if the base query doesn't already have search predicates
                    // If base query has @@@ conditions, it already provides scoring context
                    if !base_query_has_search_predicates(original_base_query, state.indexrelid) {
                        // Combine base query with join predicate using Boolean structure
                        // This provides enhanced search context for scoring while maintaining
                        // the same filtering behavior as the base query
                        enhanced_score_query = Some(SearchQueryInput::Boolean {
                            must: vec![original_base_query.clone()],
                            should: vec![join_predicate.clone()],
                            must_not: vec![],
                        });
                    }
                }

                // Store enhanced score query for use during search execution
                // This will be None for single-table queries, which is correct
                if let Some(enhanced_score_query) = enhanced_score_query {
                    builder
                        .custom_state()
                        .set_base_search_query_input(enhanced_score_query);
                }
            }

            let node = builder.target_list().as_ptr().cast();
            builder.custom_state().planning_rti = builder
                .custom_private()
                .range_table_index()
                .expect("range table index should have been set");
            builder.custom_state().snippet_generators = uses_snippets(
                builder.custom_state().planning_rti,
                &builder.custom_state().var_attname_lookup,
                node,
                snippet_funcoids,
                snippets_funcoids,
                snippet_positions_funcoids,
            )
            .into_iter()
            .map(|snippet_type| (snippet_type, None))
            .collect();

            builder.custom_state().ambulkdelete_epoch =
                builder.custom_private().ambulkdelete_epoch();

            assign_exec_method(&mut builder);

            let state = builder.build();

            // Tell ExecInitCustomScan to create the scan slot as BufferHeapTuple
            // (matching what begin_custom_scan will use via table_slot_callbacks).
            // This ensures ExecInitQual compiles plan.qual expressions for the
            // correct slot type, avoiding TTS_IS_VIRTUAL assertion failures.
            #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))]
            {
                (*state).csstate.slotOps =
                    pg_sys::table_slot_callbacks((*state).custom_state().heaprel().as_ptr());
            }

            state
        }
    }

    fn explain_custom_scan(
        state: &CustomScanStateWrapper<Self>,
        _ancestors: *mut pg_sys::List,
        explainer: &mut Explainer,
    ) {
        explainer.add_text("Table", state.custom_state().heaprelname());
        explainer.add_text("Index", state.custom_state().indexrelname());
        if explainer.is_costs() {
            explainer.add_unsigned_integer(
                "Segment Count",
                state.custom_state().segment_count as u64,
                None,
            );
        }

        if explainer.is_analyze() {
            explainer.add_unsigned_integer(
                "Heap Fetches",
                state
                    .custom_state()
                    .visibility_checker
                    .as_ref()
                    .map_or(0, |vc| vc.heap_tuple_check_count) as u64,
                None,
            );
            if explainer.is_verbose() {
                explainer.add_unsigned_integer(
                    "Virtual Tuples",
                    state.custom_state().virtual_tuple_count as u64,
                    None,
                );
                explainer.add_unsigned_integer(
                    "Invisible Tuples",
                    state
                        .custom_state()
                        .visibility_checker
                        .as_ref()
                        .map_or(0, |vc| vc.invisible_tuple_count) as u64,
                    None,
                );
                if let Some(explain_data) = &state.custom_state().parallel_explain_data {
                    explainer.add_json("Parallel Workers", &explain_data.workers);
                }
            }
        }

        explainer.add_text(
            "Exec Method",
            state
                .custom_state()
                .exec_method_name()
                .split("::")
                .last()
                .unwrap(),
        );
        exec_methods::fast_fields::explain(state, explainer);

        explainer.add_bool("Scores", state.custom_state().need_scores());
        if let Some(orderby_info) = state.custom_state().orderby_info().as_ref() {
            explainer.add_text(
                "   TopK Order By",
                orderby_info
                    .iter()
                    .map(|oi| match oi {
                        OrderByInfo {
                            feature: OrderByFeature::Field(fieldname),
                            direction,
                            ..
                        } => {
                            format!("{fieldname} {}", direction.as_ref())
                        }
                        OrderByInfo {
                            feature: OrderByFeature::Var { name, .. },
                            direction,
                            ..
                        } => {
                            format!("{} {}", name.as_deref().unwrap_or("?"), direction.as_ref())
                        }
                        OrderByInfo {
                            feature: OrderByFeature::Score { .. },
                            direction,
                            ..
                        } => {
                            format!("pdb.score() {}", direction.as_ref())
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(", "),
            );
        }

        if let Some(limit) = state.custom_state().limit() {
            explainer.add_unsigned_integer("   TopK Limit", limit as u64, None);
            if explainer.is_analyze() {
                explainer.add_unsigned_integer(
                    "   Queries",
                    state.custom_state().total_query_count().try_into().unwrap(),
                    None,
                );
            }
        }

        // Add a flag to indicate if the query is a full index scan
        let base_query = state.custom_state().base_search_query_input();

        // Only process the query if it's initialized
        // For EXPLAIN without ANALYZE, the query might not be initialized yet
        if !matches!(base_query, SearchQueryInput::Uninitialized) {
            if base_query.is_full_scan_query() {
                explainer.add_bool("Full Index Scan", true);
            }

            // Show query with integrated estimates if GUC is enabled and verbose
            if gucs::explain_recursive_estimates() && explainer.is_verbose() {
                // Get or create a search reader for estimates.
                // - EXPLAIN ANALYZE: search_reader is already initialized by begin_custom_scan
                // - EXPLAIN (without ANALYZE): search_reader is None, so we create a temporary
                //   reader using MvccSatisfies::LargestSegment for estimation purposes only
                let query_tree =
                    if let Some(search_reader) = state.custom_state().search_reader.as_ref() {
                        // EXPLAIN ANALYZE: use the existing search reader
                        search_reader
                            .build_query_tree_with_estimates(base_query.clone())
                            .expect("building query tree with estimates should not fail")
                    } else {
                        // EXPLAIN (without ANALYZE): create a temporary reader for estimates
                        let indexrel = state
                            .custom_state()
                            .indexrel
                            .as_ref()
                            .expect("indexrel should be open");

                        let temp_reader = SearchIndexReader::open_with_context(
                            indexrel,
                            base_query.clone(),
                            false,                         // don't need scores for estimates
                            MvccSatisfies::LargestSegment, // Use largest segment for estimation
                            None,                          // No expr_context needed for estimates
                            None,                          // No planstate needed for estimates
                        )
                        .expect("opening temporary search reader for estimates should not fail");

                        temp_reader
                            .build_query_tree_with_estimates(base_query.clone())
                            .expect("building query tree with estimates should not fail")
                    };

                explainer.add_query_with_estimates(&query_tree);
            } else {
                // Regular display without estimates
                explainer.add_query(base_query);
            }
        } else {
            explainer.add_text("Tantivy Query", "(query not yet initialized)");
        }
    }

    fn begin_custom_scan(
        state: &mut CustomScanStateWrapper<Self>,
        estate: *mut pg_sys::EState,
        eflags: i32,
    ) {
        unsafe {
            // open the heap and index relations with the proper locks
            let rte = pg_sys::exec_rt_fetch(state.custom_state().execution_rti, estate);
            assert!(!rte.is_null());
            let lockmode = (*rte).rellockmode as pg_sys::LOCKMODE;

            state.custom_state_mut().open_relations(lockmode);

            // For EXPLAIN ANALYZE queries, we need to continue with full initialization
            // For EXPLAIN-only (without ANALYZE), begin_custom_scan is not called at all
            if eflags & (pg_sys::EXEC_FLAG_EXPLAIN_ONLY as i32) != 0 {
                // don't do anything else if we're only explaining the query
                return;
            }

            // setup the structures we need to do mvcc checking and heap fetching
            state.custom_state_mut().visibility_checker =
                Some(VisibilityChecker::with_rel_and_snap(
                    state.custom_state().heaprel(),
                    pg_sys::GetActiveSnapshot(),
                ));
            state.custom_state_mut().doc_from_heap_state =
                Some(HeapFetchState::new(state.custom_state().heaprel()));

            // and finally, get the custom scan itself properly initialized
            let tupdesc = state.custom_state().heaptupdesc();
            let planstate = state.planstate();

            pg_sys::ExecInitScanTupleSlot(
                estate,
                addr_of_mut!(state.csstate.ss),
                tupdesc,
                pg_sys::table_slot_callbacks(state.custom_state().heaprel().as_ptr()),
            );

            // On PG15, ExecInitCustomScan hardcodes &TTSOpsVirtual for the scan slot
            // (there's no slotOps override mechanism). ExecInitQual then compiles
            // plan.qual expressions assuming virtual slots. Since we just replaced
            // the slot with BufferHeapTuple above, we must re-initialize the qual
            // so expressions are compiled for the correct slot type.
            #[cfg(feature = "pg15")]
            {
                let plan = state.csstate.ss.ps.plan;
                if !(*plan).qual.is_null() {
                    state.csstate.ss.ps.qual =
                        pg_sys::ExecInitQual((*plan).qual, state.planstate());
                }
            }

            pg_sys::ExecInitResultTypeTL(addr_of_mut!(state.csstate.ss.ps));
            pg_sys::ExecAssignProjectionInfo(
                state.planstate(),
                (*state.csstate.ss.ss_ScanTupleSlot).tts_tupleDescriptor,
            );

            state
                .custom_state_mut()
                .init_expr_context(estate, planstate);
            state.runtime_context = state.csstate.ss.ps.ps_ExprContext;
        }
    }

    fn rescan_custom_scan(state: &mut CustomScanStateWrapper<Self>) {
        Self::init_search_reader(state);
        state.custom_state_mut().reset();
    }

    #[allow(clippy::blocks_in_conditions)]
    fn exec_custom_scan(state: &mut CustomScanStateWrapper<Self>) -> *mut pg_sys::TupleTableSlot {
        if state.custom_state().search_reader.is_none() {
            Self::init_search_reader(state);
        }

        loop {
            let exec_method = state.custom_state_mut().exec_method_mut();

            // get the next matching document from our search results and look for it in the heap
            match exec_method.next(state.custom_state_mut()) {
                // reached the end of the SearchResults
                ExecState::Eof => {
                    return std::ptr::null_mut();
                }

                // SearchResults found a match
                ExecState::FromHeap {
                    ctid,
                    score,
                    doc_address: _,
                } => {
                    unsafe {
                        let slot = match check_visibility(state, ctid, state.scanslot().cast()) {
                            // the ctid is visible
                            Some(slot) => {
                                exec_method.increment_visible();
                                slot
                            }

                            // the ctid is not visible
                            None => {
                                continue;
                            }
                        };

                        // Evaluate executor-level quals (e.g., RLS policy SubPlan expressions)
                        // that couldn't be pushed into the tantivy query.
                        // These are set as plan.qual in plan_custom_path.
                        if !satisfies_subplan_quals(state, slot) {
                            continue;
                        }

                        let needs_special_projection = state.custom_state().need_scores()
                            || state.custom_state().need_snippets()
                            || state.custom_state().window_aggregate_results.is_some();

                        if !needs_special_projection {
                            //
                            // we don't need scores, snippets, or window aggregates
                            // do the projection and return
                            //

                            (*(*state.projection_info()).pi_exprContext).ecxt_scantuple = slot;
                            return pg_sys::ExecProject(state.projection_info());
                        } else {
                            //
                            // we do need scores or snippets
                            //
                            // replace their placeholder values and then rebuild the ProjectionInfo
                            // and project it
                            //

                            let mut per_tuple_context = PgMemoryContexts::For(
                                (*(*state.projection_info()).pi_exprContext).ecxt_per_tuple_memory,
                            );
                            per_tuple_context.reset();

                            if state.custom_state().need_scores() {
                                let const_score_node = state
                                    .custom_state()
                                    .const_score_node
                                    .expect("const_score_node should be set");
                                (*const_score_node).constvalue = score.into_datum().unwrap();
                                (*const_score_node).constisnull = false;
                            }

                            // Update window aggregate values
                            if let Some(agg_results) =
                                &state.custom_state().window_aggregate_results
                            {
                                for (te_idx, datum) in agg_results {
                                    if let Some(const_node) =
                                        state.custom_state().const_window_agg_nodes.get(te_idx)
                                    {
                                        (**const_node).constvalue = *datum;
                                        (**const_node).constisnull = false;
                                    }
                                }
                            }

                            // finally, do the projection
                            return per_tuple_context.switch_to(|_| {
                                // TODO: We go _back_ to the heap to get snippet information here
                                // inside of `make_snippet` and `get_snippet_positions`. It's possible
                                // that we could use a wider tuple slot to fetch the extra columns that
                                // we need during our initial lookup above (but then we'd need to copy
                                // into the correctly shaped slot for this scan).
                                maybe_project_snippets(state.custom_state(), ctid);

                                let planstate = state.planstate();

                                (*(*state.projection_info()).pi_exprContext).ecxt_scantuple = slot;
                                let proj_info = pg_sys::ExecBuildProjectionInfo(
                                    state
                                        .custom_state()
                                        .placeholder_targetlist
                                        .expect("placeholder_targetlist must be set"),
                                    (*planstate).ps_ExprContext,
                                    (*planstate).ps_ResultTupleSlot,
                                    planstate,
                                    (*state.csstate.ss.ss_ScanTupleSlot).tts_tupleDescriptor,
                                );
                                pg_sys::ExecProject(proj_info)
                            });
                        }
                    }
                }

                ExecState::Virtual { slot } => {
                    state.custom_state_mut().virtual_tuple_count += 1;
                    return slot;
                }
            }
        }
    }

    fn shutdown_custom_scan(state: &mut CustomScanStateWrapper<Self>) {
        if let Some(parallel_state) = state.custom_state().parallel_state {
            state.custom_state_mut().parallel_explain_data =
                Some(unsafe { (*parallel_state).explain_data() });
        }
    }

    fn end_custom_scan(state: &mut CustomScanStateWrapper<Self>) {
        // get some things dropped now
        drop(state.custom_state_mut().visibility_checker.take());
        drop(state.custom_state_mut().doc_from_heap_state.take());
        drop(state.custom_state_mut().search_reader.take());
        drop(std::mem::take(
            &mut state.custom_state_mut().snippet_generators,
        ));

        state.custom_state_mut().heaprel.take();
        state.custom_state_mut().indexrel.take();
    }
}

/// Returns true if LIMIT 1 was injected by PostgreSQL's MIN/MAX optimization.
///
/// PostgreSQL rewrites `SELECT MIN(col) FROM t` into a subquery like:
/// `SELECT col AS agg_target FROM t ORDER BY col LIMIT 1`
///
/// This detection is intentionally strict to avoid suppressing warnings for user-written LIMITs.
unsafe fn is_minmax_implicit_limit(root: *mut pg_sys::PlannerInfo) -> bool {
    let Some(root) = root.as_ref() else {
        return false;
    };
    let Some(parse) = root.parse.as_ref() else {
        return false;
    };

    // Must be a subquery (has parent) where parent has aggregates but this subquery doesn't.
    // This matches MIN/MAX rewrite: parent has the aggregate, child is the LIMIT 1 subquery.
    let Some(parent_parse) = root.parent_root.as_ref().and_then(|p| p.parse.as_ref()) else {
        return false;
    };
    if parse.hasAggs || !parent_parse.hasAggs || !parse.limitOffset.is_null() {
        return false;
    }

    // Must have LIMIT 1 as a constant int8 (PostgreSQL's MIN/MAX always uses exactly LIMIT 1).
    let Some(limit_const) = nodecast!(Const, T_Const, parse.limitCount) else {
        return false;
    };
    if (*limit_const).constisnull
        || (*limit_const).consttype != pg_sys::INT8OID
        || i64::from_datum((*limit_const).constvalue, false) != Some(1)
    {
        return false;
    }

    // Must have exactly one target column named "agg_target" - PostgreSQL's internal name
    // for the synthetic column created during MIN/MAX rewrite.
    if parse.targetList.is_null() || parse.sortClause.is_null() {
        return false;
    }
    let target_list = PgList::<pg_sys::TargetEntry>::from_pg(parse.targetList);
    let Some(target_entry) = target_list.get_ptr(0).filter(|_| target_list.len() == 1) else {
        return false;
    };
    if (*target_entry).resname.is_null()
        || CStr::from_ptr((*target_entry).resname).to_bytes() != b"agg_target"
    {
        return false;
    }

    // Sort clause must reference the single target column (ORDER BY for MIN/MAX).
    let sort_list = PgList::<pg_sys::SortGroupClause>::from_pg(parse.sortClause);
    let Some(sort_clause) = sort_list.get_ptr(0).filter(|_| sort_list.len() == 1) else {
        return false;
    };
    (*sort_clause).tleSortGroupRef == (*target_entry).ressortgroupref
}

///
/// Validates whether a query that should use Top K scan is actually using it.
///
/// When `paradedb.check_topk_scan` is enabled, this function checks if a query with LIMIT
/// that uses ParadeDB's search operators is using the Top K execution method. If Top K was
/// expected but not chosen, it logs a warning with diagnostic information to help developers
/// identify performance issues.
///
/// # Performance Note
/// This function has minimal overhead as it returns early when the GUC is disabled.
fn validate_topk_expectation(
    privdata: &PrivateData,
    topk_pathkey_info: &PathKeyInfo,
    limit_is_explicit: bool,
    chosen_method: &ExecMethodType,
    table_name: &str,
) {
    // Fast path: if validation is disabled, return immediately
    if !crate::gucs::check_topk_scan() {
        return;
    }

    // Check if this query should be using Top K
    let has_limit = privdata.limit().is_some();
    let has_search_query = privdata.query().is_some();
    let no_group_by = privdata.window_aggregates().is_empty();

    // Top K is expected when we have: explicit LIMIT + search query + no GROUP BY
    let should_use_topk = has_limit && limit_is_explicit && has_search_query && no_group_by;

    // Check if we actually got Top K
    let is_using_topk = matches!(chosen_method, ExecMethodType::TopK { .. });

    // If Top K is not expected or we're already using Top K, nothing to warn about
    if !should_use_topk || is_using_topk {
        return;
    }

    // At this point: should_use_topk is true AND we're not using Top K - emit warning
    let limit = privdata.limit().unwrap();
    let method_name = match chosen_method {
        ExecMethodType::Normal => "Normal",
        ExecMethodType::Columnar { .. } => "Columnar",
        ExecMethodType::TopK { .. } => "TopK",
    };

    let (reason, remedies) = match topk_pathkey_info {
        PathKeyInfo::Unusable(UnusableReason::TooManyColumns { count, max }) => (
            format!(
                "ORDER BY has {} columns but Top K supports maximum {}",
                count, max
            ),
            format!("Reduce ORDER BY columns to {} or fewer", max),
        ),
        PathKeyInfo::Unusable(UnusableReason::PrefixOnly { matched }) => (
            format!(
                "only partial prefix of ORDER BY can be pushed down ({} columns matched)",
                matched
            ),
            "Ensure all ORDER BY columns are indexed with pdb::literal tokenizer for strings, \
                 or verify that normalizer/collation matches the index"
                .to_string(),
        ),
        PathKeyInfo::Unusable(UnusableReason::NotSortable) => (
            "ORDER BY columns cannot be pushed down to the index".to_string(),
            "Ensure ORDER BY columns are indexed. Numeric columns are fast by default. \
                 For string columns, use pdb.literal tokenizer"
                .to_string(),
        ),
        PathKeyInfo::UsablePrefix(matched) => (
            format!(
                "only partial prefix of ORDER BY can be pushed down ({} of {} columns)",
                matched.len(),
                privdata
                    .maybe_orderby_info()
                    .as_ref()
                    .map_or(0, |o| o.len()),
            ),
            "Ensure all ORDER BY columns are indexed with pdb::literal tokenizer for strings, \
                 or verify that normalizer/collation matches the index"
                .to_string(),
        ),
        PathKeyInfo::None => (
            // This case should normally use Top K with no ordering
            "unknown reason (no pathkeys but should still use Top K)".to_string(),
            "This is unexpected - please report this issue".to_string(),
        ),
        PathKeyInfo::UsableAll(_) => (
            "unknown reason (pathkeys are usable)".to_string(),
            "This is unexpected - please report this issue".to_string(),
        ),
    };

    BaseScan::add_planner_warning(
        format!(
            "Query has LIMIT {} but is not using Top K scan (using {} instead). \
             Reason: {}. \
             This may cause poor performance on large datasets. \
             Remedies: {}. \
             To disable this warning: SET paradedb.check_topk_scan = false",
            limit, method_name, reason, remedies
        ),
        table_name,
    );
}

///
/// Choose and return an ExecMethodType based on the properties of the builder at planning time.
///
/// If the query can return "fast fields", make that determination here, falling back to the
/// [`NormalScanExecState`] if not.
///
/// We support [`ColumnarExecState`] when there are a mix of string and numeric fast fields.
///
/// If we have failed to extract all relevant information at planning time, then the fast-field
/// execution methods might still fall back to `Normal` at execution time: see the notes in
/// `assign_exec_method` and `compute_exec_which_fast_fields`.
///
/// `pdb.score()`, `ctid`, and `tableoid` are considered fast fields for the purposes of
/// these specialized [`ExecMethod`]s.
///
fn choose_exec_method(
    privdata: &PrivateData,
    topk_pathkey_info: &PathKeyInfo,
    limit_is_explicit: bool,
    table_name: &str,
    has_sort_by_pathkey: bool,
) -> Vec<ExecMethodType> {
    // See if we can use Top K.
    if let Some(limit) = privdata.limit() {
        if let Some(orderby_info) = privdata.maybe_orderby_info() {
            // having a valid limit and sort direction means we can do a Top K query
            // and Top K can do snippets
            let method = ExecMethodType::TopK {
                heaprelid: privdata.heaprelid().expect("heaprelid must be set"),
                limit,
                orderby_info: Some(orderby_info.clone()),
                window_aggregates: privdata.window_aggregates().clone(),
            };
            validate_topk_expectation(
                privdata,
                topk_pathkey_info,
                limit_is_explicit,
                &method,
                table_name,
            );
            return vec![method];
        }
        if matches!(topk_pathkey_info, PathKeyInfo::None) {
            // we have a limit but no pathkeys at all. we can still go through our "top k"
            // machinery, but getting "limit" (essentially) random docs, which is what the user
            // asked for
            let method = ExecMethodType::TopK {
                heaprelid: privdata.heaprelid().expect("heaprelid must be set"),
                limit,
                orderby_info: None,
                window_aggregates: privdata.window_aggregates().clone(),
            };
            validate_topk_expectation(
                privdata,
                topk_pathkey_info,
                limit_is_explicit,
                &method,
                table_name,
            );
            return vec![method];
        }
    }

    // Otherwise, see if we can use a fast fields method.
    let is_capable = fast_fields::is_columnar_capable(privdata);
    if is_capable {
        let mut methods = Vec::new();

        // Always create the Unsorted variant
        methods.push(ExecMethodType::Columnar {
            which_fast_fields: privdata.planned_which_fast_fields().clone().unwrap(),
            limit: privdata.limit(),
            sort_order: None,
        });

        // Check if the index has a sort_by configuration (and sorting is enabled)
        // and we have a matching pathkey from the query
        if gucs::is_columnar_sort_enabled() && has_sort_by_pathkey {
            let sort_order = privdata.indexrelid().and_then(|indexrelid| {
                let indexrel =
                    PgSearchRelation::with_lock(indexrelid, pg_sys::AccessShareLock as _);
                let sort_by_fields = indexrel.options().sort_by();
                // Currently only single-field sorting is supported
                sort_by_fields.into_iter().next()
            });

            if let Some(sort_order) = sort_order {
                methods.push(ExecMethodType::Columnar {
                    which_fast_fields: privdata.planned_which_fast_fields().clone().unwrap(),
                    limit: privdata.limit(),
                    sort_order: Some(sort_order),
                });
            }
        }

        // Validate expectations for the first method (Unsorted)
        validate_topk_expectation(
            privdata,
            topk_pathkey_info,
            limit_is_explicit,
            &methods[0],
            table_name,
        );

        return methods;
    }

    // Else, fall back to normal execution
    let method = ExecMethodType::Normal;
    validate_topk_expectation(
        privdata,
        topk_pathkey_info,
        limit_is_explicit,
        &method,
        table_name,
    );
    vec![method]
}

///
/// Creates and assigns the execution method which was chosen at planning time.
///
/// If a fast-fields execution method was chosen at planning time, we might still fall back to
/// NormalScanExecState if we fail to extract the superset of fields during planning time which was
/// needed at execution time.
///
fn assign_exec_method(builder: &mut CustomScanStateBuilder<BaseScan, PrivateData>) {
    match builder.custom_state_ref().exec_method_type.clone() {
        ExecMethodType::Normal => builder
            .custom_state()
            .assign_exec_method(NormalScanExecState::default(), Some(ExecMethodType::Normal)),
        ExecMethodType::TopK {
            heaprelid,
            limit,
            orderby_info,
            window_aggregates: _,
        } => builder.custom_state().assign_exec_method(
            exec_methods::top_k::TopKScanExecState::new(heaprelid, limit, orderby_info),
            None,
        ),

        ExecMethodType::Columnar {
            which_fast_fields,
            limit,
            sort_order,
        } => {
            // Compute effective sort_order at execution time from PrivateData's use_sorted_path().
            // sort_order is only populated for sorted paths, but we still guard against
            // planner/executor mismatches using use_sorted_path().
            let runtime_sorted = sort_order.is_some() && builder.custom_private().use_sorted_path();

            // Safety check: ensure we don't drop sorted output when we claimed it at planning time.
            // This catches bugs where the planner was told output would be sorted but execution
            // cannot provide it.
            if sort_order.is_some() && !runtime_sorted {
                panic!(
                    "Claimed sorted output at planning time, but unable to provide it at \
                    execution time. This indicates a planner/executor mismatch."
                );
            }

            // Pass sort_order only if we're actually using sorted execution
            let effective_sort_order = if runtime_sorted { sort_order } else { None };

            if let Some(which_fast_fields) =
                compute_exec_which_fast_fields(builder, which_fast_fields)
            {
                // Calculate extra fields needed for execution (e.g. sort column)
                // that are not in the projected fields (which_fast_fields).
                let mut extra_fast_fields = Vec::new();

                if let Some(sort) = &effective_sort_order {
                    // Check if sort field is already projected
                    // Note: This naive name check works because we only support single-column sort by name
                    let is_projected = which_fast_fields
                        .iter()
                        .any(|f| f.name() == sort.field_name.as_ref());

                    if !is_projected {
                        // Retrieve the sort field from the planned fast fields
                        let planned_fields = match &builder.custom_state_ref().exec_method_type {
                            ExecMethodType::Columnar {
                                which_fast_fields, ..
                            } => which_fast_fields,
                            _ => unreachable!(),
                        };

                        if let Some(ff) = planned_fields
                            .iter()
                            .find(|f| f.name() == sort.field_name.as_ref())
                        {
                            extra_fast_fields.push(ff.clone());
                        }
                    }
                }

                builder.custom_state().assign_exec_method(
                    exec_methods::fast_fields::columnar::ColumnarExecState::new(
                        which_fast_fields,
                        extra_fast_fields,
                        limit,
                        effective_sort_order,
                    ),
                    None,
                )
            } else {
                builder.custom_state().assign_exec_method(
                    NormalScanExecState::default(),
                    Some(ExecMethodType::Normal),
                )
            }
        }
    }
}

///
/// Computes the execution time `which_fast_fields`, which are validated to be a subset of the
/// planning time `which_fast_fields`. If it's not the case, we return `None` to indicate that
/// we should fall back to the `Normal` execution mode.
///
fn compute_exec_which_fast_fields(
    builder: &mut CustomScanStateBuilder<BaseScan, PrivateData>,
    planned_which_fast_fields: HashSet<WhichFastField>,
) -> Option<Vec<WhichFastField>> {
    let target_list = builder.target_list().as_ptr();
    let exec_which_fast_fields = unsafe {
        let custom_state = builder.custom_state();
        let indexrel = custom_state.indexrel();
        let execution_rti = custom_state.execution_rti;
        let heaprel = custom_state.heaprel();
        //
        // In order for our planned ExecMethodType to be accurate, this must always be a
        // subset of the fast fields which were extracted at planning time.
        exec_methods::fast_fields::collect_fast_fields(
            target_list,
            // At this point, all fast fields which we need to extract are listed directly
            // in our execution-time target list, so there is no need to extract from other
            // positions.
            &HashSet::default(),
            execution_rti,
            heaprel,
            indexrel,
            true,
        )
    };

    if fast_fields::is_all_special_or_junk_fields(&exec_which_fast_fields) {
        // In some cases, enough columns are pruned between planning and execution that there
        // is no point actually using fast fields, and we can fall back to `Normal`.
        //
        // TODO: In order to implement https://github.com/paradedb/paradedb/issues/2623, we will
        // need to differentiate these cases, so that we can always emit the sort order that we
        // claimed.
        return None;
    }

    let missing_fast_fields = exec_which_fast_fields
        .iter()
        .filter(|ff| !planned_which_fast_fields.contains(ff))
        .collect::<Vec<_>>();

    if !missing_fast_fields.is_empty() {
        pgrx::log!(
            "Failed to extract all fast fields at planning time: \
             was missing {missing_fast_fields:?} from {planned_which_fast_fields:?} \
             Falling back to Normal execution.",
        );
        return None;
    }

    Some(exec_which_fast_fields)
}

/// Use the [`VisibilityChecker`] to verify the ctid is visible and fetch the tuple into the slot.
/// Returns the slot if visible, None if the tuple is deleted/not visible.
/// See [`VisibilityChecker`] docs for details on two-layer visibility.
#[inline(always)]
fn check_visibility(
    state: &mut CustomScanStateWrapper<BaseScan>,
    ctid: u64,
    bslot: *mut pg_sys::BufferHeapTupleTableSlot,
) -> Option<*mut pg_sys::TupleTableSlot> {
    state
        .custom_state_mut()
        .visibility_checker()
        .exec_if_visible(ctid, bslot.cast(), move |_| bslot.cast())
}

/// Evaluate executor-level quals (e.g., RLS policy SubPlan expressions) that couldn't be pushed
/// into the tantivy query. These are set as `plan.qual` in `plan_custom_path`.
/// Returns `true` if qual passes (or no qual exists), `false` if the row should be skipped.
#[inline(always)]
unsafe fn satisfies_subplan_quals(
    state: &mut CustomScanStateWrapper<BaseScan>,
    slot: *mut pg_sys::TupleTableSlot,
) -> bool {
    let qual = state.csstate.ss.ps.qual;
    if qual.is_null() {
        return true;
    }
    let econtext = (*state.projection_info()).pi_exprContext;
    (*econtext).ecxt_scantuple = slot;
    pg_sys::slot_getallattrs(slot);
    pg_sys::ExecQual(qual, econtext)
}

/// Inject ParadeDB-specific placeholders (score, snippets, window aggregates) into the tuple slot
unsafe fn inject_pdb_placeholders(state: &mut CustomScanStateWrapper<BaseScan>) {
    let need_scores = state.custom_state().need_scores();
    let need_snippets = state.custom_state().need_snippets();
    let has_window_aggs = !state.custom_state().window_aggregates.is_empty();

    if !need_scores && !need_snippets && !has_window_aggs {
        // nothing to inject, use whatever we originally setup as our ProjectionInfo
        return;
    }

    // inject score and/or snippet placeholder [`pg_sys::Const`] nodes into what is a copy of the Plan's
    // targetlist.  We store this in our custom state's "placeholder_targetlist" for use during the
    // forced projection we must do later.
    let planstate = state.planstate();

    let (targetlist, const_score_node, const_snippet_nodes) = inject_placeholders(
        (*(*planstate).plan).targetlist,
        state.custom_state().planning_rti,
        state.custom_state().score_funcoids,
        state.custom_state().snippet_funcoids,
        state.custom_state().snippets_funcoids,
        state.custom_state().snippet_positions_funcoids,
        &state.custom_state().var_attname_lookup,
        &state.custom_state().snippet_generators,
    );

    // Now inject window aggregate placeholders
    let (targetlist, const_window_agg_nodes) = if !state.custom_state().window_aggregates.is_empty()
    {
        inject_window_aggregate_placeholders(targetlist, &state.custom_state().window_aggregates)
    } else {
        (targetlist, HashMap::default())
    };

    state.custom_state_mut().placeholder_targetlist = Some(targetlist);
    state.custom_state_mut().const_score_node = Some(const_score_node);
    state.custom_state_mut().const_snippet_nodes = const_snippet_nodes;
    state.custom_state_mut().const_window_agg_nodes = const_window_agg_nodes;
}

/// Inject placeholder Const nodes for window aggregates at execution time
/// At this point, the WindowFunc has been replaced with paradedb.window_agg(json) calls
/// This function finds those calls (which may be wrapped in other functions)
/// and replaces them with placeholder Const nodes that will be filled in during execution.
unsafe fn inject_window_aggregate_placeholders(
    targetlist: *mut pg_sys::List,
    window_aggs: &[WindowAggregateInfo],
) -> (*mut pg_sys::List, HashMap<usize, *mut pg_sys::Const>) {
    let mut const_nodes = HashMap::default();
    let tlist = PgList::<pg_sys::TargetEntry>::from_pg(targetlist);
    let window_agg_procid = window_agg_oid();

    // If window_agg function doesn't exist yet, return original targetlist
    if window_agg_procid == pg_sys::InvalidOid {
        return (targetlist, const_nodes);
    }

    // Process each window aggregate target entry
    let mut new_tlist = PgList::<pg_sys::TargetEntry>::new();

    for (idx, te) in tlist.iter_ptr().enumerate() {
        // Check if this target entry is one of our window aggregates
        let agg_info = window_aggs.iter().find(|a| a.target_entry_index == idx);

        if let Some(agg_info) = agg_info {
            // This target entry should contain a window_agg call (possibly wrapped)
            let (new_expr, const_node_opt) = replace_window_agg_with_const(
                (*te).expr as *mut pg_sys::Node,
                window_agg_procid,
                agg_info.result_type_oid(),
            );

            // Create a new target entry with the modified expression
            let new_te = pg_sys::flatCopyTargetEntry(te);
            (*new_te).expr = new_expr.cast();
            new_tlist.push(new_te);

            if let Some(const_node) = const_node_opt {
                const_nodes.insert(idx, const_node);
            }
        } else {
            // Not a window aggregate - just copy it
            new_tlist.push(te);
        }
    }

    (new_tlist.into_pg(), const_nodes)
}

// Helper function to recursively search and replace window_agg calls
//
// Note: This follows a similar recursive pattern to replace_in_node() in hook.rs,
// but operates at a different stage:
// - That function: Planning stage - replaces WindowFunc → window_agg() placeholder
// - This function: Execution stage - replaces window_agg() → Const placeholder for value injection
//
// TODO: This duplication could potentially be eliminated by moving to UPPERREL_WINDOW handling.
// See https://github.com/paradedb/paradedb/issues/3455
unsafe fn replace_window_agg_with_const(
    node: *mut pg_sys::Node,
    window_agg_procid: pg_sys::Oid,
    result_type_oid: pg_sys::Oid,
) -> (*mut pg_sys::Node, Option<*mut pg_sys::Const>) {
    if node.is_null() {
        return (node, None);
    }

    // Check if this is the window_agg FuncExpr
    if let Some(funcexpr) = nodecast!(FuncExpr, T_FuncExpr, node) {
        if (*funcexpr).funcid == window_agg_procid {
            // Found it! Replace with a Const node
            let const_node = pg_sys::makeConst(
                result_type_oid,
                -1,
                pg_sys::DEFAULT_COLLATION_OID,
                if result_type_oid == pg_sys::INT8OID {
                    8
                } else {
                    -1
                },
                pg_sys::Datum::null(),
                true,                               // constisnull
                result_type_oid == pg_sys::INT8OID, // constbyval (true for INT8)
            );

            return (const_node.cast(), Some(const_node));
        }

        // Not window_agg, but might have window_agg as an argument
        let args = PgList::<pg_sys::Node>::from_pg((*funcexpr).args);
        let mut new_args = PgList::<pg_sys::Node>::new();
        let mut found_const = None;
        let mut modified = false;

        for arg in args.iter_ptr() {
            let (new_arg, const_opt) =
                replace_window_agg_with_const(arg, window_agg_procid, result_type_oid);
            if const_opt.is_some() {
                found_const = const_opt;
                modified = true;
            }
            if new_arg != arg {
                modified = true;
            }
            new_args.push(new_arg);
        }

        if modified {
            // Create a new FuncExpr with modified arguments
            let new_funcexpr = pg_sys::makeFuncExpr(
                (*funcexpr).funcid,
                (*funcexpr).funcresulttype,
                new_args.into_pg(),
                (*funcexpr).funccollid,
                (*funcexpr).inputcollid,
                (*funcexpr).funcformat,
            );
            return (new_funcexpr.cast(), found_const);
        }
    }

    (node, None)
}

/// Determine whether there are any pathkeys at all, and whether we might be able to push down
/// ordering in Top K.
///
/// If between 1 and 3 pathkeys are declared, and are indexed as fast, then return
/// `UsableAll(Vec<OrderByStyles>)` for them for use in Top K.
///
/// This function must be kept in sync with `validate_topk_compatibility` in `hook.rs` to ensure
/// that queries validated during the planner hook phase can be executed by the custom scan.
unsafe fn pullup_topk_pathkeys(
    rti: pg_sys::Index,
    schema: &SearchIndexSchema,
    root: *mut pg_sys::PlannerInfo,
) -> PathKeyInfo {
    match extract_pathkey_styles_with_sortability_check(
        root,
        rti,
        schema,
        |search_field| search_field.is_raw_sortable(),
        |search_field| search_field.is_lower_sortable(),
    ) {
        PathKeyInfo::UsableAll(styles) if styles.len() <= MAX_TOPK_FEATURES => {
            // Top K is the base scan's only executor which supports sorting, and supports up to
            // MAX_TOPK_FEATURES order-by clauses.
            PathKeyInfo::UsableAll(styles)
        }
        PathKeyInfo::UsableAll(ref styles) => {
            // Too many pathkeys were extracted.
            PathKeyInfo::Unusable(UnusableReason::TooManyColumns {
                count: styles.len(),
                max: MAX_TOPK_FEATURES,
            })
        }
        PathKeyInfo::UsablePrefix(ref prefix) => {
            // Top K cannot execute for a prefix of pathkeys, because it eliminates results before
            // the suffix of the pathkey comes into play.
            PathKeyInfo::Unusable(UnusableReason::PrefixOnly {
                matched: prefix.len(),
            })
        }
        pki @ (PathKeyInfo::None | PathKeyInfo::Unusable(_)) => pki,
    }
}

/// Gather all columns referenced by the specified RTE (Range Table Entry) throughout the query.
/// This gives us a more complete picture than just looking at the target list.
///
/// This function is critical for issue #2505/#2556 where we need to detect all columns used in JOIN
/// conditions to ensure we select the right execution method. Previously, only looking at the
/// target list would miss columns referenced in JOIN conditions, leading to execution-time errors.
///
unsafe fn collect_maybe_fast_field_referenced_columns(
    rte_index: pg_sys::Index,
    rel: *mut pg_sys::RelOptInfo,
) -> HashSet<pg_sys::AttrNumber> {
    let mut referenced_columns = HashSet::default();

    // Check reltarget exprs.
    let reltarget_exprs = PgList::<pg_sys::Expr>::from_pg((*(*rel).reltarget).exprs);
    for rte in reltarget_exprs.iter_ptr() {
        if let Some(var) = nodecast!(Var, T_Var, rte) {
            if (*var).varno as u32 == rte_index {
                referenced_columns.insert((*var).varattno);
            }
        }
        // NOTE: Unless we encounter the fallback in `compute_exec_which_fast_fields`, then we
        // can be reasonably confident that directly inspecting Vars is sufficient. We haven't
        // seen it yet in the wild.
    }

    referenced_columns
}

#[rustfmt::skip]
/// Check if the base query has search predicates for the current table's index
fn base_query_has_search_predicates(
    query: &SearchQueryInput,
    current_index_oid: pg_sys::Oid,
) -> bool {
    match query {
        SearchQueryInput::All => false,
        SearchQueryInput::Uninitialized => false,
        SearchQueryInput::Empty => false,

        SearchQueryInput::WithIndex { oid, query } => {
            // Only consider search predicates for the current table's index
            if *oid == current_index_oid {
                // This is a search predicate for our index
                // Check the inner query directly for range vs search predicates
                base_query_has_search_predicates(query, current_index_oid)
            } else {
                // This is a search predicate for a different index, ignore it
                false
            }
        }

        // Boolean queries need recursive checking
        SearchQueryInput::Boolean {
            must,
            should,
            must_not,
        } => {
            must.iter()
                .any(|q| base_query_has_search_predicates(q, current_index_oid))
                || should
                    .iter()
                    .any(|q| base_query_has_search_predicates(q, current_index_oid))
                || must_not
                    .iter()
                    .any(|q| base_query_has_search_predicates(q, current_index_oid))
        }

        // Wrapper queries need recursive checking
        SearchQueryInput::Boost { query, .. } => {
            base_query_has_search_predicates(query, current_index_oid)
        }
        SearchQueryInput::ConstScore { query, .. } => {
            base_query_has_search_predicates(query, current_index_oid)
        }
        SearchQueryInput::ScoreFilter {
            query: Some(query), ..
        } => base_query_has_search_predicates(query, current_index_oid),
        SearchQueryInput::ScoreFilter { query: None, .. } => false,
        SearchQueryInput::DisjunctionMax { disjuncts, .. } => disjuncts
            .iter()
            .any(|q| base_query_has_search_predicates(q, current_index_oid)),

        // Despite being part of FieldedQuery, these do not use a field, as far as the user knows
        SearchQueryInput::FieldedQuery {  query: pdb::Query::All, ..} |
        SearchQueryInput::FieldedQuery {  query: pdb::Query::Empty, ..} => false,

        // These are NOT search predicates (they're range/exists/other predicates)
        SearchQueryInput::FieldedQuery { query: pdb::Query::Range { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::RangeContains { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::RangeIntersects { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::RangeTerm { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::RangeWithin { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Exists, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::FastFieldRangeWeight { .. }, .. }
        | SearchQueryInput::MoreLikeThis { .. } => false,

        // These are search predicates that use the @@@ operator
        SearchQueryInput::FieldedQuery { query: pdb::Query::ParseWithField { query_string, .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Parse { query_string, .. }, .. } => {
            // For ParseWithField, check if it's a text search or a range query
            !is_range_query_string(query_string)
        }
        SearchQueryInput::Parse { .. }
        | SearchQueryInput::TermSet { .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::UnclassifiedString { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::UnclassifiedArray { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::ScoreAdjusted { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::TermSet { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Term { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Phrase { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::PhraseArray { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Proximity { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::TokenizedPhrase { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::PhrasePrefix { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::FuzzyTerm { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Match { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::MatchArray { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::Regex { .. }, .. }
        | SearchQueryInput::FieldedQuery { query: pdb::Query::RegexPhrase { .. }, .. } => true,

        // // Term with no field is not a search predicate
        // NB:  We don't support unqualified term queries anymore
        // SearchQueryInput::Term { field: None, .. } => false,

        // Postgres expressions are unknown, assume they could be search predicates
        SearchQueryInput::PostgresExpression { .. } => true,

        // HeapFilter contains search predicates
        SearchQueryInput::HeapFilter { indexed_query, .. } => {
            base_query_has_search_predicates(indexed_query, current_index_oid)
        }
    }
}

/// Check if a query string represents a range query (contains operators like >, <, etc.)
fn is_range_query_string(query_string: &str) -> bool {
    // Range queries typically start with operators
    query_string.trim_start().starts_with('>')
        || query_string.trim_start().starts_with('<')
        || query_string.trim_start().starts_with(">=")
        || query_string.trim_start().starts_with("<=")
        || query_string.contains("..")  // Range syntax like "1..10"
        || query_string.contains(" TO ") // Range syntax like "1 TO 10"
}

/// Project configured snippets (if any).
///
/// Must be called inside the per-tuple `MemoryContext`.
unsafe fn maybe_project_snippets(state: &BaseScanState, ctid: u64) {
    if !state.need_snippets() {
        return;
    }

    for (snippet_type, const_snippet_nodes) in &state.const_snippet_nodes {
        match snippet_type {
            SnippetType::SingleText(_, _, _) => {
                let snippet = state.make_snippet(ctid, snippet_type);

                for const_ in const_snippet_nodes {
                    match &snippet {
                        Some(text) => {
                            (**const_).constvalue = text.into_datum().unwrap();
                            (**const_).constisnull = false;
                        }
                        None => {
                            (**const_).constvalue = pg_sys::Datum::null();
                            (**const_).constisnull = true;
                        }
                    }
                }
            }
            SnippetType::MultipleText(_, _, _, _) => {
                let snippets = state.make_snippets(ctid, snippet_type);

                for const_ in const_snippet_nodes {
                    match &snippets {
                        Some(array) => {
                            (**const_).constvalue = array.clone().into_datum().unwrap();
                            (**const_).constisnull = false;
                        }
                        None => {
                            (**const_).constvalue = pg_sys::Datum::null();
                            (**const_).constisnull = true;
                        }
                    }
                }
            }
            SnippetType::Positions(..) => {
                let positions = state.get_snippet_positions(ctid, snippet_type);

                for const_ in const_snippet_nodes {
                    match &positions {
                        Some(positions) => {
                            (**const_).constvalue = positions.clone().into_datum().unwrap();
                            (**const_).constisnull = false;
                        }
                        None => {
                            (**const_).constvalue = pg_sys::Datum::null();
                            (**const_).constisnull = true;
                        }
                    }
                }
            }
        }
    }
}

/// Check if the query contains a LEFT JOIN LATERAL pattern
///
/// This function verifies that the query has the specific structure:
/// `... LEFT JOIN LATERAL (...) ...`
///
/// We verify:
/// 1. The parse tree contains a LEFT JOIN node
/// 2. That LEFT JOIN's right side is marked as LATERAL in the range table
///
/// This enables Top K optimization because LEFT JOIN semantics guarantee all
/// left-side rows are preserved. If WHERE/ORDER BY/LIMIT only reference the
/// left table, we can safely apply Top K to the left scan before the join.
unsafe fn is_left_join_lateral(
    root: *mut pg_sys::PlannerInfo,
    rel: *mut pg_sys::RelOptInfo,
) -> bool {
    // Check if this is a join query
    if !(*root).hasJoinRTEs {
        return false;
    }

    // Check if this is a base relation (not a join itself)
    if (*rel).reloptkind != pg_sys::RelOptKind::RELOPT_BASEREL {
        return false;
    }

    // Check the parse tree for LEFT JOIN patterns with LATERAL
    // We need to verify:
    // 1. There's a LEFT JOIN in the query
    // 2. The right side of that LEFT JOIN is LATERAL
    //
    // We use a combination of checks:
    // - Parse tree: to find LEFT JOIN structure
    // - RTE lateral flag: to confirm the right side is LATERAL

    // First, quickly check if any LATERAL references exist at all
    let simple_rel_array = (*root).simple_rel_array;
    if simple_rel_array.is_null() {
        return false;
    }

    let mut has_lateral = false;
    let simple_rel_array_size = (*root).simple_rel_array_size;
    for i in 1..simple_rel_array_size {
        let other_rel = *simple_rel_array.add(i as usize);
        if !other_rel.is_null() && !(*other_rel).lateral_relids.is_null() {
            has_lateral = true;
            break;
        }
    }

    if !has_lateral {
        return false;
    }

    // Now check the parse tree for LEFT JOIN structure
    let jointree = (*(*root).parse).jointree;
    if jointree.is_null() || (*jointree).fromlist.is_null() {
        return false;
    }

    let fromlist = PgList::<pg_sys::Node>::from_pg((*jointree).fromlist);
    for node in fromlist.iter_ptr() {
        if has_left_join_lateral_pattern(node, (*root).parse) {
            return true;
        }
    }

    false
}

/// Recursively check if a node contains a LEFT JOIN LATERAL pattern
unsafe fn has_left_join_lateral_pattern(
    node: *mut pg_sys::Node,
    query: *mut pg_sys::Query,
) -> bool {
    if node.is_null() {
        return false;
    }

    if let Some(join_expr) = nodecast!(JoinExpr, T_JoinExpr, node) {
        // Check if this is a LEFT JOIN
        if (*join_expr).jointype == pg_sys::JoinType::JOIN_LEFT {
            // Check if the right side has LATERAL
            if is_lateral_subquery((*join_expr).rarg, query) {
                return true;
            }
        }

        // Recursively check nested joins
        if has_left_join_lateral_pattern((*join_expr).larg, query) {
            return true;
        }
        if has_left_join_lateral_pattern((*join_expr).rarg, query) {
            return true;
        }
    }

    false
}

/// Check if a node represents a LATERAL subquery
unsafe fn is_lateral_subquery(node: *mut pg_sys::Node, query: *mut pg_sys::Query) -> bool {
    if node.is_null() || query.is_null() {
        return false;
    }

    // Check if it's a RangeTblRef pointing to a LATERAL RTE
    if let Some(rtref) = nodecast!(RangeTblRef, T_RangeTblRef, node) {
        let rtable = (*query).rtable;
        if !rtable.is_null() {
            let rte = pg_sys::rt_fetch((*rtref).rtindex as pg_sys::Index, rtable);
            if !rte.is_null() && (*rte).lateral {
                return true;
            }
        }
    }

    // For nested joins, recursively check
    if let Some(join_expr) = nodecast!(JoinExpr, T_JoinExpr, node) {
        if is_lateral_subquery((*join_expr).larg, query) {
            return true;
        }
        if is_lateral_subquery((*join_expr).rarg, query) {
            return true;
        }
    }

    false
}

/// Verify WHERE clause only references the left table (current relation)
///
/// This method is used to check whether we can safely push down a LEFT LATERAL JOIN as Top K.
/// Because Top K eliminates rows _before_ the JOIN is actually executed, the WHERE clause (and
/// join condition) may only reference the left hand side of the join to avoid eliminating rows via the
/// limit which would be filtered by conditions on the right hand side.
unsafe fn where_clause_only_references_left(
    root: *mut pg_sys::PlannerInfo,
    rti: pg_sys::Index,
) -> bool {
    // Get WHERE clause
    let quals = if !(*root).parse.is_null()
        && !(*(*root).parse).jointree.is_null()
        && !(*(*(*root).parse).jointree).quals.is_null()
    {
        (*(*(*root).parse).jointree).quals
    } else {
        return true; // No WHERE clause means it only references left
    };

    // Walk the quals to check if they only reference our relation
    #[pgrx::pg_guard]
    unsafe extern "C-unwind" fn walker(
        node: *mut pg_sys::Node,
        data: *mut core::ffi::c_void,
    ) -> bool {
        if node.is_null() {
            return false;
        }

        if let Some(var) = nodecast!(Var, T_Var, node) {
            let rti = *(data as *const pg_sys::Index);
            // If we find a Var that's not from our relation, return true (fail)
            if (*var).varno as i32 != rti as i32 && (*var).varno > 0 {
                return true;
            }
        }

        pg_sys::expression_tree_walker(node, Some(walker), data)
    }

    // If walker returns true, it found a reference to another relation
    !walker(quals, &rti as *const _ as *mut _)
}
