//! Drop NULL elements from each row of a list array.
//!
//! # Why this exists
//!
//! Hamelin's `filter_null(arr)` was previously translated to
//! `array_remove_all(arr, NULL)`. DataFusion 53 changed `array_remove_all`
//! to return NULL whenever the element-to-remove argument is NULL
//! (`datafusion-functions-nested-53.1.0/src/remove.rs:392-397`), so the
//! old translation now nulls out the whole row instead of dropping null
//! elements. This UDF gives Hamelin a proper null-element filter.
//!
//! Implementation: build a boolean mask over the list's child values
//! marking non-null positions, run arrow's `filter` kernel once over the
//! child to drop the null slots, then rebuild the list offsets by
//! counting how many positions survived in each row's original window.

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

use datafusion::arrow::array::{Array, ArrayRef, BooleanArray, GenericListArray, OffsetSizeTrait};
use datafusion::arrow::buffer::{BooleanBuffer, OffsetBuffer, ScalarBuffer};
use datafusion::arrow::compute::filter;
use datafusion::arrow::datatypes::{DataType, Field, FieldRef};
use datafusion::common::{exec_err, internal_err, Result, ScalarValue};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{
    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
    Volatility,
};

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

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

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

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

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

    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
        if arg_types.len() != 1 {
            return exec_err!(
                "hamelin_filter_null requires 1 argument, got {}",
                arg_types.len()
            );
        }
        Ok(arg_types.to_vec())
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        internal_err!("return_field_from_args should be called instead")
    }

    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
        if args.arg_fields.len() != 1 {
            return exec_err!(
                "hamelin_filter_null requires 1 argument, got {}",
                args.arg_fields.len()
            );
        }

        let input_field = &args.arg_fields[0];
        let dt = input_field.data_type();

        match dt {
            DataType::List(_) | DataType::LargeList(_) => Ok(Arc::new(Field::new(
                self.name(),
                dt.clone(),
                input_field.is_nullable(),
            ))),
            DataType::Null => Ok(Arc::new(Field::new(self.name(), DataType::Null, true))),
            other => exec_err!("hamelin_filter_null: expected List or LargeList, got {other}"),
        }
    }

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

        match &args[0] {
            ColumnarValue::Array(array) => match array.data_type() {
                DataType::List(_) => {
                    let result = filter_null_list::<i32>(array.as_ref())?;
                    Ok(ColumnarValue::Array(result))
                }
                DataType::LargeList(_) => {
                    let result = filter_null_list::<i64>(array.as_ref())?;
                    Ok(ColumnarValue::Array(result))
                }
                DataType::Null => Ok(ColumnarValue::Scalar(ScalarValue::Null)),
                dt => exec_err!(
                    "hamelin_filter_null: expected List or LargeList, got {:?}",
                    dt
                ),
            },
            ColumnarValue::Scalar(scalar) => scalar_filter_null(scalar),
        }
    }
}

/// Drop null child elements from each row of a list array.
///
/// One pass to compute the new offsets and select mask, then arrow's
/// `filter` kernel applies the mask to the child values in a single
/// vectorized operation.
fn filter_null_list<O: OffsetSizeTrait>(array: &dyn Array) -> Result<ArrayRef> {
    let list_array = array
        .as_any()
        .downcast_ref::<GenericListArray<O>>()
        .ok_or_else(|| {
            DataFusionError::Execution(format!(
                "hamelin_filter_null: expected GenericListArray, got {:?}",
                array.data_type()
            ))
        })?;

    // A sliced GenericListArray keeps the full child in `values()`; only
    // the offsets buffer and validity bitmap get sliced. Restrict the child
    // to the visible window [old_offsets[0]..old_offsets[len]) before
    // masking so the filter kernel — and the new offsets — only see
    // elements this slice actually owns.
    let old_offsets = list_array.offsets();
    let window_start = old_offsets[0].to_usize().unwrap();
    let window_end = old_offsets[list_array.len()].to_usize().unwrap();
    let window_len = window_end - window_start;
    let child = list_array.values().slice(window_start, window_len);

    // Build the per-element keep mask over the windowed child. true = keep,
    // false = drop. We use `logical_nulls` rather than `nulls` so that
    // NullArray (DataType::Null) — whose every element is implicitly null
    // but carries no explicit null buffer — is filtered out correctly.
    let keep_mask = match child.logical_nulls() {
        Some(nulls) => nulls.inner().clone(),
        None => BooleanBuffer::new_set(child.len()),
    };
    let mask_array = BooleanArray::new(keep_mask.clone(), None);

    let filtered_child = filter(child.as_ref(), &mask_array)
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;

    // Recompute offsets. For each row: new_len = count of true bits in
    // keep_mask over [old_offsets[i]..old_offsets[i+1]) translated into the
    // window's coordinate space. Null rows pass through unchanged (offsets
    // repeat, the list's own null bitmap marks the row).
    let mut new_offsets: Vec<O> = Vec::with_capacity(list_array.len() + 1);
    let mut cursor: usize = 0;
    new_offsets.push(O::usize_as(cursor));

    for row in 0..list_array.len() {
        let start = old_offsets[row].to_usize().unwrap() - window_start;
        let end = old_offsets[row + 1].to_usize().unwrap() - window_start;
        // count_set_bits over [start..end) — Arrow's BooleanBuffer doesn't
        // expose a slice-and-count helper, so iterate directly.
        let kept = (start..end).filter(|j| keep_mask.value(*j)).count();
        cursor += kept;
        new_offsets.push(O::usize_as(cursor));
    }

    let new_offsets = OffsetBuffer::new(ScalarBuffer::from(new_offsets));

    // Preserve the original item field exactly. Tightening its nullability
    // to false (since we just dropped every null slot) would diverge from
    // what `return_field_from_args` declared at planning time and trip
    // DataFusion's runtime schema assertion. The result's child array
    // simply has zero null entries — same field type, narrower bitmap.
    let item_field = match list_array.data_type() {
        DataType::List(item) | DataType::LargeList(item) => item.clone(),
        _ => unreachable!("filter_null_list called with non-list type"),
    };

    let result = GenericListArray::<O>::try_new(
        item_field,
        new_offsets,
        filtered_child,
        list_array.nulls().cloned(),
    )?;

    Ok(Arc::new(result) as ArrayRef)
}

/// Scalar form: dispatch to the array kernel after extracting the inner
/// `ListArray`/`LargeListArray`, then re-wrap as a scalar.
fn scalar_filter_null(scalar: &ScalarValue) -> Result<ColumnarValue> {
    let (arr, is_large): (&dyn Array, bool) = match scalar {
        ScalarValue::List(arr) => (arr.as_ref(), false),
        ScalarValue::LargeList(arr) => (arr.as_ref(), true),
        ScalarValue::Null => return Ok(ColumnarValue::Scalar(ScalarValue::Null)),
        _ => {
            return exec_err!("hamelin_filter_null expects List type, got {:?}", scalar);
        }
    };

    let result = if is_large {
        filter_null_list::<i64>(arr)?
    } else {
        filter_null_list::<i32>(arr)?
    };

    if arr.is_empty() {
        return ScalarValue::try_from(result.data_type()).map(ColumnarValue::Scalar);
    }

    ScalarValue::try_from_array(&result, 0).map(ColumnarValue::Scalar)
}

pub fn filter_null_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinFilterNullUdf::default())
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::array::{Int64Array, ListArray};

    /// Build a `ListArray` of nullable `Int64` from rows. `None` represents
    /// a null row; `Some(elems)` represents a non-null row whose elements
    /// may themselves be `Some(i64)` (non-null) or `None` (null element).
    fn make_list_i64(rows: &[Option<&[Option<i64>]>]) -> ListArray {
        let mut values: Vec<Option<i64>> = Vec::new();
        let mut offsets: Vec<i32> = vec![0];
        let mut row_validity: Vec<bool> = Vec::new();

        for row in rows {
            match row {
                Some(elems) => {
                    values.extend_from_slice(elems);
                    offsets.push(values.len() as i32);
                    row_validity.push(true);
                }
                None => {
                    offsets.push(values.len() as i32);
                    row_validity.push(false);
                }
            }
        }

        let values_array: ArrayRef = Arc::new(Int64Array::from(values));
        let field = Arc::new(Field::new("item", DataType::Int64, true));
        let nulls = datafusion::arrow::buffer::NullBuffer::from(row_validity);
        GenericListArray::<i32>::try_new(
            field,
            OffsetBuffer::new(offsets.into()),
            values_array,
            Some(nulls),
        )
        .expect("build ListArray")
    }

    /// Materialize a `ListArray<i32>` of `Int64` into a vector of optional
    /// rows of optional values for easy comparison.
    fn list_to_vec(arr: &ListArray) -> Vec<Option<Vec<Option<i64>>>> {
        let values = arr
            .values()
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("Int64 child");
        let offsets = arr.offsets();
        (0..arr.len())
            .map(|row| {
                if arr.is_null(row) {
                    None
                } else {
                    let start = offsets[row] as usize;
                    let end = offsets[row + 1] as usize;
                    Some(
                        (start..end)
                            .map(|i| {
                                if values.is_null(i) {
                                    None
                                } else {
                                    Some(values.value(i))
                                }
                            })
                            .collect(),
                    )
                }
            })
            .collect()
    }

    #[test]
    fn drops_null_elements_unsliced() {
        let input = make_list_i64(&[
            Some(&[Some(1), None, Some(2)]),
            None,
            Some(&[None, None]),
            Some(&[Some(3)]),
        ]);

        let result = filter_null_list::<i32>(&input).expect("filter_null_list");
        let result = result.as_any().downcast_ref::<ListArray>().expect("list");

        assert_eq!(
            list_to_vec(result),
            vec![
                Some(vec![Some(1), Some(2)]),
                None,
                Some(vec![]),
                Some(vec![Some(3)]),
            ]
        );
    }

    /// Regression test: a sliced `GenericListArray` keeps the full child in
    /// `values()` while its offsets remain absolute. Earlier the cursor was
    /// initialized to 0, so rows in the sliced output picked up elements
    /// from the prefix rows (outside the slice's visible window) that
    /// happened to survive the null-filter.
    #[test]
    fn drops_null_elements_when_input_is_sliced() {
        // Prefix rows (0..2) contain entirely non-null elements. With the
        // pre-fix cursor starting at 0, the filtered child includes all
        // five prefix elements and the first sliced row would mistakenly
        // reference [1, 2] instead of its own surviving values.
        let input = make_list_i64(&[
            Some(&[Some(1), Some(2)]),          // prefix row 0
            Some(&[Some(3), Some(4), Some(5)]), // prefix row 1
            Some(&[Some(10), None, Some(20)]),  // visible row 0
            Some(&[None, Some(30)]),            // visible row 1
            Some(&[None, None]),                // visible row 2
        ]);

        // `GenericListArray::slice` returns a new `ListArray` whose offsets
        // remain absolute into the unsliced child — exactly the shape that
        // tripped the original bug.
        let sliced = input.slice(2, 3);
        // Sanity-check the test setup: the slice's first offset is non-zero.
        assert_ne!(sliced.offsets()[0], 0);

        let result = filter_null_list::<i32>(&sliced).expect("filter_null_list");
        let result = result.as_any().downcast_ref::<ListArray>().expect("list");

        assert_eq!(
            list_to_vec(result),
            vec![
                Some(vec![Some(10), Some(20)]),
                Some(vec![Some(30)]),
                Some(vec![]),
            ]
        );
    }

    #[test]
    fn preserves_null_rows_when_input_is_sliced() {
        let input = make_list_i64(&[
            Some(&[Some(1)]),
            Some(&[Some(2), None]),
            None,
            Some(&[None, Some(3)]),
        ]);

        let sliced = input.slice(1, 3);
        let result = filter_null_list::<i32>(&sliced).expect("filter_null_list");
        let result = result.as_any().downcast_ref::<ListArray>().expect("list");

        assert_eq!(
            list_to_vec(result),
            vec![Some(vec![Some(2)]), None, Some(vec![Some(3)])]
        );
    }
}
