// 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 serde::{Deserialize, Serialize};

use crate::api::FieldName;
use crate::postgres::pdb_owned_value::PdbOwnedValue;
use crate::query::pdb_query::pdb::Query;
use crate::query::SearchQueryInput;

/// Defines the logical boundary split points for scanning the index. When provided,
/// the DataFusion execution plan uses these points to statically partition the scan
/// into sequential chunks, rather than relying on dynamic segment checkout.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RangePartitioning {
    /// The index field used to define the boundaries.
    pub partition_by: FieldName,
    /// The values that split the data space into separate partitions. A length of N
    /// produces N+1 partitions.
    pub split_points: Vec<PdbOwnedValue>,
}

impl RangePartitioning {
    /// Returns the static boundary constraint for the given partition as a single RangeQuery.
    ///
    /// **Consumer Caveats**:
    /// - A row whose partition field is NULL will be deterministically routed to partition 0.
    /// - A multi-valued field can fall into multiple partition ranges and duplicate the row. This is statically prevented during index configuration for `partition_by` columns.
    pub fn partition_bounds(&self, partition: usize) -> SearchQueryInput {
        let lower = if partition > 0 {
            let val = &self.split_points[partition - 1];
            if matches!(val, PdbOwnedValue::Null) {
                std::ops::Bound::Unbounded
            } else {
                std::ops::Bound::Included(val.clone())
            }
        } else {
            std::ops::Bound::Unbounded
        };

        let mut is_empty_range = false;
        let upper = if partition < self.split_points.len() {
            let val = &self.split_points[partition];
            if matches!(val, PdbOwnedValue::Null) {
                is_empty_range = true;
                std::ops::Bound::Unbounded
            } else {
                std::ops::Bound::Excluded(val.clone())
            }
        } else {
            std::ops::Bound::Unbounded
        };

        let range_query = if is_empty_range {
            SearchQueryInput::Empty
        } else {
            SearchQueryInput::FieldedQuery {
                field: self.partition_by.clone(),
                query: Query::Range {
                    lower_bound: lower,
                    upper_bound: upper,
                },
            }
        };

        // NULLs fall into partition 0
        if partition == 0 {
            let null_query = SearchQueryInput::Boolean {
                must: vec![],
                should: vec![],
                must_not: vec![SearchQueryInput::FieldedQuery {
                    field: self.partition_by.clone(),
                    query: Query::Exists,
                }],
                minimum_should_match: None,
            };

            if is_empty_range {
                null_query
            } else {
                SearchQueryInput::Boolean {
                    must: vec![],
                    should: vec![range_query, null_query],
                    must_not: vec![],
                    minimum_should_match: None,
                }
            }
        } else {
            range_query
        }
    }
}

/// A representation of the raw partition split points for range partitioning.
/// Rather than directly instantiating a static `RangePartitioning` with these
/// points, this sample is kept so that `TaskEstimator`s and distributed execution
/// engines can ask for exactly their preferred number of partitions at runtime.
///
/// **Precondition**: `sample_points` must be sorted ascending, otherwise the generated
/// boundaries will produce overlapping or gapped ranges that silently drop or duplicate rows.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RangePartitioningSample {
    /// The index field used to define the boundaries.
    pub partition_by: FieldName,
    /// A sample of values from the data space. This sample is typically much larger
    /// than the target number of partitions, allowing us to safely down-sample to compute
    /// relatively uniform distribution boundaries.
    pub sample_points: Vec<PdbOwnedValue>,
}

impl RangePartitioningSample {
    /// Generates a concrete `RangePartitioning` bounding exactly `target_partitions`.
    ///
    /// If `target_partitions` is smaller than the sample's inherent size, the
    /// sample points are evenly down-sampled (effectively merging contiguous partitions).
    /// To avoid scheduling unnecessary tasks scanning empty ranges, `target_partitions`
    /// is capped at `sample_points.len() + 1`.
    pub fn build(&self, target_partitions: usize) -> RangePartitioning {
        debug_assert!(
            self.sample_points.windows(2).all(|w| w[0] <= w[1]),
            "RangePartitioningSample requires sample_points to be sorted ascending"
        );

        let num_samples = self.sample_points.len();

        // Cap target_partitions to avoid scheduling unnecessary empty ranges
        let actual_partitions = if target_partitions > num_samples + 1 {
            num_samples + 1
        } else {
            target_partitions
        };

        if actual_partitions <= 1 {
            return RangePartitioning {
                partition_by: self.partition_by.clone(),
                split_points: vec![],
            };
        }

        let mut new_split_points = Vec::with_capacity(actual_partitions - 1);

        // Down-sample evenly. We want `actual_partitions - 1` split points.
        for i in 1..actual_partitions {
            let split_idx = (i * num_samples) / actual_partitions;
            new_split_points.push(self.sample_points[split_idx].clone());
        }

        RangePartitioning {
            partition_by: self.partition_by.clone(),
            split_points: new_split_points,
        }
    }
}
