//! DataFusion translations for map functions.

use datafusion_functions::core::expr_ext::FieldAccessor;
use datafusion_functions_nested::expr_fn as nested_fn;

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_from_arrays_udf, map_variant_get_values_udf,
    map_variant_to_json_values_udf,
};

pub fn register(registry: &mut DataFusionTranslationRegistry) {
    // map() -> hamelin_map_from_arrays([], [])
    registry.register::<MapEmpty>(|_params| {
        let empty_keys = datafusion_functions_nested::expr_fn::make_array(vec![]);
        let empty_values = datafusion_functions_nested::expr_fn::make_array(vec![]);
        Ok(map_from_arrays_udf().call(vec![empty_keys, empty_values]))
    });

    // map(keys_array, values_array) -> hamelin_map_from_arrays(keys_array, values_array)
    registry.register::<MapFromArrays>(|mut params| {
        let keys = params.take()?.expr;
        let values = params.take()?.expr;
        Ok(map_from_arrays_udf().call(vec![keys, values]))
    });

    // 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_from_arrays_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]))
    });
}
