// 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 datafusion::common::{Column, JoinType, Result, ScalarValue, TableReference};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{col, lit, BinaryExpr, Expr, Operator};
use datafusion::prelude::DataFrame;
use pgrx::{pg_sys, PgList};

use crate::api::{HashMap, HashSet};
use crate::postgres::customscan::joinscan::build::{
    JoinLevelExpr, JoinLevelSearchPredicate, JoinNode, JoinSource, JoinType as PgJoinType,
    RelationAlias,
};
use crate::postgres::customscan::joinscan::privdat::{OutputColumnInfo, SCORE_COL_NAME};
use crate::postgres::customscan::opexpr::lookup_operator;
use crate::scan::SearchPredicateUDF;

pub trait ColumnMapper {
    /// Map a PostgreSQL variable to a DataFusion Column expression
    fn map_var(&self, varno: pg_sys::Index, varattno: pg_sys::AttrNumber) -> Option<Expr>;
}

/// Helper struct for translating PostgreSQL expression trees into DataFusion `Expr`s.
pub struct PredicateTranslator<'a> {
    pub sources: &'a [&'a JoinSource],
    mapper: Option<Box<dyn ColumnMapper + 'a>>,
}

impl<'a> PredicateTranslator<'a> {
    pub fn new(sources: &'a [&'a JoinSource]) -> Self {
        Self {
            sources,
            mapper: None,
        }
    }

    pub fn with_mapper(mut self, mapper: Box<dyn ColumnMapper + 'a>) -> Self {
        self.mapper = Some(mapper);
        self
    }

    /// Translate a `JoinLevelExpr` tree to a DataFusion `Expr`.
    ///
    /// This creates `SearchPredicateUDF` expressions for single-table predicates,
    /// which can be pushed down to `PgSearchTableProvider` via DataFusion's
    /// filter pushdown mechanism.
    /// Translate a `JoinLevelExpr` tree to a DataFusion `Expr`.
    ///
    /// `deferred_positions` contains plan_positions whose ctid columns still hold
    /// packed DocAddresses at this point in the plan. Each `SingleTablePredicate`
    /// checks whether its `plan_position` is in the set to determine whether its
    /// `SearchPredicateUDF` should emit packed DocAddresses or real ctids.
    pub unsafe fn translate_join_level_expr(
        expr: &JoinLevelExpr,
        custom_exprs: &[Expr],
        ctid_map: &HashMap<pg_sys::Index, Expr>,
        predicates: &[JoinLevelSearchPredicate],
        deferred_positions: &crate::api::HashSet<usize>,
        sources: &[&JoinSource],
    ) -> Option<Expr> {
        match expr {
            JoinLevelExpr::SingleTablePredicate {
                plan_position,
                predicate_idx,
            } => {
                let predicate = predicates.get(*predicate_idx)?;
                let col = ctid_map.get(&(*plan_position as pg_sys::Index))?;
                let deferred = deferred_positions.contains(plan_position);
                let udf = SearchPredicateUDF::with_deferred_visibility(
                    predicate.indexrelid,
                    predicate.heaprelid,
                    predicate.query.clone(),
                    predicate.display_string.clone(),
                    deferred,
                    Some(*plan_position),
                );
                Some(udf.into_expr(col.clone()))
            }
            JoinLevelExpr::MultiTablePredicate { predicate_idx } => {
                custom_exprs.get(*predicate_idx).cloned()
            }
            JoinLevelExpr::And(children) => {
                if children.is_empty() {
                    return None;
                }
                let mut result = Self::translate_join_level_expr(
                    &children[0],
                    custom_exprs,
                    ctid_map,
                    predicates,
                    deferred_positions,
                    sources,
                )?;
                for child in &children[1..] {
                    let right = Self::translate_join_level_expr(
                        child,
                        custom_exprs,
                        ctid_map,
                        predicates,
                        deferred_positions,
                        sources,
                    )?;
                    result = Expr::BinaryExpr(BinaryExpr::new(
                        Box::new(result),
                        Operator::And,
                        Box::new(right),
                    ));
                }
                Some(result)
            }
            JoinLevelExpr::Or(children) => {
                if children.is_empty() {
                    return None;
                }
                let mut result = Self::translate_join_level_expr(
                    &children[0],
                    custom_exprs,
                    ctid_map,
                    predicates,
                    deferred_positions,
                    sources,
                )?;
                for child in &children[1..] {
                    let right = Self::translate_join_level_expr(
                        child,
                        custom_exprs,
                        ctid_map,
                        predicates,
                        deferred_positions,
                        sources,
                    )?;
                    result = Expr::BinaryExpr(BinaryExpr::new(
                        Box::new(result),
                        Operator::Or,
                        Box::new(right),
                    ));
                }
                Some(result)
            }
            JoinLevelExpr::Not(child) => {
                let inner = Self::translate_join_level_expr(
                    child,
                    custom_exprs,
                    ctid_map,
                    predicates,
                    deferred_positions,
                    sources,
                )?;
                Some(Expr::Not(Box::new(inner)))
            }
            JoinLevelExpr::MarkOrNull {
                is_anti,
                null_test_varno,
                null_test_attno,
            } => {
                // Resolve the outer column name from source metadata.
                let source = sources.iter().find(|s| s.contains_rti(*null_test_varno));
                let col_name = source.and_then(|s| s.column_name(*null_test_attno))?;

                // Build: mark = true OR col IS NULL  (for IN)
                //        mark = false OR col IS NULL  (for NOT IN)
                let mark_check = if *is_anti {
                    col("mark").eq(lit(false))
                } else {
                    col("mark").eq(lit(true))
                };
                let null_check = Expr::IsNull(Box::new(col(col_name)));
                Some(mark_check.or(null_check))
            }
        }
    }

    /// Translate a PostgreSQL expression to a DataFusion `Expr`.
    ///
    /// Returns `None` if the expression cannot be translated.
    ///
    /// IMPORTANT: This translator is used to check if a predicate CAN be translated,
    /// but the actual predicate evaluation happens via heap fetch + PostgreSQL evaluation.
    /// Cross-type comparisons (e.g., INT < NUMERIC) involve type casts that change value
    /// semantics - we cannot simply look through them because the underlying storage
    /// representations differ (e.g., INT 95 vs Numeric64 5225 for 52.25).
    ///
    /// For predicates involving type casts, we return None to indicate that the predicate
    /// cannot be evaluated purely in DataFusion and must fall back to PostgreSQL evaluation.
    pub unsafe fn translate(&self, node: *mut pg_sys::Node) -> Option<Expr> {
        if node.is_null() {
            return None;
        }

        match (*node).type_ {
            pg_sys::NodeTag::T_OpExpr => self.translate_op_expr(node as *mut pg_sys::OpExpr),
            pg_sys::NodeTag::T_Var => self.translate_var(node as *mut pg_sys::Var),
            pg_sys::NodeTag::T_Const => self.translate_const(node as *mut pg_sys::Const),
            pg_sys::NodeTag::T_BoolExpr => self.translate_bool_expr(node as *mut pg_sys::BoolExpr),
            // Type casts (RelabelType, CoerceViaIO, FuncExpr) are not supported because
            // they may change value semantics. Cross-type comparisons like INT < NUMERIC
            // require proper type coercion that DataFusion cannot perform correctly when
            // the underlying fast field storage uses different scales/representations.
            _ => None,
        }
    }

    unsafe fn translate_op_expr(&self, op_expr: *mut pg_sys::OpExpr) -> Option<Expr> {
        let args = PgList::<pg_sys::Node>::from_pg((*op_expr).args);
        if args.len() != 2 {
            return None; // Only support binary operators for now
        }

        let left = self.translate(args.get_ptr(0)?)?;
        let right = self.translate(args.get_ptr(1)?)?;

        let opno = (*op_expr).opno;
        let op_str = lookup_operator(opno)?;

        let op = match op_str {
            "=" => Operator::Eq,
            "<>" => Operator::NotEq,
            "<" => Operator::Lt,
            "<=" => Operator::LtEq,
            ">" => Operator::Gt,
            ">=" => Operator::GtEq,
            _ => return None,
        };

        Some(Expr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            op,
            Box::new(right),
        )))
    }

    unsafe fn translate_var(&self, var: *mut pg_sys::Var) -> Option<Expr> {
        let varno = (*var).varno as pg_sys::Index;
        let varattno = (*var).varattno;

        // INDEX_VAR is allowed because it represents a reference to the custom scan's output
        if varno != pg_sys::INDEX_VAR as pg_sys::Index
            && !self.sources.iter().any(|s| s.contains_rti(varno))
        {
            return None;
        }

        if let Some(ref mapper) = self.mapper {
            if let Some(expr) = mapper.map_var(varno, varattno) {
                return Some(expr);
            }
            return None;
        }

        Some(col("placeholder"))
    }

    unsafe fn translate_const(&self, c: *mut pg_sys::Const) -> Option<Expr> {
        if (*c).constisnull {
            return Some(lit(ScalarValue::Null));
        }

        let type_oid = (*c).consttype;
        let datum = (*c).constvalue;

        // Simple mapping for common types
        let scalar_value = match type_oid {
            pg_sys::INT2OID => {
                let val = datum.value() as i16;
                ScalarValue::Int16(Some(val))
            }
            pg_sys::INT4OID => {
                let val = datum.value() as i32;
                ScalarValue::Int32(Some(val))
            }
            pg_sys::INT8OID => {
                let val = datum.value() as i64;
                ScalarValue::Int64(Some(val))
            }
            pg_sys::FLOAT4OID => {
                let val = f32::from_bits(datum.value() as u32);
                ScalarValue::Float32(Some(val))
            }
            pg_sys::FLOAT8OID => {
                let val = f64::from_bits(datum.value() as u64);
                ScalarValue::Float64(Some(val))
            }
            pg_sys::BOOLOID => {
                let val = datum.value() != 0;
                ScalarValue::Boolean(Some(val))
            }
            _ => return None,
        };

        Some(lit(scalar_value))
    }

    unsafe fn translate_bool_expr(&self, bool_expr: *mut pg_sys::BoolExpr) -> Option<Expr> {
        let args = PgList::<pg_sys::Node>::from_pg((*bool_expr).args);

        match (*bool_expr).boolop {
            pg_sys::BoolExprType::AND_EXPR => {
                if args.len() < 2 {
                    return None;
                }
                let mut expr = self.translate(args.get_ptr(0)?)?;
                for i in 1..args.len() {
                    let right = self.translate(args.get_ptr(i)?)?;
                    expr = Expr::BinaryExpr(BinaryExpr::new(
                        Box::new(expr),
                        Operator::And,
                        Box::new(right),
                    ));
                }
                Some(expr)
            }
            pg_sys::BoolExprType::OR_EXPR => {
                if args.len() < 2 {
                    return None;
                }
                let mut expr = self.translate(args.get_ptr(0)?)?;
                for i in 1..args.len() {
                    let right = self.translate(args.get_ptr(i)?)?;
                    expr = Expr::BinaryExpr(BinaryExpr::new(
                        Box::new(expr),
                        Operator::Or,
                        Box::new(right),
                    ));
                }
                Some(expr)
            }
            pg_sys::BoolExprType::NOT_EXPR => {
                if args.len() != 1 {
                    return None;
                }
                let child = self.translate(args.get_ptr(0)?)?;
                Some(Expr::Not(Box::new(child)))
            }
            _ => None,
        }
    }
}

/// Translate a `JoinLevelExpr` predicate into a DataFusion filter expression
/// and apply it to `df`. This is the shared spine of the `RelNode::Filter`
/// arm in JoinScan's and AggregateScan's `build_relnode_df` recursions —
/// both modules previously inlined the same translate-and-filter sequence.
///
/// `handle_mark` controls JoinScan-specific cleanup: when `true`, a
/// `MarkOrNull` predicate triggers a post-filter projection that drops the
/// synthetic `mark` column from the schema. AggregateScan never produces
/// `MarkOrNull` predicates and passes `false`.
#[allow(clippy::too_many_arguments)]
pub fn apply_join_level_filter(
    mut df: DataFrame,
    predicate: &JoinLevelExpr,
    translated_exprs: &[Expr],
    ctid_map: &HashMap<pg_sys::Index, Expr>,
    join_level_predicates: &[JoinLevelSearchPredicate],
    deferred_positions: &HashSet<usize>,
    sources: &[&JoinSource],
    handle_mark: bool,
) -> Result<DataFrame> {
    let filter_expr = unsafe {
        PredicateTranslator::translate_join_level_expr(
            predicate,
            translated_exprs,
            ctid_map,
            join_level_predicates,
            deferred_positions,
            sources,
        )
    }
    .ok_or_else(|| {
        DataFusionError::Internal(format!(
            "Failed to translate join level expression tree: {:?}",
            predicate
        ))
    })?;

    df = df.filter(filter_expr)?;

    // For MarkOrNull filters in JoinScan, drop the synthetic "mark" column
    // post-filter so it doesn't leak into the projection.
    if handle_mark && matches!(predicate, JoinLevelExpr::MarkOrNull { .. }) {
        let schema = df.schema().clone();
        let proj_cols: Vec<Expr> = schema
            .columns()
            .into_iter()
            .filter(|c| c.name != "mark")
            .map(col)
            .collect();
        if !proj_cols.is_empty() {
            df = df.select(proj_cols)?;
        }
    }

    Ok(df)
}

/// Creates a DataFusion column expression with a bare table reference.
/// This is preferred over `datafusion::logical_expr::col()` because `col()` parses the input string,
pub fn make_col(relation: &str, name: &str) -> Expr {
    Expr::Column(Column::new(
        Some(TableReference::Bare {
            table: relation.into(),
        }),
        name,
    ))
}

/// Allow-list of `JoinNode::join_type` variants accepted by [`build_join_df`].
///
/// JoinScan supports the full set; AggregateScan only the four equi-join
/// variants. The variant choice is exposed at the call site so the policy
/// difference is grep-able.
#[derive(Copy, Clone, Debug)]
pub enum JoinTypeAllowList {
    /// JoinScan: Inner, Left, Right, Full, Semi, Anti, LeftMark, RightMark,
    /// RightSemi, RightAnti.
    All,
    /// AggregateScan: Inner, Left, Right, Full only.
    EquiOnly,
}

/// Lower a `JoinNode` over already-built left/right `DataFrame`s.
///
/// Builds equi-join keys with [`build_equi_join_exprs`], maps the
/// [`super::build::JoinType`] into DataFusion's `JoinType` (subject to
/// `allowed_join_types`), and dispatches to `join_on` (when there are
/// equi keys) or `join` (cross join). Caller is responsible for any
/// post-join filter handling — `JoinNode::filter` is not consulted here.
///
/// Returns `Err(NotImplemented)` if the join type is outside the allow-list.
pub fn build_join_df(
    left: DataFrame,
    right: DataFrame,
    join: &JoinNode,
    allowed_join_types: JoinTypeAllowList,
) -> Result<DataFrame> {
    let on = build_equi_join_exprs(join)?;

    let df_join_type = match (join.join_type, allowed_join_types) {
        (PgJoinType::Inner, _) => JoinType::Inner,
        (PgJoinType::Left, _) => JoinType::Left,
        (PgJoinType::Right, _) => JoinType::Right,
        (PgJoinType::Full, _) => JoinType::Full,
        (PgJoinType::Semi, JoinTypeAllowList::All) => JoinType::LeftSemi,
        (PgJoinType::Anti, JoinTypeAllowList::All) => JoinType::LeftAnti,
        (PgJoinType::LeftMark, JoinTypeAllowList::All) => JoinType::LeftMark,
        (PgJoinType::RightMark, JoinTypeAllowList::All) => JoinType::RightMark,
        (PgJoinType::RightSemi, JoinTypeAllowList::All) => JoinType::RightSemi,
        (PgJoinType::RightAnti, JoinTypeAllowList::All) => JoinType::RightAnti,
        (jt, JoinTypeAllowList::EquiOnly) => {
            return Err(DataFusionError::NotImplemented(format!(
                "Aggregate-on-join does not support {} JOIN",
                jt
            )));
        }
        (jt, JoinTypeAllowList::All) => {
            return Err(DataFusionError::NotImplemented(format!(
                "{} JOIN is not supported in JoinScan execution",
                jt
            )));
        }
    };

    if on.is_empty() {
        left.join(right, df_join_type, &[], &[], None)
    } else {
        left.join_on(right, df_join_type, on)
    }
}

/// Build equi-join filter expressions from a [`JoinNode`]'s key pairs.
///
/// Shared between JoinScan and AggregateScan `build_relnode_df` implementations.
/// Each key pair is resolved against the join's left/right subtrees and converted
/// to a `left_col = right_col` DataFusion expression.
pub fn build_equi_join_exprs(join: &JoinNode) -> Result<Vec<Expr>> {
    let mut on = Vec::with_capacity(join.equi_keys.len());
    for jk in &join.equi_keys {
        let ((left_source, left_attno), (right_source, right_attno)) =
            jk.resolve_against(&join.left, &join.right).ok_or_else(|| {
                DataFusionError::Internal(format!(
                    "Failed to resolve join key to current join sides: outer_rti={}, inner_rti={}",
                    jk.outer_rti, jk.inner_rti
                ))
            })?;

        let left_col_name = left_source
            .column_name(left_attno)
            .ok_or_else(|| DataFusionError::Internal("Missing left join-key column".into()))?;
        let right_col_name = right_source
            .column_name(right_attno)
            .ok_or_else(|| DataFusionError::Internal("Missing right join-key column".into()))?;

        let left_expr = make_source_col(left_source, &left_col_name);
        let right_expr = make_source_col(right_source, &right_col_name);
        on.push(left_expr.eq(right_expr));
    }
    Ok(on)
}

/// Build a DataFusion column expression for a `(source, field_name)` pair.
///
/// The execution alias is built from the source's optional alias and its
/// `plan_position` (the DFS index assigned by `JoinCSClause::new`). Both
/// JoinScan and AggregateScan use this exact pattern; sharing it keeps the
/// alias-construction policy in one place.
pub fn make_source_col(source: &JoinSource, field_name: &str) -> Expr {
    let alias =
        RelationAlias::new(source.scan_info.alias.as_deref()).execution(source.plan_position);
    make_col(&alias, field_name)
}

/// Build a DataFusion column expression for the synthetic score column on the
/// given source. Equivalent to `make_source_col(source, SCORE_COL_NAME)`.
pub fn make_source_score_col(source: &JoinSource) -> Expr {
    make_source_col(source, SCORE_COL_NAME)
}

pub struct CombinedMapper<'a> {
    pub sources: &'a [&'a JoinSource],
    pub output_columns: &'a [OutputColumnInfo],
}

impl<'a> ColumnMapper for CombinedMapper<'a> {
    fn map_var(&self, varno: pg_sys::Index, varattno: pg_sys::AttrNumber) -> Option<Expr> {
        let (rti, attno, is_score) = if varno == pg_sys::INDEX_VAR as pg_sys::Index {
            let idx = (varattno - 1) as usize;
            match self.output_columns.get(idx)? {
                OutputColumnInfo::Var {
                    rti,
                    original_attno,
                    ..
                } => (*rti, *original_attno, false),
                OutputColumnInfo::Score { rti, .. } => (*rti, 0, true),
                OutputColumnInfo::Pruned => return None,
            }
        } else {
            (varno, varattno, false)
        };

        let (plan_position, source) = self
            .sources
            .iter()
            .enumerate()
            .find(|(_, s)| s.contains_rti(rti))?;

        let alias = RelationAlias::new(source.scan_info.alias.as_deref()).execution(plan_position);

        if is_score {
            if let Some(col_idx) = source.map_var(rti, 0) {
                if let Some(name) = source.column_name(col_idx) {
                    return Some(make_col(&alias, &name));
                }
            }
            return Some(make_col(&alias, SCORE_COL_NAME));
        }

        let mapped_attno = source.map_var(rti, attno)?;
        let col_name = source.column_name(mapped_attno)?;
        Some(make_col(&alias, &col_name))
    }
}
