//! DataFusion translations for map functions.

use datafusion::logical_expr::expr::Case as DFCase;
use datafusion::logical_expr::{lit, Expr as DFExpr};
use datafusion_functions::core::expr_ext::FieldAccessor;
use datafusion_functions_nested::expr_fn as nested_fn;
use datafusion_functions_nested::map::map_udf;

use hamelin_lib::func::defs::{
    GetMap, MapEmpty, MapFromArrays, MapFromKeyValue, MapFromPairs, MapKeys, MapValues,
    MapVariantGetValues, MapVariantToJsonValues, TransformValues,
};

use super::DataFusionTranslationRegistry;
use crate::udf::{map_extract_udf, map_variant_get_values_udf, map_variant_to_json_values_udf};

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // map() -> map_udf().call([])
    // Empty map - call the UDF with empty key/value arrays
    registry.register::<MapEmpty>(|_params| {
        // Create empty arrays for keys and values
        let empty_keys = datafusion_functions_nested::expr_fn::make_array(vec![]);
        let empty_values = datafusion_functions_nested::expr_fn::make_array(vec![]);
        Ok(map_udf().call(vec![empty_keys, empty_values]))
    });

    // map(keys_array, values_array) ->
    //   CASE WHEN keys IS NOT NULL AND values IS NOT NULL
    //             AND array_length(keys) = array_length(values)
    //        THEN map(keys, values)
    //   END
    //
    // Guards against:
    // - Null arrays (scalar-vs-array mismatch after constant folding)
    // - Mismatched lengths (DataFusion's map() panics instead of returning null)
    registry.register::<MapFromArrays>(|mut params| {
        let keys = params.take()?.expr;
        let values = params.take()?.expr;
        let not_null = keys.clone().is_not_null().and(values.clone().is_not_null());
        let array_len = crate::udf::hamelin_array_length_udf();
        let same_len = array_len
            .call(vec![keys.clone()])
            .eq(array_len.call(vec![values.clone()]));
        let condition = not_null.and(same_len);
        let map_expr = map_udf().call(vec![keys, values]);
        Ok(DFExpr::Case(DFCase {
            expr: None,
            when_then_expr: vec![(Box::new(condition), Box::new(map_expr))],
            else_expr: Some(Box::new(lit(datafusion::common::ScalarValue::Null))),
        }))
    });

    // map(array_of_pairs) -> hamelin_map_from_entries(array_of_pairs)
    registry.register::<MapFromPairs>(|mut params| {
        let pairs_array = params.take()?.expr;
        Ok(crate::udf::map_from_entries_udf().call(vec![pairs_array]))
    });

    // map(pair1, pair2, ...) -> map([k1, k2, ...], [v1, v2, ...])
    // Each pair is a tuple/row - need to build key and value arrays
    registry.register::<MapFromKeyValue>(|mut params| {
        // Collect all pairs and build separate key/value arrays
        let mut keys = vec![];
        let mut values = vec![];

        while let Ok(pair) = params.take() {
            // Extract first and second elements from the row/tuple
            // In DataFusion, struct field access uses get_field
            let pair_expr = pair.expr;
            keys.push(pair_expr.clone().field("f0"));
            values.push(pair_expr.field("f1"));
        }

        let keys_array = datafusion_functions_nested::expr_fn::make_array(keys);
        let values_array = datafusion_functions_nested::expr_fn::make_array(values);
        Ok(map_udf().call(vec![keys_array, values_array]))
    });

    // get(map, key) -> hamelin_map_extract(map, key)
    //
    // Hamelin's MapKeyMatcher restricts keys to byte-comparable types, so the
    // generic map_extract UDF's MutableArrayData path is never needed. This
    // UDF uses Arrow's type-specialized eq + take kernels and returns the
    // value directly (no list-wrap that array_element would have to undo).
    // See hamelin_datafusion/src/udf/map_extract.rs for details.
    registry.register::<GetMap>(|mut params| {
        let map_expr = params.take()?.expr;
        let key = params.take()?.expr;
        Ok(map_extract_udf().call(vec![map_expr, key]))
    });

    // map_keys(map) -> map_keys(map)
    registry.register::<MapKeys>(|mut params| {
        let map_expr = params.take()?.expr;
        Ok(nested_fn::map_keys(map_expr))
    });

    // map_values(map) -> map_values(map)
    registry.register::<MapValues>(|mut params| {
        let map_expr = params.take()?.expr;
        Ok(nested_fn::map_values(map_expr))
    });

    // transform_values should never reach the DataFusion translator: the
    // `rewrite_transform_values` normalizer runs before translation and
    // either fast-paths it to a vectorized UDF or unfolds it to
    // `map(map_keys(m), transform(map_values(m), λ))`. If one slips through,
    // the pipeline is misconfigured — surface a clear error rather than the
    // generic "No translation registered" so whoever's debugging knows where
    // to look.
    registry.register::<TransformValues>(|_params| {
        anyhow::bail!(
            "transform_values reached the DataFusion translator — this should \
             have been rewritten by the `rewrite_transform_values` pass. \
             Check that the pass is wired into normalize_pipeline before \
             translation and that ctx.lower_transform is set.",
        )
    });

    // map_variant_get_values(map, path) -> hamelin_map_variant_get_values(map, path)
    registry.register::<MapVariantGetValues>(|mut params| {
        let map_expr = params.take_by_name("map")?.expr;
        let path = params.take_by_name("path")?.expr;
        Ok(map_variant_get_values_udf().call(vec![map_expr, path]))
    });

    // map_variant_to_json_values(map) -> hamelin_map_variant_to_json_values(map)
    registry.register::<MapVariantToJsonValues>(|mut params| {
        let map_expr = params.take_by_name("map")?.expr;
        Ok(map_variant_to_json_values_udf().call(vec![map_expr]))
    });
}
