//! SQL translations for conditional functions

use hamelin_lib::func::defs::{ArrayCoalesce, Case, Coalesce, If2, If3};
use hamelin_lib::sql::expression::apply::{FunctionCallApply, Lambda};
use hamelin_lib::sql::expression::identifier::SimpleIdentifier;
use hamelin_lib::sql::expression::literal::{ColumnReference, IntegerLiteral};
use hamelin_lib::sql::expression::IndexLookup;
use hamelin_lib::sql::expression::SQLExpression;

use crate::utils::direct_function_translation;
use crate::TranslationRegistry;

/// Register all conditional function translations.
pub fn register(registry: &mut TranslationRegistry) {
    // If functions use direct pass-through
    registry.register::<If2>(direct_function_translation);
    registry.register::<If3>(direct_function_translation);

    // Case has custom translation
    registry.register::<Case>(|_, bindings| {
        let mut cases = vec![];
        for v in bindings.into_iter() {
            match v.sql {
                SQLExpression::RowLiteral(rl) if rl.values.len() == 2 => {
                    cases.push((rl.values[0].clone(), rl.values[1].clone()));
                }
                v => {
                    cases.push((
                        IndexLookup::new(v.clone(), IntegerLiteral::new("1").into()).into(),
                        IndexLookup::new(v.clone(), IntegerLiteral::new("2").into()).into(),
                    ));
                }
            }
        }
        Ok(hamelin_lib::sql::expression::Case::new(&cases[..], None).into())
    });

    // Coalesce uses direct pass-through
    registry.register::<Coalesce>(direct_function_translation);

    // array_coalesce(array<T>, T) -> transform(arr, x -> coalesce(x, default))
    // Trino/DuckDB support lambdas natively; element-wise coalesce is just a
    // transform whose body is the scalar coalesce. The DataFusion backend
    // overrides this with a vectorized UDF instead.
    registry.register::<ArrayCoalesce>(|_, mut bindings| {
        let array = bindings.take()?.sql;
        let default = bindings.take()?.sql;
        let lambda_param = SimpleIdentifier::new("x");
        let param_ref: SQLExpression = ColumnReference::new(lambda_param.clone().into()).into();
        let body: SQLExpression =
            FunctionCallApply::with_two("coalesce", param_ref, default).into();
        let lambda: SQLExpression = Lambda::new(vec![lambda_param], body).into();
        Ok(FunctionCallApply::with_two("transform", array, lambda).into())
    });
}
