//! DataFusion translations for membership operators (IN, NOT IN).

use datafusion::logical_expr::expr::InList;
use datafusion::logical_expr::{BinaryExpr, Expr as DFExpr, Operator as DFOperator};
use datafusion_functions::core::expr_ext::FieldAccessor;
use datafusion_functions::core::expr_fn as core_fn;
use datafusion_functions_nested::expr_fn as array_fn;

use hamelin_lib::func::defs::{
    InArray, InMap, InRange, InTimestampInterval, InTimestampTimestamp, InTuple, NotInArray,
    NotInMap, NotInRange, NotInTimestampInterval, NotInTimestampTimestamp, NotInTuple,
    NumericRange, NumericRangeInclusive, TimestampRange, TimestampRangeInclusive,
};
use hamelin_lib::types::Type;

use super::DataFusionTranslationRegistry;

/// Try to extract tuple elements from a DataFusion expression.
///
/// When a Hamelin tuple is translated, it becomes `named_struct("f0", e0, "f1", e1, ...)`
/// which is represented as a ScalarFunction with the "named_struct" UDF. This function
/// extracts the original value elements (every other arg starting from index 1).
///
/// Returns None if the expression is not a named_struct() call, or if the elements cannot
/// be extracted.
fn try_extract_tuple_elements(expr: DFExpr) -> Option<Vec<DFExpr>> {
    match expr {
        DFExpr::ScalarFunction(scalar_func) => {
            if scalar_func.func.name() == "named_struct" {
                // named_struct args are: name0, val0, name1, val1, ...
                // Extract only the value args (odd indices)
                Some(scalar_func.args.into_iter().skip(1).step_by(2).collect())
            } else {
                None
            }
        }
        _ => None,
    }
}

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // left IN array -> array_has(array, left)
    registry.register::<InArray>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(array_fn::array_has(right, left))
    });

    // left NOT IN array -> NOT array_has(array, left)
    registry.register::<NotInArray>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(DFExpr::Not(Box::new(array_fn::array_has(right, left))))
    });

    // left IN map -> map_contains_key(map, left)
    registry.register::<InMap>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        // DataFusion doesn't have map_contains_key directly
        // Use: array_has(map_keys(map), key)
        Ok(array_fn::array_has(array_fn::map_keys(right), left))
    });

    // left NOT IN map -> NOT map_contains_key(map, left)
    registry.register::<NotInMap>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        // Use: NOT array_has(map_keys(map), key)
        Ok(DFExpr::Not(Box::new(array_fn::array_has(
            array_fn::map_keys(right),
            left,
        ))))
    });

    // left..right -> named_struct("begin", left, "end", right)
    registry.register::<NumericRange>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(core_fn::named_struct(vec![
            datafusion::logical_expr::lit("begin"),
            left,
            datafusion::logical_expr::lit("end"),
            right,
        ]))
    });

    // left..=right -> named_struct("begin", left, "end", right) (inclusive)
    registry.register::<NumericRangeInclusive>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(core_fn::named_struct(vec![
            datafusion::logical_expr::lit("begin"),
            left,
            datafusion::logical_expr::lit("end"),
            right,
        ]))
    });

    // timestamp/interval range
    registry.register::<TimestampRange>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(core_fn::named_struct(vec![
            datafusion::logical_expr::lit("begin"),
            left,
            datafusion::logical_expr::lit("end"),
            right,
        ]))
    });

    // timestamp/interval inclusive range
    registry.register::<TimestampRangeInclusive>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        Ok(core_fn::named_struct(vec![
            datafusion::logical_expr::lit("begin"),
            left,
            datafusion::logical_expr::lit("end"),
            right,
        ]))
    });

    // left IN range -> (begin IS NULL OR left >= begin) AND (end IS NULL OR left </<= end)
    // End operator depends on range type's end_inclusive flag: .. uses <, ..= uses <=
    registry.register::<InRange>(|mut params| {
        let left = params.take()?;
        let right = params.take()?;
        let end_op = if matches!(&*right.typ, Type::RangeInclusive(_)) {
            DFOperator::LtEq
        } else {
            DFOperator::Lt
        };
        let begin = right.expr.clone().field("begin");
        let end = right.expr.field("end");

        let begin_is_null = DFExpr::IsNull(Box::new(begin.clone()));
        let gte_begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left.expr.clone()),
            DFOperator::GtEq,
            Box::new(begin),
        ));
        let begin_check = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(begin_is_null),
            DFOperator::Or,
            Box::new(gte_begin),
        ));

        let end_is_null = DFExpr::IsNull(Box::new(end.clone()));
        let end_cmp =
            DFExpr::BinaryExpr(BinaryExpr::new(Box::new(left.expr), end_op, Box::new(end)));
        let end_check = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(end_is_null),
            DFOperator::Or,
            Box::new(end_cmp),
        ));

        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(begin_check),
            DFOperator::And,
            Box::new(end_check),
        )))
    });

    // left NOT IN range -> NOT((begin IS NULL OR left >= begin) AND (end IS NULL OR left </<= end))
    registry.register::<NotInRange>(|mut params| {
        let left = params.take()?;
        let right = params.take()?;
        let end_op = if matches!(&*right.typ, Type::RangeInclusive(_)) {
            DFOperator::LtEq
        } else {
            DFOperator::Lt
        };
        let begin = right.expr.clone().field("begin");
        let end = right.expr.field("end");

        let begin_is_null = DFExpr::IsNull(Box::new(begin.clone()));
        let gte_begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left.expr.clone()),
            DFOperator::GtEq,
            Box::new(begin),
        ));
        let begin_check = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(begin_is_null),
            DFOperator::Or,
            Box::new(gte_begin),
        ));

        let end_is_null = DFExpr::IsNull(Box::new(end.clone()));
        let end_cmp =
            DFExpr::BinaryExpr(BinaryExpr::new(Box::new(left.expr), end_op, Box::new(end)));
        let end_check = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(end_is_null),
            DFOperator::Or,
            Box::new(end_cmp),
        ));

        let in_range = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(begin_check),
            DFOperator::And,
            Box::new(end_check),
        ));
        Ok(DFExpr::Not(Box::new(in_range)))
    });

    // timestamp IN interval -> timestamp >= now() + interval AND timestamp <= now()
    // Inclusive end: user wrote `IN interval`, not `..` syntax.
    registry.register::<InTimestampInterval>(|mut params| {
        let timestamp = params.take()?.expr;
        let interval = params.take()?.expr;
        let now = datafusion_functions::datetime::expr_fn::now();
        let begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(now.clone()),
            DFOperator::Plus,
            Box::new(interval),
        ));
        let gte_begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(timestamp.clone()),
            DFOperator::GtEq,
            Box::new(begin),
        ));
        let lte_now = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(timestamp),
            DFOperator::LtEq,
            Box::new(now),
        ));
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(gte_begin),
            DFOperator::And,
            Box::new(lte_now),
        )))
    });

    // timestamp NOT IN interval
    // Inclusive end: user wrote `NOT IN interval`, not `..` syntax.
    registry.register::<NotInTimestampInterval>(|mut params| {
        let timestamp = params.take()?.expr;
        let interval = params.take()?.expr;
        let now = datafusion_functions::datetime::expr_fn::now();
        let begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(now.clone()),
            DFOperator::Plus,
            Box::new(interval),
        ));
        let gte_begin = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(timestamp.clone()),
            DFOperator::GtEq,
            Box::new(begin),
        ));
        let lte_now = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(timestamp),
            DFOperator::LtEq,
            Box::new(now),
        ));
        let in_interval = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(gte_begin),
            DFOperator::And,
            Box::new(lte_now),
        ));
        Ok(DFExpr::Not(Box::new(in_interval)))
    });

    // timestamp IN timestamp -> left >= right AND left <= now()
    // Inclusive end: user wrote `IN timestamp`, not `..` syntax.
    registry.register::<InTimestampTimestamp>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        let now = datafusion_functions::datetime::expr_fn::now();
        let gte_right = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left.clone()),
            DFOperator::GtEq,
            Box::new(right),
        ));
        let lte_now = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::LtEq,
            Box::new(now),
        ));
        Ok(DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(gte_right),
            DFOperator::And,
            Box::new(lte_now),
        )))
    });

    // timestamp NOT IN timestamp -> NOT (left >= right AND left <= now())
    // Inclusive end: user wrote `NOT IN timestamp`, not `..` syntax.
    registry.register::<NotInTimestampTimestamp>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;
        let now = datafusion_functions::datetime::expr_fn::now();
        let gte_right = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left.clone()),
            DFOperator::GtEq,
            Box::new(right),
        ));
        let lte_now = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(left),
            DFOperator::LtEq,
            Box::new(now),
        ));
        let in_range = DFExpr::BinaryExpr(BinaryExpr::new(
            Box::new(gte_right),
            DFOperator::And,
            Box::new(lte_now),
        ));
        Ok(DFExpr::Not(Box::new(in_range)))
    });

    // left IN (a, b, c) -> left IN [a, b, c]
    // The tuple is translated as struct(a, b, c), so we extract the elements
    registry.register::<InTuple>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;

        // Extract tuple elements from the struct() expression
        let list = try_extract_tuple_elements(right).ok_or_else(|| {
            anyhow::anyhow!(
                "IN tuple: expected struct() expression for tuple, got a different expression type"
            )
        })?;

        Ok(DFExpr::InList(InList::new(Box::new(left), list, false)))
    });

    // left NOT IN (a, b, c) -> left NOT IN [a, b, c]
    // The tuple is translated as struct(a, b, c), so we extract the elements
    registry.register::<NotInTuple>(|mut params| {
        let left = params.take()?.expr;
        let right = params.take()?.expr;

        // Extract tuple elements from the struct() expression
        let list = try_extract_tuple_elements(right).ok_or_else(|| {
            anyhow::anyhow!(
                "NOT IN tuple: expected struct() expression for tuple, got a different expression type"
            )
        })?;

        Ok(DFExpr::InList(InList::new(Box::new(left), list, true)))
    });
}
