//! Vectorized element-wise coalesce over a list array with a scalar default.
//!
//! Implements `coalesce(array<T>, T) -> array<T>`: for each element of the
//! input list, return the element if non-null, else the scalar default. The
//! outer list's offsets and null bitmap are preserved verbatim — null rows
//! stay null, empty rows stay empty.
//!
//! # Why this exists
//!
//! Without this UDF, `coalesce(arr, default)` over an array column has no
//! direct vectorized form: arrow's scalar `coalesce` operates row-wise on
//! same-length array columns, not on the inner elements of a list. Hamelin
//! used to broadcast `coalesce(arr, lit)` into `transform(arr, x -> coalesce(x, lit))`
//! and lower the transform via explode/agg, which incurs a UUID-keyed
//! reaggregate. Routing through this UDF skips all of that.
//!
//! Operates on the list's flat child values: a single `arrow_select::zip::zip`
//! call selects between the child and the broadcast scalar default based on
//! a per-element non-null mask, then re-wraps with the original offsets.

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

use datafusion::arrow::array::{
    Array, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait, Scalar,
};
use datafusion::arrow::compute::{is_not_null, kernels::zip::zip};
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 HamelinArrayCoalesceDefaultUdf {
    signature: Signature,
}

impl Default for HamelinArrayCoalesceDefaultUdf {
    fn default() -> Self {
        Self {
            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
        }
    }
}

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

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        match arg_types.first() {
            Some(dt @ DataType::List(_)) => Ok(dt.clone()),
            Some(dt @ DataType::LargeList(_)) => Ok(dt.clone()),
            other => exec_err!(
                "hamelin_array_coalesce_default: expected List or LargeList as first argument, got {:?}",
                other
            ),
        }
    }

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

        let default_scalar = match &args[1] {
            ColumnarValue::Scalar(s) => s.clone(),
            ColumnarValue::Array(_) => {
                return exec_err!(
                    "hamelin_array_coalesce_default: default value must be a scalar, got an array"
                );
            }
        };

        match &args[0] {
            ColumnarValue::Array(array) => match array.data_type() {
                DataType::List(_) => {
                    let result = apply_coalesce_to_list(array.as_list::<i32>(), &default_scalar)?;
                    Ok(ColumnarValue::Array(result))
                }
                DataType::LargeList(_) => {
                    let result = apply_coalesce_to_list(array.as_list::<i64>(), &default_scalar)?;
                    Ok(ColumnarValue::Array(result))
                }
                dt => exec_err!(
                    "hamelin_array_coalesce_default expects List or LargeList, got {:?}",
                    dt
                ),
            },
            ColumnarValue::Scalar(scalar) => scalar_array_coalesce(scalar, &default_scalar),
        }
    }
}

/// Apply element-wise coalesce to the flat child values of a ListArray, then re-wrap.
///
/// Builds a non-null mask over the children and uses `zip` to substitute the
/// scalar default at null positions. The list's offsets and null buffer are
/// preserved unchanged — `zip` is applied at the child level, so null/empty
/// rows of the outer list never participate.
fn apply_coalesce_to_list<O: OffsetSizeTrait>(
    list_array: &GenericListArray<O>,
    default_scalar: &ScalarValue,
) -> Result<ArrayRef> {
    let child_values = list_array.values();
    let child_len = child_values.len();

    // The default's Arrow type must match the child element type. The Hamelin
    // typechecker enforces this, but verify here so a malformed plan doesn't
    // produce silently wrong data.
    let default_arr = default_scalar.to_array_of_size(1)?;
    if default_arr.data_type() != child_values.data_type() {
        return exec_err!(
            "hamelin_array_coalesce_default: default scalar type {:?} does not match list element type {:?}",
            default_arr.data_type(),
            child_values.data_type()
        );
    }

    let new_child = if child_len == 0 {
        // zip rejects mask/array length mismatches; short-circuit to avoid
        // building empty intermediates.
        Arc::clone(child_values)
    } else {
        let mask = is_not_null(child_values.as_ref())?;
        let default_datum = Scalar::new(default_arr);
        zip(&mask, child_values, &default_datum)?
    };

    // Reuse the original list's field, offsets, and null buffer. The element
    // type is unchanged, so the field carries over verbatim.
    let field = match list_array.data_type() {
        DataType::List(f) | DataType::LargeList(f) => f.clone(),
        other => return exec_err!("expected List or LargeList, got {:?}", other),
    };

    let result_list = GenericListArray::<O>::try_new(
        field,
        list_array.offsets().clone(),
        new_child,
        list_array.nulls().cloned(),
    )?;
    Ok(Arc::new(result_list) as ArrayRef)
}

/// Handle scalar list inputs by extracting the inner array, processing, and re-wrapping.
fn scalar_array_coalesce(
    scalar: &ScalarValue,
    default_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_array_coalesce_default expects List type, got {:?}",
                scalar
            );
        }
    };

    if arr.is_empty() || arr.is_null(0) {
        return Ok(ColumnarValue::Scalar(scalar.clone()));
    }

    let result = if is_large {
        apply_coalesce_to_list(arr.as_list::<i64>(), default_scalar)?
    } else {
        apply_coalesce_to_list(arr.as_list::<i32>(), default_scalar)?
    };

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

pub fn array_coalesce_default_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinArrayCoalesceDefaultUdf::default())
}
