// 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 std::sync::Arc;

use datafusion::common::config::ConfigOptions;
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::common::Result;
use datafusion::physical_expr::equivalence::EquivalenceProperties;
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::joins::utils::JoinFilter;
use datafusion::physical_plan::joins::{HashJoinExec, SortMergeJoinExec};
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};

use crate::scan::filter_passthrough_exec::FilterPassthroughExec;

/// A DataFusion physical optimizer rule that replaces `HashJoinExec` with `SortMergeJoinExec`
/// when the inputs are already sorted by the join keys.
///
/// # Motivation
///
/// DataFusion's default physical planner prefers `HashJoinExec` for most joins, especially when
/// `prefer_hash_join` is true (default). However, `pg_search` indexes (BM25) can provide pre-sorted
/// data via `MultiSegmentPlan` when configured with `sort_by`. In these cases, a `SortMergeJoinExec`
/// is significantly more efficient as it streams data without needing to build a hash table or sort.
///
/// # Capabilities
///
/// This rule performs two key optimizations:
/// 1.  **HashJoin to SortMergeJoin Conversion**: It detects if the inputs to a `HashJoinExec` are
///     physically sorted by the join keys. If so, it replaces the HashJoin with a SortMergeJoin.
/// 2.  **Partitioning Recovery**: When `target_partitions=1`, DataFusion inserts `CoalescePartitionsExec`
///     to merge parallel streams. This destroys ordering. This rule detects if `CoalescePartitionsExec`
///     wraps a sorted multi-partition input and replaces it with `SortPreservingMergeExec`, restoring
///     the sort property and enabling the SortMergeJoin optimization.
///
/// # Usage
///
/// This rule should be registered in the `SessionState` physical optimizer rules. It runs post-planning,
/// allowing it to "upgrade" the default plan generated by DataFusion.
pub struct SortMergeJoinEnforcer {}

impl SortMergeJoinEnforcer {
    pub fn new() -> Self {
        Self {}
    }

    /// Attempts to convert a `HashJoinExec` node into a `SortMergeJoinExec`.
    ///
    /// This method checks if both the left and right inputs are sorted by their respective join keys.
    /// It handles three scenarios for each side to ensure a single sorted stream is provided
    /// to the `SortMergeJoinExec`:
    /// 1.  The input is already sorted and has a single partition.
    /// 2.  The input is a `CoalescePartitionsExec` wrapping sorted partitions. In this case,
    ///     it replaces the destructive Coalesce with `SortPreservingMergeExec`.
    /// 3.  The input is sorted but has multiple partitions. It wraps it in `SortPreservingMergeExec`
    ///     to ensure a single sorted stream, satisfying `SortMergeJoinExec` requirements and
    ///     avoiding partition count mismatches.
    fn try_convert_to_smj(
        &self,
        hash_join: &HashJoinExec,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        let left = hash_join.left();
        let right = hash_join.right();
        let on = hash_join.on();

        if on.is_empty() {
            return Ok(None);
        }

        // Extract join keys from the HashJoin
        let left_keys: Vec<Arc<dyn PhysicalExpr>> = on.iter().map(|(l, _)| l.clone()).collect();
        let right_keys: Vec<Arc<dyn PhysicalExpr>> = on.iter().map(|(_, r)| r.clone()).collect();

        // Helper to check sort or upgrade Coalesce -> SPM
        // Returns (NewPlan, SortOptions) if sorted/sortable, else None
        let check_side =
            |plan: Arc<dyn ExecutionPlan>,
             keys: &[Arc<dyn PhysicalExpr>]|
             -> Option<(Arc<dyn ExecutionPlan>, Vec<arrow_schema::SortOptions>)> {
                // 1. Check if CoalescePartitions wrapping sorted partitions
                // We check this first to "rescue" sortedness from Coalesce
                if let Some(coalesce) = plan.as_any().downcast_ref::<CoalescePartitionsExec>() {
                    let child = coalesce.input().clone();
                    if let Some(opts) = check_ordering(child.equivalence_properties(), keys) {
                        // Replace Coalesce with SortPreservingMerge
                        let ordering_vec: Vec<PhysicalSortExpr> = keys
                            .iter()
                            .zip(&opts)
                            .map(|(expr, opt)| PhysicalSortExpr {
                                expr: expr.clone(),
                                options: *opt,
                            })
                            .collect();

                        let ordering = LexOrdering::new(ordering_vec).expect("valid ordering");
                        let spm = Arc::new(SortPreservingMergeExec::new(ordering, child));
                        return Some((spm, opts));
                    }
                }

                // 2. Check if plan itself is sorted
                if let Some(opts) = check_ordering(plan.equivalence_properties(), keys) {
                    // If plan is sorted but has multiple partitions, wrap in SortPreservingMerge
                    // to ensure we provide a single sorted stream to SortMergeJoin.
                    // This is necessary because HashJoin might have left one side partitioned
                    // (CollectLeft), leading to partition mismatch if we don't merge.
                    if plan.output_partitioning().partition_count() > 1 {
                        let ordering_vec: Vec<PhysicalSortExpr> = keys
                            .iter()
                            .zip(&opts)
                            .map(|(expr, opt)| PhysicalSortExpr {
                                expr: expr.clone(),
                                options: *opt,
                            })
                            .collect();

                        let ordering = LexOrdering::new(ordering_vec).expect("valid ordering");
                        let spm = Arc::new(SortPreservingMergeExec::new(ordering, plan));
                        return Some((spm, opts));
                    }

                    return Some((plan, opts));
                }

                None
            };

        // Try to upgrade both sides
        if let Some((new_left, left_opts)) = check_side(left.clone(), &left_keys) {
            if let Some((new_right, right_opts)) = check_side(right.clone(), &right_keys) {
                // Ensure SortOptions match between left and right
                if left_opts == right_opts {
                    let null_equality = hash_join.null_equality();

                    // SortMergeJoinExec's filter evaluation (get_filter_column) assumes
                    // column_indices are ordered Left-first, Right-second. Reorder if needed.
                    let filter = hash_join.filter().map(normalize_join_filter_ordering);

                    let exec = SortMergeJoinExec::try_new(
                        new_left,
                        new_right,
                        on.to_vec(),
                        filter,
                        *hash_join.join_type(),
                        left_opts, // sort options match
                        null_equality,
                    )?;

                    // TODO: See https://github.com/apache/datafusion/pull/20455
                    let exec = Arc::new(FilterPassthroughExec::new(Arc::new(exec)))
                        as Arc<dyn ExecutionPlan>;

                    // HashJoinExec might have an internal projection (pruning columns).
                    // SortMergeJoinExec outputs all columns from both sides.
                    // If HashJoinExec has a projection, we must wrap the SortMergeJoinExec in a ProjectionExec
                    // to match the output schema expected by the parent node.
                    if let Some(projection) = hash_join.projection.as_ref() {
                        let schema = exec.schema();
                        let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = projection
                            .iter()
                            .map(|&i| {
                                let field = schema.field(i);
                                let col = Arc::new(Column::new(field.name(), i));
                                (col as Arc<dyn PhysicalExpr>, field.name().clone())
                            })
                            .collect();
                        let proj = ProjectionExec::try_new(exprs, exec)?;
                        return Ok(Some(Arc::new(proj)));
                    }

                    return Ok(Some(exec));
                }
            }
        }

        Ok(None)
    }
}

impl PhysicalOptimizerRule for SortMergeJoinEnforcer {
    fn optimize(
        &self,
        plan: Arc<dyn ExecutionPlan>,
        _config: &ConfigOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        plan.transform_up(|plan| {
            if let Some(hash_join) = plan.as_any().downcast_ref::<HashJoinExec>() {
                if let Some(smj) = self.try_convert_to_smj(hash_join)? {
                    return Ok(Transformed::yes(smj));
                }
            }
            Ok(Transformed::no(plan))
        })
        .map(|t| t.data)
    }

    fn name(&self) -> &str {
        "SortMergeJoinEnforcer"
    }

    fn schema_check(&self) -> bool {
        true
    }
}

impl std::fmt::Debug for SortMergeJoinEnforcer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SortMergeJoinEnforcer").finish()
    }
}

/// Checks if the plan's ordering satisfies the given keys.
///
/// This function determines if the physical output of a plan is pre-sorted by the
/// specified `keys` (e.g., join keys). It uses `EquivalenceProperties` to account
/// for column equivalences established by prior joins or filters. For example,
/// if a prior join established that `t1.id == t2.t1_id`, and the input is sorted
/// by `t1.id`, this function will recognize it as also being sorted by `t2.t1_id`.
///
/// Returns the detected `SortOptions` (ASC/DESC, Nulls First/Last) for the keys
/// if the ordering is satisfied.
fn check_ordering(
    eq_properties: &EquivalenceProperties,
    keys: &[Arc<dyn PhysicalExpr>],
) -> Option<Vec<arrow_schema::SortOptions>> {
    let output_ordering = eq_properties.output_ordering()?;

    if output_ordering.len() < keys.len() {
        return None;
    }

    // To use `ordering_satisfy`, we need a `LexOrdering` (expressions + sort options).
    // Since join keys don't inherently have sort options, we "borrow" the options
    // from the actual physical output ordering at the corresponding positions.
    let mut proposed_ordering = Vec::with_capacity(keys.len());
    let mut options = Vec::with_capacity(keys.len());

    for (i, key) in keys.iter().enumerate() {
        let sort_options = output_ordering[i].options;
        proposed_ordering.push(PhysicalSortExpr {
            expr: key.clone(),
            options: sort_options,
        });
        options.push(sort_options);
    }

    // Verify if the proposed ordering (keys + borrowed options) is logically
    // satisfied by the actual physical ordering, considering equivalences.
    if eq_properties
        .ordering_satisfy(proposed_ordering)
        .unwrap_or(false)
    {
        return Some(options);
    }

    None
}

/// Reorders a `JoinFilter`'s column_indices to place Left-side columns before Right-side
/// columns, which is required by `SortMergeJoinExec`'s `get_filter_column` implementation.
///
/// DataFusion's `SortMergeJoinExec` filter evaluation builds the intermediate filter batch
/// by collecting all Left columns first, then all Right columns. However, the filter's
/// schema and expression reference columns by their original position in `column_indices`.
/// If `column_indices` has Right entries interleaved before Left entries, the filter batch
/// columns won't match the schema, causing type mismatch errors.
///
/// This function normalizes the ordering by:
/// 1. Grouping column_indices into Left-first, Right-second order
/// 2. Remapping column references in the expression to the new positions
/// 3. Rebuilding the schema to match the new column order
fn normalize_join_filter_ordering(filter: &JoinFilter) -> JoinFilter {
    use datafusion::common::JoinSide;

    let column_indices = filter.column_indices();

    // Check if already in Left-first, Right-second order (fast path)
    let mut seen_right = false;
    let mut needs_reorder = false;
    for ci in column_indices {
        match ci.side {
            JoinSide::Right => seen_right = true,
            JoinSide::Left if seen_right => {
                needs_reorder = true;
                break;
            }
            _ => {}
        }
    }

    if !needs_reorder {
        return filter.clone();
    }

    // Build the reordered column_indices: Left entries first, then Right entries,
    // preserving relative order within each group.
    let mut left_old_positions = Vec::new();
    let mut right_old_positions = Vec::new();
    for (old_idx, ci) in column_indices.iter().enumerate() {
        match ci.side {
            JoinSide::Left => left_old_positions.push(old_idx),
            JoinSide::Right => right_old_positions.push(old_idx),
            _ => left_old_positions.push(old_idx),
        }
    }

    // Build mapping: old_index -> new_index
    let mut old_to_new = vec![0usize; column_indices.len()];
    let mut new_column_indices = Vec::with_capacity(column_indices.len());

    for (new_idx, &old_idx) in left_old_positions
        .iter()
        .chain(right_old_positions.iter())
        .enumerate()
    {
        old_to_new[old_idx] = new_idx;
        new_column_indices.push(column_indices[old_idx].clone());
    }

    // Build new schema in the reordered column order
    let old_schema = filter.schema();
    let new_fields: Vec<arrow_schema::Field> = left_old_positions
        .iter()
        .chain(right_old_positions.iter())
        .map(|&old_idx| old_schema.field(old_idx).clone())
        .collect();
    let new_schema = Arc::new(arrow_schema::Schema::new(new_fields));

    // Remap Column references in the expression to use new indices
    let new_expression = filter
        .expression()
        .clone()
        .transform_up(|expr| {
            if let Some(col) = expr.as_any().downcast_ref::<Column>() {
                let new_index = old_to_new[col.index()];
                if new_index != col.index() {
                    return Ok(Transformed::yes(Arc::new(Column::new(
                        col.name(),
                        new_index,
                    ))));
                }
            }
            Ok(Transformed::no(expr))
        })
        .expect("filter expression rewrite should not fail")
        .data;

    JoinFilter::new(new_expression, new_column_indices, new_schema)
}
