//! Vectorized `map[key]` extraction (`hamelin_map_extract`).
//!
//! Drop-in replacement for the `map_extract(...)[1]` lowering Hamelin uses for
//! the `map[key]` / `get(map, key)` syntax. Returns the value associated with
//! `key` for each row, or null when the key is absent.
//!
//! # Why this exists
//!
//! DataFusion's upstream `map_extract` UDF (in `datafusion-functions-nested`)
//! is implemented generically with `MutableArrayData`:
//!
//!   * Preallocates an output buffer sized for the entire input map's value
//!     array — `total_keys × value_element_size` — then writes only `num_rows`
//!     entries. For prom-metrics-shaped batches (~33K rows × ~5 entries each ×
//!     16 B `Utf8View`), that's a multi-MB allocation per call, which on macOS
//!     goes through `_xzm_malloc_large_huge` → mmap → high sys time.
//!   * Per-key probe is `keys.slice(start + i, 1).as_ref() == query.as_ref()`
//!     — an Arc-bumped, vtable-dispatched dyn comparison per element, ~500K
//!     such comparisons per batch.
//!
//! Profiling on a 24.8 MB perpetualmetrics ingest showed 99% of CPU below
//! `MapExtract::invoke_with_args` and 49% wall time in kernel sys.
//!
//! # The fast path
//!
//! Hamelin's `MapKeyMatcher` (see `hamelin_lib/src/types/matcher.rs`)
//! restricts map keys to byte-comparable / SIMD-eligible types: `Binary`,
//! `Boolean`, `Interval`, `Int`, `String`, `Timestamp`, `Decimal`. Floats,
//! lists, structs, and nested types are rejected at type-check. Every map
//! Hamelin can construct is therefore safe to compare with Arrow's
//! type-specialized `eq` kernel.
//!
//! Algorithm:
//!
//!   1. Slice the map's flat keys/values to the visible window. `MapArray`
//!      slicing only narrows the outer offsets; `keys()` / `values()` still
//!      point at the full underlying buffers. Arrow's `Array::slice` is O(1),
//!      so this is free, and after it every downstream kernel operates in
//!      child-relative coordinates regardless of whether the input was sliced.
//!   2. Build a `BooleanArray` over the visible keys by `eq`-ing them against
//!      the query — a `Scalar`-broadcast for literal keys, or a row-aligned
//!      expansion of the column for non-literal keys.
//!   3. For each row, scan that bitmap within the row's (child-relative)
//!      offset window for the first true position — a tight, allocation-free
//!      loop.
//!   4. `take` the matching value once at the end — a single type-specialized
//!      gather over the visible values array.
//!
//! Allocation budget per call: a bitmap of `total_keys` bits + an output of
//! exactly `num_rows` × value-element-size + a `Vec<Option<u32>>` of
//! `num_rows`. None of these hit the multi-MB mmap path the generic
//! implementation forces.
//!
//! Returns the value type directly (not wrapped in a `List`), so the caller
//! does not need a follow-up `array_element` to scalarize.

use std::any::Any;
use std::sync::Arc;

use datafusion::arrow::array::{
    Array, ArrayRef, AsArray, BooleanArray, MapArray, Scalar, UInt32Array,
};
use datafusion::arrow::compute::{kernels::cmp::eq, take};
use datafusion::arrow::datatypes::DataType;
use datafusion::common::utils::take_function_args;
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};

#[derive(Debug, PartialEq, Eq, Hash)]
pub struct HamelinMapExtractUdf {
    signature: Signature,
}

impl Default for HamelinMapExtractUdf {
    fn default() -> Self {
        Self {
            signature: Signature::user_defined(Volatility::Immutable),
        }
    }
}

/// True for key types the byte-comparable fast path can handle. Mirrors
/// `MapKeyMatcher` (`hamelin_lib::types::matcher::MapKeyMatcher`): every type
/// Hamelin allows as a map key has a SIMD-friendly Arrow `eq` specialization,
/// and floats are excluded because their equality is unreliable. The list is
/// also a hard guard against complex-key inputs (`Struct`, `List`, `Map`,
/// nested) that would otherwise reach Arrow's `eq` kernel and produce an
/// opaque "not supported" error — codec round-trip and downstream embedders
/// can both bypass the Hamelin frontend, so we re-state the contract here.
fn is_byte_comparable_key_type(dt: &DataType) -> bool {
    use DataType::*;
    matches!(
        dt,
        Boolean
            | Int8
            | Int16
            | Int32
            | Int64
            | UInt8
            | UInt16
            | UInt32
            | UInt64
            | Decimal128(_, _)
            | Decimal256(_, _)
            | Timestamp(_, _)
            | Date32
            | Date64
            | Time32(_)
            | Time64(_)
            | Duration(_)
            | Interval(_)
            | Utf8
            | LargeUtf8
            | Utf8View
            | Binary
            | LargeBinary
            | BinaryView
            | FixedSizeBinary(_)
    )
}

impl ScalarUDFImpl for HamelinMapExtractUdf {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "hamelin_map_extract"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        let Some(DataType::Map(field, _)) = arg_types.first() else {
            return exec_err!(
                "hamelin_map_extract: first argument must be Map, got {:?}",
                arg_types.first()
            );
        };
        let DataType::Struct(entry_fields) = field.data_type() else {
            return exec_err!("hamelin_map_extract: expected Struct inside Map");
        };
        Ok(entry_fields[1].data_type().clone())
    }

    /// Coerce the query key to the map's declared key type. This pushes any
    /// Utf8/Utf8View/LargeUtf8 reconciliation into a planner-inserted cast,
    /// so the runtime always sees matching types and the byte-comparable
    /// fast path stays a fast path. Also rejects key types Arrow's `eq`
    /// kernel can't handle (structs, lists, nested maps, floats).
    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
        let [map_type, _] = take_function_args(self.name(), arg_types)?;
        let DataType::Map(field, _) = map_type else {
            return exec_err!("hamelin_map_extract: first argument must be Map, got {map_type:?}");
        };
        let DataType::Struct(entry_fields) = field.data_type() else {
            return exec_err!("hamelin_map_extract: expected Struct inside Map");
        };
        let key_type = entry_fields[0].data_type().clone();
        if !is_byte_comparable_key_type(&key_type) {
            return exec_err!(
                "hamelin_map_extract: key type {key_type:?} is not supported. \
                 Only byte-comparable types (Boolean, integers, Decimal, \
                 Timestamp/Date/Time, Duration, Interval, Utf8/Utf8View/LargeUtf8, \
                 Binary variants) can be used as map keys."
            );
        }
        Ok(vec![map_type.clone(), key_type])
    }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let args = args.args;
        if args.len() != 2 {
            return exec_err!(
                "hamelin_map_extract expects 2 arguments, got {}",
                args.len()
            );
        }

        // Materialize the map. Scalar map inputs go through to_array_of_size(1)
        // and we extract index 0 of the result before returning.
        let (map_arr, scalar_map): (ArrayRef, bool) = match &args[0] {
            ColumnarValue::Array(a) => (Arc::clone(a), false),
            ColumnarValue::Scalar(ScalarValue::Map(m)) => (Arc::clone(m) as ArrayRef, true),
            ColumnarValue::Scalar(s) if s.is_null() => {
                let dt = self.return_type(&[args[0].data_type(), args[1].data_type()])?;
                return Ok(ColumnarValue::Scalar(ScalarValue::try_from(&dt)?));
            }
            ColumnarValue::Scalar(other) => {
                return exec_err!(
                    "hamelin_map_extract: first argument must be a Map, got {:?}",
                    other.data_type()
                );
            }
        };

        let DataType::Map(_, _) = map_arr.data_type() else {
            return exec_err!(
                "hamelin_map_extract: first argument must be a Map, got {:?}",
                map_arr.data_type()
            );
        };
        let map = map_arr.as_map();

        let result = match &args[1] {
            ColumnarValue::Scalar(key) => extract_with_scalar_key(map, key)?,
            ColumnarValue::Array(keys) => {
                if keys.len() != map.len() {
                    return exec_err!(
                        "hamelin_map_extract: key column length {} does not match map column length {}",
                        keys.len(),
                        map.len()
                    );
                }
                extract_with_array_key(map, keys)?
            }
        };

        if scalar_map {
            Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
                &result, 0,
            )?))
        } else {
            Ok(ColumnarValue::Array(result))
        }
    }
}

/// Narrow the map's flat keys/values to just the visible window, returning the
/// child-relative `base` offset that must be subtracted from absolute
/// `value_offsets()` entries.
///
/// `MapArray::slice` only narrows the outer offsets buffer; `keys()` and
/// `values()` still return the full underlying flat arrays, and
/// `value_offsets()` returns absolute indices into them. Slicing the children
/// once up front (Arrow's `Array::slice` is O(1) — bumps offset+len, no copy)
/// lets every downstream kernel operate in child-relative coordinates without
/// caring whether the input was sliced.
fn visible_children(map: &MapArray) -> (ArrayRef, ArrayRef, usize) {
    // A valid `MapArray` has `value_offsets().len() == map.len() + 1`, so both
    // indexes are guaranteed in-bounds by Arrow's MapArray invariant.
    let offsets = map.value_offsets();
    let n = map.len();
    let base = offsets[0] as usize;
    let total = offsets[n] as usize;
    let len = total - base;
    (
        map.keys().slice(base, len),
        map.values().slice(base, len),
        base,
    )
}

fn extract_with_scalar_key(map: &MapArray, key: &ScalarValue) -> Result<ArrayRef> {
    let key_arr = key.to_array_of_size(1)?;
    if key_arr.data_type() != map.key_type() {
        return exec_err!(
            "hamelin_map_extract: key type {:?} does not match map key type {:?}",
            key_arr.data_type(),
            map.key_type()
        );
    }
    let (visible_keys, visible_values, base) = visible_children(map);
    let scalar = Scalar::new(key_arr);
    let key_match = eq(&visible_keys, &scalar)?;
    gather_first_match(map, base, &visible_values, &key_match)
}

fn extract_with_array_key(map: &MapArray, keys: &ArrayRef) -> Result<ArrayRef> {
    if keys.data_type() != map.key_type() {
        return exec_err!(
            "hamelin_map_extract: key type {:?} does not match map key type {:?}",
            keys.data_type(),
            map.key_type()
        );
    }

    let (visible_keys, visible_values, base) = visible_children(map);

    // Expand the per-row query keys to align with the visible flat keys array,
    // then `eq` once across the whole batch. Using `resize` (vs. a per-entry
    // push loop) lets LLVM lower the fill to a memset for `u32`.
    let offsets = map.value_offsets();
    let mut row_indices: Vec<u32> = Vec::with_capacity(visible_keys.len());
    for row in 0..map.len() {
        let len = (offsets[row + 1] - offsets[row]) as usize;
        row_indices.resize(row_indices.len() + len, row as u32);
    }
    let row_index_arr = UInt32Array::from(row_indices);
    let expanded_keys = take(keys.as_ref(), &row_index_arr, None)?;
    let key_match = eq(&visible_keys, &expanded_keys)?;
    gather_first_match(map, base, &visible_values, &key_match)
}

/// For each row, find the first index in the visible-window keys array where
/// `key_match` is true and gather the corresponding value. `base` is the
/// child-relative offset of the visible window — subtract it from
/// `value_offsets()` entries to get indices into `visible_values` / `key_match`.
fn gather_first_match(
    map: &MapArray,
    base: usize,
    visible_values: &ArrayRef,
    key_match: &BooleanArray,
) -> Result<ArrayRef> {
    let num_rows = map.len();
    let mut indices: Vec<Option<u32>> = vec![None; num_rows];
    for (row, win) in map.value_offsets().windows(2).enumerate() {
        if map.is_null(row) {
            continue;
        }
        let start = win[0] as usize - base;
        let end = win[1] as usize - base;
        for j in start..end {
            if key_match.value(j) && !key_match.is_null(j) {
                indices[row] = Some(j as u32);
                break;
            }
        }
    }
    let indices_arr = UInt32Array::from(indices);
    let values = take(visible_values.as_ref(), &indices_arr, None)?;
    Ok(values)
}

pub fn map_extract_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinMapExtractUdf::default())
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::array::{Int64Array, StringArray, StructArray};
    use datafusion::arrow::buffer::OffsetBuffer;
    use datafusion::arrow::datatypes::{Field, Fields};
    use datafusion::logical_expr::ScalarUDFImpl;

    /// Build a `MapArray<Utf8, Int64>` from per-row slices of `(key, value)`
    /// pairs. Rows of `None` produce a null map row. The MapArray is the only
    /// way to introduce duplicate keys or directly-sliced inputs into the UDF
    /// — Hamelin's surface syntax validates both.
    fn build_map_string_int(rows: &[Option<&[(&str, i64)]>]) -> MapArray {
        let mut keys: Vec<&str> = Vec::new();
        let mut values: Vec<i64> = Vec::new();
        let mut lengths: Vec<usize> = Vec::with_capacity(rows.len());
        let mut row_validity: Vec<bool> = Vec::with_capacity(rows.len());

        for row in rows {
            match row {
                Some(entries) => {
                    for (k, v) in *entries {
                        keys.push(k);
                        values.push(*v);
                    }
                    lengths.push(entries.len());
                    row_validity.push(true);
                }
                None => {
                    lengths.push(0);
                    row_validity.push(false);
                }
            }
        }

        let entry_fields = Fields::from(vec![
            Field::new("key", DataType::Utf8, false),
            Field::new("value", DataType::Int64, true),
        ]);
        let entries = StructArray::new(
            entry_fields.clone(),
            vec![
                Arc::new(StringArray::from(keys)) as ArrayRef,
                Arc::new(Int64Array::from(values)) as ArrayRef,
            ],
            None,
        );
        let nulls = datafusion::arrow::buffer::NullBuffer::from(row_validity);
        MapArray::new(
            Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)),
            OffsetBuffer::from_lengths(lengths),
            entries,
            Some(nulls),
            false,
        )
    }

    fn invoke_with_scalar_string_key(map: MapArray, key: &str) -> ArrayRef {
        let udf = HamelinMapExtractUdf::default();
        let result = udf
            .invoke_with_args(ScalarFunctionArgs {
                args: vec![
                    ColumnarValue::Array(Arc::new(map.clone())),
                    ColumnarValue::Scalar(ScalarValue::Utf8(Some(key.to_string()))),
                ],
                arg_fields: vec![],
                number_rows: map.len(),
                return_field: Arc::new(Field::new("v", DataType::Int64, true)),
                config_options: Arc::new(datafusion::config::ConfigOptions::default()),
            })
            .expect("invoke");
        match result {
            ColumnarValue::Array(a) => a,
            ColumnarValue::Scalar(s) => s.to_array_of_size(map.len()).expect("to array"),
        }
    }

    fn assert_int64_eq(arr: &ArrayRef, expected: &[Option<i64>]) {
        let arr = arr
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("Int64Array");
        assert_eq!(arr.len(), expected.len(), "length mismatch");
        for (i, want) in expected.iter().enumerate() {
            match want {
                Some(v) => {
                    assert!(!arr.is_null(i), "row {i} unexpectedly null");
                    assert_eq!(arr.value(i), *v, "row {i}");
                }
                None => assert!(arr.is_null(i), "row {i} should be null"),
            }
        }
    }

    /// Hamelin's `map(...)` constructors validate uniqueness, so duplicate
    /// keys can only enter via external sources (parquet ingest, raw arrow).
    /// The `gather_first_match` loop has an explicit "first match wins"
    /// contract — if a refactor changes that to e.g. last-match-wins, this
    /// test catches it. Builds the `MapArray` directly to bypass Hamelin's
    /// uniqueness check.
    #[test]
    fn duplicate_keys_first_match_wins() {
        // One row with three entries: 'a' → 1, 'a' → 2, 'b' → 3.
        let map = build_map_string_int(&[Some(&[("a", 1), ("a", 2), ("b", 3)])]);
        let result = invoke_with_scalar_string_key(map, "a");
        assert_int64_eq(&result, &[Some(1)]);
    }

    /// Regression test: `MapArray::slice` produces a row-window over the same
    /// underlying flat keys/values, with `value_offsets()[0]` non-zero. The
    /// UDF must operate in the visible window only — without `visible_children`,
    /// `eq(map.keys(), ...)` errors on length mismatch and `gather_first_match`
    /// uses out-of-bounds absolute indices into a shorter `key_match`. Mirrors
    /// the analogous regression in `filter_null::tests`.
    #[test]
    fn handles_sliced_map_array_scalar_key() {
        // Five rows: [{a:1,b:2}, {a:3}, {a:10,b:20,c:30}, {a:40}, {b:50}].
        // Slice to rows 2..5 — `value_offsets()[0]` becomes 3 (sum of pre-slice
        // entries), `keys()` and `values()` still see all 8 underlying entries.
        let full = build_map_string_int(&[
            Some(&[("a", 1), ("b", 2)]),
            Some(&[("a", 3)]),
            Some(&[("a", 10), ("b", 20), ("c", 30)]),
            Some(&[("a", 40)]),
            Some(&[("b", 50)]),
        ]);
        let sliced = full.slice(2, 3);
        assert_ne!(
            sliced.value_offsets()[0],
            0,
            "test setup: slice's first offset must be non-zero"
        );

        let result = invoke_with_scalar_string_key(sliced, "a");
        assert_int64_eq(&result, &[Some(10), Some(40), None]);
    }

    /// Same as above but for the non-literal-key path. Both branches use the
    /// same `visible_children` helper, so this guards against either drifting
    /// independently.
    #[test]
    fn handles_sliced_map_array_array_key() {
        let full = build_map_string_int(&[
            Some(&[("a", 1), ("b", 2)]),
            Some(&[("a", 3)]),
            Some(&[("a", 10), ("b", 20), ("c", 30)]),
            Some(&[("a", 40)]),
            Some(&[("b", 50), ("c", 60)]),
        ]);
        let sliced = full.slice(2, 3);
        let keys: ArrayRef = Arc::new(StringArray::from(vec!["b", "a", "c"]));

        let udf = HamelinMapExtractUdf::default();
        let result = udf
            .invoke_with_args(ScalarFunctionArgs {
                args: vec![
                    ColumnarValue::Array(Arc::new(sliced.clone())),
                    ColumnarValue::Array(Arc::clone(&keys)),
                ],
                arg_fields: vec![],
                number_rows: sliced.len(),
                return_field: Arc::new(Field::new("v", DataType::Int64, true)),
                config_options: Arc::new(datafusion::config::ConfigOptions::default()),
            })
            .expect("invoke");
        let arr = match result {
            ColumnarValue::Array(a) => a,
            ColumnarValue::Scalar(s) => s.to_array_of_size(sliced.len()).expect("to array"),
        };
        assert_int64_eq(&arr, &[Some(20), Some(40), Some(60)]);
    }
}
