//! Fast dimension-1 array length UDF.
//!
//! # Why this exists
//!
//! DataFusion's built-in `array_length` (in `datafusion-functions-nested`) computes
//! element counts by iterating over the `ListArray` via `.iter()`, which calls
//! `.value(i)` on each row. For a `ListArray` whose children are deeply nested
//! structs (e.g. `List<Struct<BinaryView, BinaryView>>` from variant/OTEL data),
//! each `.value(i)` call **slices** the child `StructArray`, which recursively
//! slices every child column — cloning `Arc`s for every buffer in every
//! `BinaryViewArray`. The sliced view is then immediately dropped, walking the
//! entire `Arc` tree again.
//!
//! Profiling an OTEL log ingestion query showed **85% of execution time** in
//! `ArrayLength::invoke_with_args`:
//!   - 48% in `StructArray::slice` → `GenericByteViewArray::slice` (Arc clones)
//!   - 37% in `Arc::drop_slow` (dropping the temporary sliced views)
//!
//! But a `ListArray` already stores an offsets buffer! The length of element `i`
//! is simply `offsets[i+1] - offsets[i]` — a single integer subtraction with
//! **zero allocation**. This UDF reads offsets directly, making the operation
//! O(N) with no Arc traffic, regardless of how deeply nested the child arrays are.
//!
//! # Scope
//!
//! This UDF only computes **dimension-1** length (top-level element count), which
//! is the only dimension Hamelin ever uses. It does not support the optional
//! `dimension` parameter that DataFusion's built-in offers.

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

use datafusion::arrow::array::{
    Array, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait, UInt64Array,
};
use datafusion::arrow::datatypes::DataType;
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
    Volatility,
};

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

impl Default for HamelinArrayLengthUdf {
    fn default() -> Self {
        Self::new()
    }
}

impl HamelinArrayLengthUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
        }
    }
}

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

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

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

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::UInt64)
    }

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

        match &args[0] {
            ColumnarValue::Scalar(scalar) => scalar_array_length(scalar),
            ColumnarValue::Array(array) => match array.data_type() {
                DataType::List(_) => Ok(ColumnarValue::Array(lengths_from_offsets(
                    array.as_list::<i32>(),
                ))),
                DataType::LargeList(_) => Ok(ColumnarValue::Array(lengths_from_offsets(
                    array.as_list::<i64>(),
                ))),
                dt => exec_err!(
                    "hamelin_array_length expects List or LargeList, got {:?}",
                    dt
                ),
            },
        }
    }
}

/// Compute array lengths directly from the list offsets buffer.
///
/// For a `GenericListArray` with offsets `[o0, o1, o2, ..., oN]`,
/// the length of element `i` is `offsets[i+1] - offsets[i]`.
/// Null rows produce null output.
fn lengths_from_offsets<O: OffsetSizeTrait>(list_array: &GenericListArray<O>) -> ArrayRef {
    let offsets = list_array.value_offsets();
    let nulls = list_array.nulls();

    let lengths: UInt64Array = match nulls {
        Some(null_buffer) => offsets
            .windows(2)
            .enumerate()
            .map(|(i, w)| {
                null_buffer
                    .is_valid(i)
                    .then(|| (w[1] - w[0]).as_usize() as u64)
            })
            .collect(),
        None => offsets
            .windows(2)
            .map(|w| Some((w[1] - w[0]).as_usize() as u64))
            .collect(),
    };

    Arc::new(lengths)
}

/// Handle scalar (single-value) list inputs.
///
/// The scalar path still calls `.value(0)` which does a slice, but this is
/// acceptable — it processes exactly one row and is not a hot path.
fn scalar_array_length(scalar: &ScalarValue) -> Result<ColumnarValue> {
    let arr: &dyn Array = match scalar {
        ScalarValue::List(arr) => arr.as_ref(),
        ScalarValue::LargeList(arr) => arr.as_ref(),
        _ => return exec_err!("hamelin_array_length expects List type, got {:?}", scalar),
    };
    if arr.is_empty() || arr.is_null(0) {
        Ok(ColumnarValue::Scalar(ScalarValue::UInt64(None)))
    } else {
        // For a scalar list, the single-element .value(0) slice is acceptable.
        let length = match arr.data_type() {
            DataType::List(_) => arr.as_list::<i32>().value(0).len(),
            DataType::LargeList(_) => arr.as_list::<i64>().value(0).len(),
            dt => return exec_err!("hamelin_array_length expects List type, got {:?}", dt),
        };
        Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(
            length as u64,
        ))))
    }
}

pub fn hamelin_array_length_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinArrayLengthUdf::new())
}
