// 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::Cardinality;

use crate::api::HashSet;

use crate::postgres::ParallelScanState;

pub use crate::scan::info::RowEstimate;

use pgrx::pg_sys;

use tantivy::index::SegmentId;

fn clamp_to_gather_limits(nworkers: usize) -> usize {
    unsafe {
        nworkers
            .min(pg_sys::max_parallel_workers_per_gather as usize)
            .min(pg_sys::max_parallel_workers as usize)
    }
}

/// Compute the number of workers that should be used for the given ExecMethod.
///
/// This calculation determines the "Parallel Awareness" of the path:
/// - If it returns `0`, the path is marked as `parallel_safe` but NOT `parallel_aware`.
///   PostgreSQL may run this scan in a worker (e.g. inner side of a join), but it will
///   be a "replicated" scan where every worker processes the full data set.
/// - If it returns `> 0`, the path is marked as BOTH `parallel_safe` and `parallel_aware`.
///   It becomes a "partial" path that coordinates with other workers via DSM to
///   partition segments and avoid duplicate work.
///
/// Note: PostgreSQL asserts that `parallel_aware` paths must have `parallel_workers > 0`.
pub fn compute_nworkers(
    declares_sorted_output: bool,
    limit: Option<Cardinality>,
    estimated_total_rows: RowEstimate,
    segment_count: usize,
    has_external_quals: bool,
    has_correlated_param: bool,
    is_join_context: bool,
) -> usize {
    // Start with segment-based parallelism. The leader is not included in `nworkers`,
    // so exclude it here. For example: if we expect to need to query 1 segment, then
    // we don't need any workers.
    let mut nworkers = segment_count.saturating_sub(1);

    // For scans with reliable row estimates (RowEstimate::Known) we cap workers two
    // ways; an Unknown estimate (table not ANALYZEd) caps nothing, since we can't
    // trust it:
    //
    // 1. Limit-based (UNSORTED only): cap to the segments needed to reach LIMIT.
    //    Sorted output must scan every segment to produce a correct global order, so
    //    it is exempt (#4457).
    // 2. Row-based: cap so each worker processes at least `min_rows_per_worker` rows
    //    (~300K default), for both sorted and unsorted output. Skipped in join
    //    contexts to avoid preventing Parallel Hash Join.
    //
    // See: https://github.com/paradedb/paradedb/issues/3055
    if let RowEstimate::Known(total_rows) = estimated_total_rows {
        // Cap to the number of segments needed to reach the LIMIT. Unsorted only:
        // sorted output needs every segment, so it is exempt (#4457).
        if let (false, Some(limit)) = (declares_sorted_output, limit) {
            let rows_per_segment = total_rows as f64 / segment_count.max(1) as f64;
            let segments_to_reach_limit = (limit / rows_per_segment).ceil() as usize;
            // The leader is not included in `nworkers`, so subtract 1.
            let nworkers_for_limited_segments = segments_to_reach_limit.saturating_sub(1);
            nworkers = nworkers.min(nworkers_for_limited_segments);
        }

        // Cap so each worker processes at least min_rows_per_worker rows.
        // Skipped for joins: failing to claim workers can prevent the planner from
        // choosing Parallel Hash Join, leading to inefficient serial plans.
        if !is_join_context {
            let min_rows_per_worker = crate::gucs::min_rows_per_worker() as u64;
            #[allow(clippy::manual_checked_ops)]
            if min_rows_per_worker > 0 {
                let max_workers_for_rows = (total_rows / min_rows_per_worker) as usize;
                nworkers = nworkers.min(max_workers_for_rows);
            }
        }
    }

    nworkers = clamp_to_gather_limits(nworkers);

    if has_external_quals {
        // Don't attempt to parallelize if we depend on external variables (e.g. inner side of a nested loop join).
        // This occurs when a qual contains a Param that references a value from another relation
        // (e.g. t1.val @@@ t2.val). In this case, we are likely executing a parameterized scan
        // where we are re-executed for every row of the outer relation. Parallelism here is
        // complex and often not desired.
        //
        // This is distinct from `is_join_context`, which indicates we are part of a join query
        // (e.g. Hash Join) but our scan keys are independent.
        //
        // TODO: Re-evaluate.
        nworkers = 0;
    }

    if has_correlated_param {
        // Don't attempt to parallelize when we have correlated PARAM_EXEC nodes. Uncorrelated
        // params are solved during BeginCustomScan and pushed down to parallel workers, but
        // correlated params need to be evaluated during the scan itself.
        // TODO: Implement proper correlated PARAM_EXEC param sharing with parallel workers.
        nworkers = 0;
    }

    #[cfg(not(feature = "pg15"))]
    unsafe {
        if nworkers == 0 && pg_sys::debug_parallel_query != 0 {
            // force a parallel worker if the `debug_parallel_query` GUC is on
            nworkers = 1;
        }
    }

    nworkers
}

/// Maximum number of useful parallel workers given structural constraints only.
///
/// Unlike `compute_nworkers`, this does NOT gate on row count or the
/// `min_rows_per_worker` GUC. The caller uses the returned upper bound in its
/// serial-vs-parallel cost comparison before emitting one chosen path (see
/// #4664).
///
/// Returns 0 if parallelism is structurally impossible:
/// - External quals (parameterized scan in a nested loop).
/// - Correlated PARAM_EXEC nodes.
///
/// Otherwise returns `min(segment_count - 1, max_parallel_workers_per_gather,
/// max_parallel_workers)`. The leader is excluded from `segment_count - 1`.
pub fn max_useful_workers(
    segment_count: usize,
    has_external_quals: bool,
    has_correlated_param: bool,
) -> usize {
    if has_external_quals || has_correlated_param {
        return 0;
    }

    // Only the pg15+ `debug_parallel_query` block below mutates this; on pg15
    // that block is compiled out, leaving the binding immutable.
    #[cfg_attr(feature = "pg15", allow(unused_mut))]
    let mut nworkers = clamp_to_gather_limits(segment_count.saturating_sub(1));

    #[cfg(not(feature = "pg15"))]
    unsafe {
        if nworkers == 0 && pg_sys::debug_parallel_query != 0 {
            nworkers = 1;
        }
    }

    nworkers
}

pub unsafe fn checkout_segment_for_source(
    pscan_state: *mut ParallelScanState,
    source_idx: usize,
) -> Option<SegmentId> {
    (*pscan_state).checkout_segment_for_source(source_idx)
}

pub unsafe fn list_segment_ids(pscan_state: *mut ParallelScanState) -> HashSet<SegmentId> {
    (*pscan_state).segments().into_keys().collect()
}
