// 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::index::reader::index::SearchIndexReader;
use crate::postgres::ParallelScanState;
use pgrx::{pg_guard, pg_sys};
use tantivy::index::SegmentId;

/// Safety margin for DSM allocation: actual segment count can exceed target_segment_count due to
/// concurrent writes between plan and execute.
const PARALLEL_SEGMENT_ESTIMATE_MULTIPLIER: usize = 64;

#[pg_guard]
pub unsafe extern "C-unwind" fn aminitparallelscan(target: *mut ::core::ffi::c_void) {
    let state = target.cast::<ParallelScanState>();
    (*state).create();
}

#[pg_guard]
pub unsafe extern "C-unwind" fn amparallelrescan(scan: pg_sys::IndexScanDesc) {
    // PostgreSQL calls this before a rescan to reset the parallel scan state.
    // Mark as uninitialized so workers wait for leader to re-populate.
    if let Some(state) = get_bm25_scan_state(scan) {
        let _mutex = state.acquire_mutex();
        state.mark_uninitialized();
    }
}

#[cfg(any(feature = "pg15", feature = "pg16"))]
#[pg_guard]
pub unsafe extern "C-unwind" fn amestimateparallelscan() -> pg_sys::Size {
    ParallelScanState::size_of(&[u16::MAX as usize], &[], false)
}

#[cfg(feature = "pg17")]
#[pg_guard]
pub unsafe extern "C-unwind" fn amestimateparallelscan(
    _nkeys: i32,
    _norderbys: i32,
) -> pg_sys::Size {
    // NB:  in this function, we have no idea how many segments we have.  We don't even know which
    // index we're querying.  So we choose a, hopefully, large enough value at 65536, or u16::MAX
    // TODO: This will result in a ~1MB allocation.
    ParallelScanState::size_of(&[u16::MAX as usize], &[], false)
}

#[cfg(feature = "pg18")]
#[pg_guard]
pub unsafe extern "C-unwind" fn amestimateparallelscan(
    rel: *mut pg_sys::RelationData,
    _nkeys: i32,
    _norderbys: i32,
) -> pg_sys::Size {
    // In PG18, we have access to the relation, so we can get a better estimate
    // using target_segment_count() instead of the worst-case u16::MAX
    let nsegments = if rel.is_null() {
        u16::MAX as usize
    } else {
        estimated_parallel_segments(rel)
    };
    ParallelScanState::size_of(&[nsegments], &[], false)
}

unsafe fn bm25_shared_state(
    scan: &mut pg_sys::IndexScanDescData,
) -> Option<&mut ParallelScanState> {
    if scan.parallel_scan.is_null() {
        None
    } else {
        scan.parallel_scan
            .cast::<std::ffi::c_void>()
            .add({
                #[cfg(any(feature = "pg15", feature = "pg16", feature = "pg17"))]
                {
                    (*scan.parallel_scan).ps_offset
                }
                #[cfg(feature = "pg18")]
                {
                    (*scan.parallel_scan).ps_offset_am
                }
            })
            .cast::<ParallelScanState>()
            .as_mut()
    }
}

/// Initialize parallel scan state if not already done.
/// Only the leader calls this function to populate the segment pool with its
/// snapshot-visible segments. Workers wait for this initialization and then
/// use ParallelWorker visibility to see the same segments.
/// Segments are NOT claimed here - they're claimed lazily in amgettuple/amgetbitmap.
pub unsafe fn maybe_init_parallel_scan(
    scan: pg_sys::IndexScanDesc,
    searcher: &SearchIndexReader,
) -> Option<i32> {
    if unsafe { (*scan).parallel_scan.is_null() } {
        // not a parallel scan, so there's nothing to initialize
        return None;
    }

    // Read indexRelation before mutable borrow of scan
    let rel = (*scan).indexRelation;

    let state = get_bm25_scan_state(scan)?;

    let _mutex = state.acquire_mutex();

    if !state.is_initialized() {
        // If concurrent writes created more segments than the DSM can hold, fall back to serial
        // scan to avoid a shared memory overflow.
        let estimated_segments = estimated_parallel_segments(rel);
        let actual_segments = searcher.segment_readers().len();
        if actual_segments > estimated_segments {
            state.mark_initialized_empty();
            pgrx::warning!(
                "Actual segment count ({}) exceeds estimated ({}), falling back to serial scan",
                actual_segments,
                estimated_segments
            );
            return None;
        }
        state.populate(&[searcher.segment_readers()], &[], false);
    }
    Some(unsafe { pg_sys::ParallelWorkerNumber })
}

/// Claim a segment from the shared pool.
/// Both leader and workers use this to get work.
/// All participants wait for initialization before attempting to claim.
pub unsafe fn maybe_claim_segment(scan: pg_sys::IndexScanDesc) -> Option<SegmentId> {
    get_bm25_scan_state(scan)?.checkout_segment_for_source(0)
}

pub fn get_bm25_scan_state<'a>(scan: pg_sys::IndexScanDesc) -> Option<&'a mut ParallelScanState> {
    unsafe {
        assert!(!scan.is_null());
        let scan = scan.as_mut().unwrap_unchecked();
        bm25_shared_state(scan)
    }
}

unsafe fn estimated_parallel_segments(rel: pg_sys::Relation) -> usize {
    crate::postgres::options::BM25IndexOptions::from_relation(rel).target_segment_count()
        * PARALLEL_SEGMENT_ESTIMATE_MULTIPLIER
}
