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

use crate::api::version::Version;
use crate::gucs;
use crate::index::kdtree::KdTree;
use crate::index::mvcc::MvccSatisfies;
use crate::index::writer::index::{
    DiskSpaceGuard, IndexWriterConfig, Mergeable, SearchIndexMerger, SerialIndexWriter,
};
use crate::launch_parallel_process;
use crate::parallel_worker::mqueue::MessageQueueSender;
use crate::parallel_worker::{
    chunk_range, ParallelProcess, ParallelState, ParallelStateManager, ParallelStateType,
    ParallelWorker, WorkerStyle,
};
use crate::postgres::build_partitioning::plan_partition_boundaries;
use crate::postgres::composite::CompositeSlotValues;
use crate::postgres::locks::Spinlock;
use crate::postgres::merge::garbage_collect_index;
use crate::postgres::ps_status::{
    set_ps_display_remove_suffix, set_ps_display_suffix, COMMITTING, FINALIZING,
    GARBAGE_COLLECTING, INDEXING, MERGING,
};
use crate::postgres::rel::PgSearchRelation;
use crate::postgres::storage::buffer::BufferManager;
use crate::postgres::storage::metadata::MetaPage;
use crate::postgres::utils::{
    collect_composites_for_unpacking, get_field_value, row_to_search_document,
};
use crate::schema::{CategorizedFieldData, SearchField};
use pgrx::pg_sys::panic::ErrorReport;
use pgrx::{
    check_for_interrupts, function_name, pg_guard, pg_sys, PgLogLevel, PgMemoryContexts,
    PgSqlErrorCode,
};
use std::num::NonZeroUsize;
use std::ptr::{addr_of_mut, NonNull};
use std::sync::OnceLock;
use tantivy::index::SegmentId;
use tantivy::{SegmentMeta, TantivyDocument};

const TUPLES_DONE_BATCH_SIZE: usize = 5;
/// General, immutable configuration used for the workers
#[derive(Copy, Clone, Debug)]
#[repr(C)]
struct WorkerConfig {
    heaprelid: pg_sys::Oid,
    indexrelid: pg_sys::Oid,
    concurrent: bool,
    current_xid: pg_sys::FullTransactionId,
    need_wal: bool,
    next_xid: pg_sys::FullTransactionId,
}
impl ParallelStateType for WorkerConfig {}

/// Type alias that holds a pointer to a [`pg_sys::ParallelTableScanDescData`] which is over-allocated,
/// so the [`usize`] field tells us how big it really is, in bytes
type ScanDesc = (usize, *mut pg_sys::ParallelTableScanDescData);
impl ParallelStateType for pg_sys::ParallelTableScanDescData {}

#[derive(Copy, Clone, Default)]
#[repr(C)]
struct WorkerCoordination {
    mutex: Spinlock,
    nstarted: usize,
    nlaunched: usize,
    ntuples_done: usize,
    nsegments_written: usize,
}

impl ParallelStateType for WorkerCoordination {}
impl WorkerCoordination {
    fn inc_nstarted(&mut self) {
        let _lock = self.mutex.acquire();
        self.nstarted += 1;
    }
    fn nstarted(&mut self) -> usize {
        let _lock = self.mutex.acquire();
        self.nstarted
    }
    fn set_nlaunched(&mut self, nlaunched: usize) {
        let _lock = self.mutex.acquire();
        self.nlaunched = nlaunched;
    }
    fn nlaunched(&mut self) -> usize {
        let _lock = self.mutex.acquire();
        self.nlaunched
    }
    fn add_tuples_done(&mut self, count: usize) {
        let _lock = self.mutex.acquire();
        self.ntuples_done += count;
    }
    fn tuples_done(&mut self) -> usize {
        let _lock = self.mutex.acquire();
        self.ntuples_done
    }
    fn add_segments_written(&mut self) -> usize {
        let _lock = self.mutex.acquire();
        self.nsegments_written += 1;
        self.nsegments_written
    }
}

/// The parallel process for setting up a parallel index build
struct ParallelBuild {
    config: WorkerConfig,
    scandesc: ScanDesc,
    coordination: WorkerCoordination,
    /// The leader's global partition boundaries, serialized with `postcard`. Empty when the
    /// index has no `partition_by`.
    partitioning: Vec<u8>,
}

impl ParallelState for ScanDesc {
    fn type_name(&self) -> &'static str {
        std::any::type_name::<pg_sys::ParallelTableScanDescData>()
    }

    fn size_of(&self) -> usize {
        self.0
    }

    fn as_bytes(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.1 as *const _ as *const u8, self.size_of()) }
    }
}

impl ParallelBuild {
    fn new(
        heaprel: &PgSearchRelation,
        snapshot: pg_sys::Snapshot,
        config: WorkerConfig,
        partitioning: Vec<u8>,
    ) -> Self {
        let scandesc = unsafe {
            let size = size_of::<pg_sys::ParallelTableScanDescData>()
                + pg_sys::table_parallelscan_estimate(heaprel.as_ptr(), snapshot) as usize;
            let scandesc = pg_sys::palloc0(size).cast();
            pg_sys::table_parallelscan_initialize(heaprel.as_ptr(), scandesc, snapshot);
            (size, scandesc)
        };
        Self {
            config,
            scandesc,
            coordination: Default::default(),
            partitioning,
        }
    }
}

impl ParallelProcess for ParallelBuild {
    fn state_values(&self) -> Vec<&dyn ParallelState> {
        vec![
            &self.config,
            &self.scandesc,
            &self.coordination,
            &self.partitioning,
        ]
    }
}

#[derive(serde::Serialize, serde::Deserialize)]
struct WorkerResponse {
    reltuples: f64,
    nmerges: usize,
}

fn deserialize_partitioning(bytes: &[u8]) -> Option<KdTree> {
    if bytes.is_empty() {
        return None;
    }
    Some(
        postcard::from_bytes(bytes)
            .unwrap_or_else(|e| panic!("could not deserialize partition boundaries: {e}")),
    )
}

struct BuildWorker<'a> {
    config: WorkerConfig,
    table_scan_desc: Option<NonNull<pg_sys::TableScanDescData>>,
    coordination: &'a mut WorkerCoordination,
    heaprel: PgSearchRelation,
    indexrel: PgSearchRelation,
    /// Global partition boundaries computed by the leader; `None` without `partition_by`.
    partitioning: Option<KdTree>,
}

impl ParallelWorker for BuildWorker<'_> {
    fn new_parallel_worker(state_manager: ParallelStateManager) -> Self
    where
        Self: Sized,
    {
        let config = state_manager
            .object::<WorkerConfig>(0)
            .expect("should be able to get ParallelBuildConfig from state manager")
            .expect("ParallelBuildConfig should not be NULL");
        let scandesc = state_manager
            .object::<pg_sys::ParallelTableScanDescData>(1)
            .expect("should be able to get ParallelTableScanDesc")
            .expect("ParallelTableScanDesc should not be NULL");
        let coordination = state_manager
            .object::<WorkerCoordination>(2)
            .expect("should be able to get WorkerCoordination")
            .expect("WorkerCoordination should not be NULL");
        let partitioning = state_manager
            .slice::<u8>(3)
            .expect("should be able to get partitioning bytes")
            .expect("partitioning bytes should not be NULL");
        let partitioning = deserialize_partitioning(partitioning);

        unsafe {
            let (heap_lock, index_lock) = if !config.concurrent {
                (pg_sys::ShareLock, pg_sys::AccessExclusiveLock)
            } else {
                (pg_sys::ShareUpdateExclusiveLock, pg_sys::RowExclusiveLock)
            };

            let heaprel =
                PgSearchRelation::with_lock(config.heaprelid, heap_lock as pg_sys::LOCKMODE);
            let mut indexrel =
                PgSearchRelation::with_lock(config.indexrelid, index_lock as pg_sys::LOCKMODE);
            let table_scan_desc = pg_sys::table_beginscan_parallel(heaprel.as_ptr(), scandesc);

            indexrel.set_is_create_index();
            indexrel.set_need_wal(config.need_wal);

            Self {
                config: *config,
                table_scan_desc: NonNull::new(table_scan_desc),
                coordination,
                heaprel,
                indexrel,
                partitioning,
            }
        }
    }

    fn run(mut self, mq_sender: &MessageQueueSender, worker_number: i32) -> anyhow::Result<()> {
        // wait for the leader to tell us how many total workers have been launched
        while self.coordination.nlaunched() == 0 {
            check_for_interrupts!();
            std::thread::yield_now();
        }

        // communicate to the group that we've started
        self.coordination.inc_nstarted();

        let (reltuples, nmerges) = self.do_build(worker_number, false)?;
        Ok(mq_sender.send(serde_json::to_vec(&WorkerResponse { reltuples, nmerges })?)?)
    }
}

impl<'a> BuildWorker<'a> {
    fn new(
        heaprel: &PgSearchRelation,
        indexrel: &PgSearchRelation,
        config: WorkerConfig,
        coordination: &'a mut WorkerCoordination,
        partitioning: Option<KdTree>,
    ) -> Self {
        Self {
            config,
            table_scan_desc: None,
            heaprel: Clone::clone(heaprel),
            indexrel: Clone::clone(indexrel),
            coordination,
            partitioning,
        }
    }

    fn do_build(&mut self, worker_number: i32, is_leader: bool) -> anyhow::Result<(f64, usize)> {
        unsafe {
            let index_info = self.indexrel.index_info();
            (*index_info).ii_Concurrent = self.config.concurrent;
            let nlaunched = self.coordination.nlaunched();
            let per_worker_memory_budget =
                gucs::adjust_maintenance_work_mem(nlaunched).get() / nlaunched;
            let target_segment_count =
                plan::adjusted_target_segment_count(&self.heaprel, &self.indexrel);
            let (_, worker_segment_target) =
                chunk_range(target_segment_count, nlaunched, worker_number as usize);

            pgrx::debug1!(
                "build_worker {worker_number}: target_segment_count: {target_segment_count}, nlaunched: {nlaunched}, worker_segment_target: {worker_segment_target}"
            );
            if let Some(partitioning) = &self.partitioning {
                pgrx::debug1!("build_worker {worker_number}: partition boundaries: {partitioning}");
            }

            let mut build_state = WorkerBuildState::new(
                &self.heaprel,
                &self.indexrel,
                NonZeroUsize::new(per_worker_memory_budget)
                    .expect("per worker memory budget should be non-zero"),
                self.config.current_xid,
                self.config.next_xid,
                worker_segment_target.max(1),
                target_segment_count,
                self.coordination,
                worker_number,
                is_leader,
            )?;

            set_ps_display_suffix(INDEXING.as_ptr());
            let reltuples = pg_sys::table_index_build_scan(
                self.heaprel.as_ptr(),
                self.indexrel.as_ptr(),
                index_info,
                true,
                true,
                Some(build_callback),
                addr_of_mut!(build_state).cast(),
                self.table_scan_desc
                    .as_ref()
                    .map(|x| x.as_ptr())
                    .unwrap_or(std::ptr::null_mut()),
            );

            build_state.commit()?;
            Ok((reltuples as f64, build_state.nmerges))
        }
    }
}

/// Internal state used by each parallel build worker
struct WorkerBuildState<'a> {
    writer: Option<SerialIndexWriter>,
    categorized_fields: Vec<(SearchField, CategorizedFieldData)>,
    per_row_context: PgMemoryContexts,
    current_xid: pg_sys::FullTransactionId,
    next_xid: pg_sys::FullTransactionId,
    indexrel: PgSearchRelation,
    heaprel: PgSearchRelation,
    index_created_by_version: Option<Version>,
    // the following statistics are used to determine when and what to merge:
    //
    // 1. how many segments does this worker expect to make, assuming no merges?
    estimated_nsegments: OnceLock<usize>,
    //
    // 2. how many segments is this worker supposed to make? (assigned by the leader)
    worker_segment_target: usize,
    //
    // 2b. how many segments are all workers together supposed to make?
    target_segment_count: usize,
    //
    // 3. how many merges has this worker done so far? (incrementing counter)
    nmerges: usize,
    //
    // 4. how many workers are there in total? (including the leader) - utilizing the `nlaunched` field in WorkerCoordination
    coordination: &'a mut WorkerCoordination, // passing in `WorkerCoordination` to have a shared view of `ntuples_done`, used for reporting `tuples_done` in progress monitoring view
    //
    // 5. unmerged segment metas that this worker has created so far
    unmerged_metas: Vec<SegmentMeta>,
    local_tuple_done_count: usize, // worker-local number of tuples done - used to updated the shared `ntuples_done` in `coordination`
    is_leader: bool,
    // whether the first flushed segment's on-disk size has been reported to the disk guard
    recorded_segment_bytes: bool,
}

impl<'a> WorkerBuildState<'a> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        heaprel: &PgSearchRelation,
        indexrel: &PgSearchRelation,
        per_worker_memory_budget: NonZeroUsize,
        current_xid: pg_sys::FullTransactionId,
        next_xid: pg_sys::FullTransactionId,
        worker_segment_target: usize,
        target_segment_count: usize,
        coordination: &'a mut WorkerCoordination,
        worker_number: i32,
        is_leader: bool,
    ) -> anyhow::Result<Self> {
        // If we're making more than one segment, do an early cutoff based on doc
        // count in case the memory budget is so high that all the docs fit into one
        // segment. For a single-segment target, leave it unbounded (memory-driven).
        // Any vector-specific doc cap is applied in `SerialIndexWriter::open`.
        let max_docs_per_segment = if worker_segment_target > 1 {
            Some(
                plan::estimate_heap_reltuples(heaprel) as u32
                    / coordination.nlaunched as u32
                    / worker_segment_target as u32,
            )
        } else {
            None
        };
        let config = IndexWriterConfig {
            memory_budget: per_worker_memory_budget,
            max_docs_per_segment,
        };
        // Abort the build early if it is projected not to fit on the tablespace volume.
        let disk_guard = indexrel
            .is_create_index()
            .then(|| DiskSpaceGuard::new(indexrel));
        let writer =
            SerialIndexWriter::open(indexrel, config, worker_number)?.with_disk_guard(disk_guard);
        let schema = writer.schema();
        let categorized_fields = schema.categorized_fields().clone();
        let created_by_version = indexrel.created_by_version();
        Ok(Self {
            writer: Some(writer),
            categorized_fields,
            per_row_context: PgMemoryContexts::new("pg_search ambuild context"),
            indexrel: indexrel.clone(),
            heaprel: heaprel.clone(),
            index_created_by_version: created_by_version,
            current_xid,
            next_xid,
            worker_segment_target,
            target_segment_count,
            coordination,
            estimated_nsegments: OnceLock::new(),
            nmerges: Default::default(),
            unmerged_metas: Default::default(),
            local_tuple_done_count: 0,
            is_leader,
            recorded_segment_bytes: false,
        })
    }

    /// After this worker flushes a segment, refresh its disk guard with a build-wide view: bump
    /// the shared count of segments written across all workers, and tell the guard how many
    /// segments the whole build still has to write. On the first flush, also hand the guard the
    /// segment's on-disk size (segments are memory-budget bound, so one sample is representative).
    fn on_segment_flushed(&mut self, segment_id: SegmentId) {
        let written = self.coordination.add_segments_written();
        let remaining = self.target_segment_count.saturating_sub(written);

        let mut segment_bytes = None;
        if !self.recorded_segment_bytes {
            unsafe {
                MetaPage::open(&self.indexrel)
                    .segment_metas()
                    .for_each(|_, entry| {
                        if entry.segment_id() == segment_id {
                            segment_bytes = Some(entry.byte_size());
                        }
                    });
            }
            self.recorded_segment_bytes = segment_bytes.is_some();
        }

        let Some(writer) = self.writer.as_mut() else {
            return;
        };
        writer.set_remaining_segments(remaining);
        if let Some(bytes) = segment_bytes {
            writer.set_segment_byte_size(bytes);
        }
    }

    fn commit(&mut self) -> anyhow::Result<()> {
        unsafe {
            set_ps_display_suffix(FINALIZING.as_ptr());
        }
        let writer = self.writer.take().expect("writer should be set");
        if let Some((segment_meta, _)) = writer.commit()? {
            self.unmerged_metas.push(segment_meta);
        }
        self.try_merge(true)?;

        unsafe { set_ps_display_remove_suffix() };
        Ok(())
    }

    /// Based on our calculated chunk size, merge down a chunk of segments into a single segment
    /// if we have created at least that many segments.
    fn try_merge(&mut self, is_last_merge: bool) -> anyhow::Result<()> {
        // which segments should me merge together? if there's not enough, return early
        let segment_ids_to_merge = {
            if self.unmerged_metas.is_empty() {
                return Ok(());
            }

            let chunk_size = if !is_last_merge {
                // calculate the chunk size for this merge iteration
                //
                // chunk_range gives us chunks with the larger ones at the front
                // we want the larger ones at the back, because the smallest "straggler" segment will be written last
                //
                // for instance, imagine we have 3 segments of size [100, 100, 5]
                // we would want the chunks to be [1,2] (merging together [100, 5]) and not [2,1] (merging together [100, 100])
                let (_, chunk_size) = chunk_range(
                    self.estimated_nsegments(self.unmerged_metas[0].max_doc()),
                    self.worker_segment_target,
                    // this achieves the effect of reversing the chunks
                    self.worker_segment_target.saturating_sub(self.nmerges),
                );

                if chunk_size <= 1 || self.unmerged_metas.len() < chunk_size {
                    pgrx::debug1!(
                        "try_merge: skipping merge because chunk_size: {chunk_size}, unmerged_metas: {:?}",
                        self.unmerged_metas.len()
                    );
                    return Ok(());
                }

                if self.nmerges == self.worker_segment_target - 1 {
                    pgrx::debug1!("try_merge: skipping merge because this is not the last merge, and we can only do one more");
                    return Ok(());
                }

                chunk_size
            } else {
                // if it's the last merge, ignore the chunk size and instead solve the following equation for chunk_size:
                //
                // worker_segment_target = self.nmerges + unmerged_segments_len - chunk_size + 1
                //
                // this guarantees we hit the target segment count exactly, assuming we haven't already exceeded it
                // convert to i32 because it's possible that chunk_size comes out negative
                let chunk_size: i32 = self.nmerges as i32 - self.worker_segment_target as i32
                    + self.unmerged_metas.len() as i32
                    + 1;
                if chunk_size <= 1 || self.unmerged_metas.len() < chunk_size as usize {
                    return Ok(());
                }

                chunk_size as usize
            };

            self.unmerged_metas.sort_by_key(|entry| entry.max_doc());
            self.unmerged_metas
                .drain(..chunk_size)
                .map(|entry| entry.id())
                .collect::<Vec<_>>()
        };

        // do the merge
        pgrx::debug1!(
            "do_merge: last merge {}, about to merge {} segments: {:?}",
            is_last_merge,
            segment_ids_to_merge.len(),
            segment_ids_to_merge
        );
        let mut merger = SearchIndexMerger::open(&self.indexrel, MvccSatisfies::Mergeable)?;
        unsafe { set_ps_display_suffix(MERGING.as_ptr()) };
        merger.merge_segments(&segment_ids_to_merge)?;

        // garbage collect the index, returning to the fsm
        pgrx::debug1!("do_merge: garbage collecting");
        unsafe {
            set_ps_display_suffix(GARBAGE_COLLECTING.as_ptr());
        };
        unsafe { garbage_collect_index(&self.indexrel, self.current_xid, self.next_xid) };

        self.nmerges += 1;

        Ok(())
    }

    /// Estimates how many segments this worker will make if no merging happens.
    ///
    /// This is used to determine how many segments to merge down in chunks.
    fn estimated_nsegments(&self, docs_per_segment: u32) -> usize {
        *self.estimated_nsegments.get_or_init(|| {
            let reltuples = plan::estimate_heap_reltuples(&self.heaprel);
            let reltuples_per_worker = reltuples / self.coordination.nlaunched as f64;
            let nsegments = (reltuples_per_worker / docs_per_segment as f64).ceil() as usize;
            pgrx::debug1!("estimated that this worker will make {nsegments} segments, based on reltuples: {reltuples}, nlaunched: {}, reltuples_per_worker: {reltuples_per_worker}, docs_per_segment: {docs_per_segment}", self.coordination.nlaunched);
            nsegments
        })
    }
}

#[pg_guard]
unsafe extern "C-unwind" fn build_callback(
    _indexrel: pg_sys::Relation,
    ctid: pg_sys::ItemPointer,
    values: *mut pg_sys::Datum,
    isnull: *mut bool,
    _tuple_is_alive: bool,
    state: *mut std::os::raw::c_void,
) {
    check_for_interrupts!();

    let build_state = &mut *state.cast::<WorkerBuildState>();
    let ctid_u64 = crate::postgres::utils::item_pointer_to_u64(*ctid);

    let segment_meta = build_state.per_row_context.switch_to(|_| {
        let mut doc = TantivyDocument::new();

        // Unpack all composites upfront
        let unpacked_composites =
            CompositeSlotValues::from_composites(collect_composites_for_unpacking(
                build_state.categorized_fields.iter().map(|(_, cat)| cat),
                values,
                isnull,
            ));

        row_to_search_document(
            build_state
                .categorized_fields
                .iter()
                .map(|(field, categorized)| {
                    let (datum, is_null) = get_field_value(
                        &categorized.source,
                        categorized.attno,
                        values,
                        isnull,
                        &unpacked_composites,
                    );
                    (datum, is_null, field, categorized)
                }),
            &mut doc,
            build_state.index_created_by_version,
        )
        .unwrap_or_else(|e| panic!("{e}"));

        build_state
            .writer
            .as_mut()
            .expect("build_callback: writer should be set")
            .insert(doc, ctid_u64, || set_ps_display_suffix(COMMITTING.as_ptr()))
            .unwrap_or_else(|e| panic!("{e}"))
    });
    build_state.per_row_context.reset();

    build_state.local_tuple_done_count += 1;

    if build_state.local_tuple_done_count % TUPLES_DONE_BATCH_SIZE == 0 {
        build_state
            .coordination
            .add_tuples_done(TUPLES_DONE_BATCH_SIZE);

        if build_state.is_leader {
            pg_sys::pgstat_progress_update_param(
                pg_sys::PROGRESS_CREATEIDX_TUPLES_DONE as i32,
                build_state.coordination.tuples_done() as i64,
            );
        }
    }

    if let Some(segment_meta) = segment_meta {
        build_state.on_segment_flushed(segment_meta.id());
        build_state.unmerged_metas.push(segment_meta);
        build_state
            .try_merge(false)
            .unwrap_or_else(|e| panic!("{e}"));
        set_ps_display_suffix(INDEXING.as_ptr());
    }
}

/// Build an index.  This is the workhorse behind `CREATE INDEX` and `REINDEX`.
///
/// If the system allows, it will build the index in parallel.  Otherwise the index is built in
/// serially in this connected backend.
pub(super) fn build_index(
    heaprel: PgSearchRelation,
    indexrel: PgSearchRelation,
    concurrent: bool,
) -> anyhow::Result<f64> {
    struct SnapshotDropper(pg_sys::Snapshot);
    crate::impl_safe_drop!(SnapshotDropper, |self| {
        unsafe {
            let snapshot = self.0;
            // if it's an mvcc snapshot we must unregister it
            if (*snapshot).snapshot_type == pg_sys::SnapshotType::SNAPSHOT_MVCC
                || (*snapshot).snapshot_type == pg_sys::SnapshotType::SNAPSHOT_HISTORIC_MVCC
            {
                pg_sys::UnregisterSnapshot(snapshot);
            }
        }
    });

    let snapshot = SnapshotDropper(unsafe {
        if concurrent {
            pg_sys::RegisterSnapshot(pg_sys::GetTransactionSnapshot())
        } else {
            &raw mut pg_sys::SnapshotAnyData
        }
    });

    let config = WorkerConfig {
        heaprelid: heaprel.oid(),
        indexrelid: indexrel.oid(),
        concurrent,
        current_xid: unsafe { pg_sys::GetCurrentFullTransactionId() },
        need_wal: indexrel.need_wal(),
        next_xid: unsafe { pg_sys::ReadNextFullTransactionId() },
    };

    // Boundaries are fixed before any worker starts, so every worker cuts on the same ones.
    // TODO(M3): the target segment count doubles as the partition count for now; rename the
    // reloption to `partition_count` once partitioned storage lands.
    let partitioning = plan_partition_boundaries(
        &heaprel,
        &indexrel,
        snapshot.0,
        plan::adjusted_target_segment_count(&heaprel, &indexrel),
    )?;
    if let Some(partitioning) = &partitioning {
        pgrx::debug1!(
            "build_index: {} partition boundaries:\n{}",
            partitioning.partition_count(),
            partitioning.bounds_listing()
        );
    }
    let partitioning_bytes = match &partitioning {
        Some(partitioning) => postcard::to_allocvec(partitioning)?,
        None => Vec::new(),
    };

    let process = ParallelBuild::new(&heaprel, snapshot.0, config, partitioning_bytes);
    let nworkers = plan::create_index_nworkers(&heaprel, &indexrel);
    pgrx::debug1!("build_index: asked for {nworkers} workers");

    // This is updating `tuples_total` in the `pg_stat_progress_create_index` view - `tuples_done` is incremented in `build_callback`
    unsafe {
        pg_sys::pgstat_progress_update_param(
            pg_sys::PROGRESS_CREATEIDX_TUPLES_TOTAL as i32,
            plan::estimate_heap_reltuples(&heaprel) as i64,
        );
    }

    let total_tuples = if let Some(mut process) = launch_parallel_process!(
        ParallelBuild<BuildWorker>,
        process,
        WorkerStyle::Maintenance,
        nworkers,
        1024
    ) {
        let nlaunched = process.launched_workers();
        pgrx::debug1!("build_index: launched {nworkers} workers (not including leader)");
        let coordination = process
            .state_manager_mut()
            .object::<WorkerCoordination>(2)
            .expect("process coordination")
            .expect("process coordination should not be NULL");

        // account for the leader in the coordination
        let mut nlaunched_plus_leader = nlaunched;
        let leader_participating = unsafe { pg_sys::parallel_leader_participation };
        if leader_participating {
            nlaunched_plus_leader += 1;
        }

        // set_nlaunched last, because workers wait for this to be set
        coordination.set_nlaunched(nlaunched_plus_leader);
        pgrx::debug1!("build_index: has {nlaunched_plus_leader} workers (including leader)");

        let (mut total_tuples, mut total_merges) = if leader_participating {
            // if the leader is to participate too, it's nice for it to wait until all the other workers
            // have indicated that they're running.  Otherwise, it's likely the leader will get ahead
            // of the workers, which doesn't allow for "evenly" distributing the work
            while coordination.nstarted() != nlaunched {
                check_for_interrupts!();
                std::thread::yield_now();
            }

            // directly instantiate a worker for the leader and have it do its build
            let mut worker = BuildWorker::new_parallel_worker(*process.state_manager());
            worker.do_build(nlaunched_plus_leader as i32, true)?
        } else {
            pgrx::debug1!("build_index: leader is not participating");
            (0.0, 0)
        };

        // wait for the workers to finish by collecting all their response messages
        for (_, message) in process {
            check_for_interrupts!();
            let worker_response = serde_json::from_slice::<WorkerResponse>(&message)?;
            total_tuples += worker_response.reltuples;
            total_merges += worker_response.nmerges;
        }

        pgrx::debug1!("build_index: total_tuples: {total_tuples}, total_merges: {total_merges}");
        total_tuples
    } else {
        pgrx::debug1!("build_index: not doing a parallel build");
        // not doing a parallel build, so directly instantiate a BuildWorker and serially run the
        // whole build here in this connected backend
        let mut coordination: WorkerCoordination = Default::default();
        coordination.set_nlaunched(1);

        let mut worker =
            BuildWorker::new(&heaprel, &indexrel, config, &mut coordination, partitioning);

        let (total_tuples, total_merges) = worker.do_build(1, true)?;
        pgrx::debug1!("build_index: total_tuples: {total_tuples}, total_merges: {total_merges}");
        total_tuples
    };

    unsafe { set_ps_display_remove_suffix() };
    Ok(total_tuples)
}

mod plan {
    use super::*;

    pub(super) const MAX_VECTOR_BUILD_WORKERS: usize = 4;

    /// Determine the number of workers to use for a given CREATE INDEX/REINDEX statement.
    ///
    /// The number of workers is determined by max_parallel_maintenance_workers. However, if max_parallel_maintenance_workers
    /// is greater than available parallelism, we use available parallelism.
    ///
    /// If the leader is participating, we subtract 1 from the number of workers because the leader also counts as a worker.
    pub(super) fn create_index_nworkers(
        heaprel: &PgSearchRelation,
        indexrel: &PgSearchRelation,
    ) -> usize {
        // We don't want a parallel build to happen if we're creating a single segment
        let target_segment_count = plan::adjusted_target_segment_count(heaprel, indexrel);
        if target_segment_count == 1 {
            return 0;
        }

        // NB: we _could_ use pg_sys::plan_create_index_workers(), or on v17+ accept IndexIndex::ii_ParallelWorkers,
        // but doing either of these would prohibit the user from having direct control over the number of
        // workers used for a given CREATE INDEX/REINDEX statement.  Internal discussions led to that
        // being more important that us trying to be "smart"
        let maintenance_workers = unsafe {
            if !heaprel.rd_options.is_null() {
                let options = heaprel.rd_options.cast::<pg_sys::StdRdOptions>();
                if (*options).parallel_workers <= 0 {
                    pg_sys::max_parallel_maintenance_workers as usize
                } else {
                    (*options).parallel_workers as usize
                }
            } else {
                pg_sys::max_parallel_maintenance_workers as usize
            }
        };

        // must also be less than max_parallel_workers and max_worker_processes
        let maintenance_workers = maintenance_workers
            .min(unsafe { pg_sys::max_parallel_workers as usize })
            .min(unsafe { pg_sys::max_worker_processes as usize });

        // For vector builds, beyond 4 workers there's no improvement to index build time:
        // clustering dominates and parallelizes internally (rayon + BLAS threads), so
        // parallelizing the heap scan/segment flushing isn't the bottleneck. Constrain to 4
        // because fewer concurrent merges means less memory.
        let maintenance_workers = if indexrel
            .schema()
            .is_ok_and(|schema| schema.has_vector_field())
        {
            maintenance_workers.min(MAX_VECTOR_BUILD_WORKERS)
        } else {
            maintenance_workers
        };

        if maintenance_workers < 3 {
            ErrorReport::new(
                PgSqlErrorCode::ERRCODE_INSUFFICIENT_RESOURCES,
                format!("only {maintenance_workers} parallel workers were available for index build"),
                function_name!(),
            )
            .set_detail("for large tables, increasing the number of workers can reduce the time it takes to build the index")
            .set_hint("`SET max_parallel_maintenance_workers = <number>`")
            .report(PgLogLevel::WARNING);
        }

        if maintenance_workers == 0 {
            return 0;
        }

        // Ensure that we never have more workers (including the leader) than the max allowed number of workers.
        //
        // We also want nworkers to be at most 1/2 of the target segment count. To illustrate why:
        //
        // Imagine we have 8 workers, a target segment count of 8, and a table size such that each worker produces 4 segments.
        // In this scenario, each worker would do one big merge of all 4 segments at the very end, which means none of the
        // merges would be able to reuse the FSM.
        //
        // On the other hand, imagine we have only 4 workers, over the same table and target segment count.
        // In this scenario, each worker would target 2 segments, meaning it would do 2 merges -- once when it's about halfway done
        // and once at the end. The merge at the end would be able to use the free space created by the first merge.
        let max_workers = target_segment_count.div_ceil(2);
        let mut nworkers = maintenance_workers.min(max_workers);

        if unsafe { pg_sys::parallel_leader_participation } && nworkers == max_workers {
            nworkers -= 1;
        }

        nworkers
    }

    /// If we determine that the table is very small, we should just create a single segment
    pub(super) fn adjusted_target_segment_count(
        heaprel: &PgSearchRelation,
        indexrel: &PgSearchRelation,
    ) -> usize {
        // If there are fewer rows than number of CPUs, use 1 worker
        let reltuples = plan::estimate_heap_reltuples(heaprel);
        let target_segment_count = indexrel.options().target_segment_count();
        if reltuples <= target_segment_count as f64 {
            pgrx::debug1!("number of reltuples ({reltuples}) is less than target segment count ({target_segment_count}), creating a single segment");
            return 1;
        }

        // If the entire heap fits inside the smallest allowed Tantivy segment memory budget of 15MB, use 1 worker
        let byte_size = plan::estimate_heap_byte_size(heaprel);
        if byte_size <= 15 * 1024 * 1024 {
            pgrx::debug1!(
                "heap byte size ({byte_size}) is less than 15MB, creating a single segment"
            );
            return 1;
        }

        target_segment_count
    }

    // TODO: Convert to use RowEstimate.
    pub(super) fn estimate_heap_reltuples(heap_relation: &PgSearchRelation) -> f64 {
        let mut reltuples = unsafe { (*heap_relation.rd_rel).reltuples };

        // if the reltuples estimate is not available, estimate the number of tuples in the heap
        // by multiplying the number of pages by the max offset number of the first page
        if reltuples <= 0.0 {
            let npages = unsafe {
                pg_sys::RelationGetNumberOfBlocksInFork(
                    heap_relation.as_ptr(),
                    heap_relation.fork_number(),
                )
            };

            if npages == 0 {
                // the tuple count actually is 0
                return 0.0;
            }

            let bman = BufferManager::new(heap_relation);
            let buffer = bman.get_buffer(0);
            let page = buffer.page();
            let max_offset = page.max_offset_number();
            reltuples = npages as f32 * max_offset as f32;
        }

        reltuples as f64
    }

    pub(super) fn estimate_heap_byte_size(heap_relation: &PgSearchRelation) -> usize {
        let npages = unsafe {
            pg_sys::RelationGetNumberOfBlocksInFork(
                heap_relation.as_ptr(),
                heap_relation.fork_number(),
            )
        };

        npages as usize * pg_sys::BLCKSZ as usize
    }
}

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

    fn setup_parallel_build_large_table() {
        Spi::run(
            r#"
            DROP TABLE IF EXISTS parallel_build_large;
            CREATE TABLE parallel_build_large (
                id SERIAL PRIMARY KEY,
                name TEXT
            );
            INSERT INTO parallel_build_large (name)
            SELECT 'lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.'
            FROM generate_series(1, 35000);
            "#,
        )
        .expect("failed to setup parallel_build_large table");
    }

    fn cleanup_parallel_build_large_table() {
        Spi::run("DROP TABLE IF EXISTS parallel_build_large;")
            .expect("failed to cleanup parallel_build_large table");
    }

    /// Tests that parallel index building fails with insufficient memory.
    #[pg_test]
    #[should_panic(expected = "maintenance_work_mem")]
    fn test_parallel_build_large_insufficient_memory() {
        setup_parallel_build_large_table();

        Spi::run("SET max_parallel_workers = 8;").unwrap();
        Spi::run("SET maintenance_work_mem = '64MB';").unwrap();
        Spi::run("SET max_parallel_maintenance_workers = 8;").unwrap();

        // This should panic with a "maintenance_work_mem is not high enough" error
        Spi::run(
            "CREATE INDEX parallel_build_large_idx ON parallel_build_large USING bm25 (id, name) WITH (key_field = 'id', target_segment_count = 16);",
        ).unwrap();
    }

    /// Tests parallel index building with various configurations.
    ///
    /// This test creates 35,000 rows and tests 32 different configuration
    /// combinations of maintenance_work_mem, workers, leader participation,
    /// and target segment count.
    #[pg_test]
    fn test_parallel_build_large_configurations() {
        setup_parallel_build_large_table();

        Spi::run("SET max_parallel_workers = 8;").unwrap();

        let maintenance_work_mem = ["2GB", "128MB"];
        let maintenance_workers = [6, 2];
        let leader_participation = [true, false];
        let target_segments = [4, 32];

        for mwm in &maintenance_work_mem {
            for mw in &maintenance_workers {
                for lp in &leader_participation {
                    for ts in &target_segments {
                        // Set configuration
                        Spi::run(&format!("SET max_parallel_maintenance_workers = {};", mw))
                            .unwrap();
                        Spi::run(&format!("SET parallel_leader_participation = {};", lp)).unwrap();
                        Spi::run(&format!("SET maintenance_work_mem = '{}';", mwm)).unwrap();

                        // Create index
                        Spi::run(&format!(
                            "CREATE INDEX parallel_build_large_idx ON parallel_build_large USING bm25 (id, name) WITH (key_field = 'id', target_segment_count = {});",
                            ts
                        ))
                        .unwrap_or_else(|e| {
                            panic!(
                                "CREATE INDEX failed with workers={}, leader={}, mem={}, segments={}: {:?}",
                                mw, lp, mwm, ts, e
                            )
                        });

                        // Verify segment count
                        let count: i64 = Spi::get_one(
                            "SELECT COUNT(*)::bigint FROM paradedb.index_info('parallel_build_large_idx');",
                        )
                        .unwrap()
                        .unwrap();

                        if *ts == 4 {
                            assert_eq!(
                                count, 4,
                                "Expected 4 segments with workers={}, leader={}, mem={}, segments={}, got {}",
                                mw, lp, mwm, ts, count
                            );
                        } else if *ts == 32 {
                            assert!(
                                (28..=33).contains(&count),
                                "Expected 28-33 segments with workers={}, leader={}, mem={}, segments={}, got {}",
                                mw, lp, mwm, ts, count
                            );
                        }

                        // Verify total document count
                        let num_docs: i64 = Spi::get_one(
                            "SELECT COALESCE(SUM(num_docs), 0)::bigint FROM paradedb.index_info('parallel_build_large_idx');",
                        )
                        .unwrap()
                        .unwrap();

                        assert_eq!(
                            num_docs, 35000,
                            "Expected 35000 docs with workers={}, leader={}, mem={}, segments={}, got {}",
                            mw, lp, mwm, ts, num_docs
                        );

                        // Drop index for next iteration
                        Spi::run("DROP INDEX parallel_build_large_idx;").unwrap();
                    }
                }
            }
        }

        cleanup_parallel_build_large_table();
    }

    /// A row count that is an exact multiple of the vector doc-count cap leaves the writer with
    /// no pending segment at commit; the final merge must still run.
    #[pg_test]
    fn test_single_segment_build_exact_doc_cap_multiple() {
        let nrows = crate::index::writer::index::DEFAULT_MAX_DOCS_PER_SEGMENT * 5;
        Spi::run("CREATE EXTENSION IF NOT EXISTS vector;").unwrap();
        Spi::run(&format!(
            r#"
            CREATE TABLE exact_cap_multiple (id SERIAL8, emb vector(8));
            INSERT INTO exact_cap_multiple (emb)
            SELECT ('[' || array_to_string(array(SELECT random() FROM generate_series(1, 8)), ',') || ']')::vector
            FROM generate_series(1, {nrows});
            "#
        ))
        .unwrap();
        Spi::run(
            "CREATE INDEX exact_cap_multiple_idx ON exact_cap_multiple USING paradedb (id, emb vector_cosine_ops) WITH (key_field = 'id', target_segment_count = 1);",
        )
        .unwrap();

        let count: i64 = Spi::get_one(
            "SELECT COUNT(*)::bigint FROM paradedb.index_info('exact_cap_multiple_idx');",
        )
        .unwrap()
        .unwrap();
        assert_eq!(count, 1, "expected a single segment, got {count}");
    }
}
