//! DataFusion translations for unary operators.

use datafusion::common::ScalarValue;
use datafusion::logical_expr::Expr as DFExpr;
use datafusion_functions::core::expr_fn as core_fn;

use hamelin_lib::func::defs::{
    UnaryMinus, UnaryPlus, UnaryRangePostfix, UnaryRangePrefix, UnaryRangePrefixInclusive,
};

use super::DataFusionTranslationRegistry;

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // -x -> NEGATIVE(x)
    registry.register::<UnaryMinus>(|mut params| {
        let x = params.take()?.expr;
        Ok(DFExpr::Negative(Box::new(x)))
    });

    // +x -> x (no-op)
    registry.register::<UnaryPlus>(|mut params| Ok(params.take()?.expr));

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

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

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