//! Vectorized struct-field extraction over list arrays.
//!
//! # Why this exists
//!
//! Hamelin represents `array_of_structs.field_name` structurally as a
//! `FieldAccess::BroadcastStructField`. DataFusion's built-in `get_field`
//! rejects `List<Struct>` at planning time — it only accepts Struct or
//! Map base types. Without this UDF we'd fall back to an EXPLODE →
//! per-row projection → `sliding_array_agg BY __row_id` pipeline, which
//! triggers a hash repartition on __row_id even though the row id is
//! partition-local.
//!
//! Arrow stores `ListArray<StructArray>` with each struct field as its
//! own contiguous child column, so `.field` is a pure metadata operation:
//! grab the named child via `StructArray::column_by_name`, re-wrap it
//! with the parent list's offsets and null bitmap. Zero explode, zero
//! repartition. The UDF handles both `List` and `LargeList`, and both
//! `ColumnarValue::Array` and `ColumnarValue::Scalar` inputs.
//!
//! Tuple elements share this UDF: Hamelin's tuple-as-struct convention
//! uses synthetic `fN` field names at the Arrow layer, so
//! `BroadcastTupleElement(i)` emits a call with `"fi"` as the field name.

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

use datafusion::arrow::array::{
    make_array, Array, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait, StructArray,
};
use datafusion::arrow::buffer::NullBuffer;
use datafusion::arrow::datatypes::{DataType, Field, FieldRef};
use datafusion::common::{exec_err, internal_err, plan_datafusion_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 HamelinArrayGetFieldUdf {
    signature: Signature,
}

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

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

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

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

    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
        if arg_types.len() != 2 {
            return exec_err!(
                "hamelin_array_get_field requires 2 arguments, 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() != 2 {
            return exec_err!(
                "hamelin_array_get_field requires 2 arguments, got {}",
                args.arg_fields.len()
            );
        }

        let list_field = &args.arg_fields[0];
        let field_name_scalar = args
            .scalar_arguments
            .get(1)
            .and_then(|s| s.as_ref())
            .ok_or_else(|| {
                plan_datafusion_err!(
                    "hamelin_array_get_field: field name argument must be a literal string"
                )
            })?;

        let field_name = field_name_scalar
            .try_as_str()
            .flatten()
            .filter(|s| !s.is_empty())
            .ok_or_else(|| {
                plan_datafusion_err!(
                    "hamelin_array_get_field: field name must be a non-empty string literal"
                )
            })?;

        let (outer_large, struct_field) = match list_field.data_type() {
            DataType::List(inner) => (false, inner.clone()),
            DataType::LargeList(inner) => (true, inner.clone()),
            DataType::Null => {
                return Ok(Arc::new(Field::new(self.name(), DataType::Null, true)));
            }
            other => {
                return exec_err!(
                    "hamelin_array_get_field: expected List or LargeList, got {other}"
                );
            }
        };

        let struct_fields = match struct_field.data_type() {
            DataType::Struct(fields) => fields,
            DataType::Null => {
                return Ok(Arc::new(Field::new(self.name(), DataType::Null, true)));
            }
            other => {
                return exec_err!(
                    "hamelin_array_get_field: list element must be Struct, got {other}"
                );
            }
        };

        let child_field = struct_fields
            .iter()
            .find(|f| f.name() == field_name)
            .ok_or_else(|| plan_datafusion_err!("Field {field_name} not found in struct"))?;

        // Inherit nullability: if the outer list or the struct element is
        // nullable, the output's item field must also be nullable.
        let item_nullable = child_field.is_nullable() || struct_field.is_nullable();
        let item_field = Arc::new(
            child_field
                .as_ref()
                .clone()
                .with_name("item")
                .with_nullable(item_nullable),
        );

        let result_type = if outer_large {
            DataType::LargeList(item_field)
        } else {
            DataType::List(item_field)
        };

        Ok(Arc::new(Field::new(
            self.name(),
            result_type,
            list_field.is_nullable(),
        )))
    }

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

        let field_name = match &args[1] {
            ColumnarValue::Scalar(s) => s
                .try_as_str()
                .flatten()
                .map(|s| s.to_string())
                .ok_or_else(|| {
                    DataFusionError::Execution(
                        "hamelin_array_get_field: field name must be a non-null string literal"
                            .into(),
                    )
                })?,
            _ => {
                return exec_err!(
                    "hamelin_array_get_field: field name must be a scalar string literal"
                );
            }
        };

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

/// Extract a struct field from the flat child of a ListArray and re-wrap.
///
/// `ListArray<StructArray>` stores the struct's fields as separate contiguous
/// child columns. Grabbing the named column is a pointer operation — no data
/// is copied. The result list inherits the parent's offsets and null bitmap.
fn apply_get_field_to_list<O: OffsetSizeTrait>(
    list_array: &GenericListArray<O>,
    field_name: &str,
) -> Result<ArrayRef> {
    let child_values = list_array.values();
    let struct_array = child_values.as_struct_opt().ok_or_else(|| {
        DataFusionError::Execution(format!(
            "hamelin_array_get_field: list element must be Struct, got {:?}",
            child_values.data_type()
        ))
    })?;

    let (child_idx, child_field_ref) = struct_array
        .fields()
        .iter()
        .enumerate()
        .find(|(_, f)| f.name() == field_name)
        .map(|(i, f)| (i, f.clone()))
        .ok_or_else(|| {
            DataFusionError::Execution(format!("Field '{field_name}' not found in struct"))
        })?;

    let field_values = struct_array.column(child_idx).clone();

    // Apply the struct-level null bitmap to the child column so that null
    // struct entries produce null output elements. StructArray's child
    // columns do not inherit the struct's null mask automatically.
    let field_values = apply_parent_nulls(field_values, struct_array)?;

    // Item nullability must match what `return_field_from_args` planned:
    // nullable if either the child field or the outer struct element
    // (list's item field) is nullable. Hardcoding `true` would diverge
    // from the planned schema when both are non-nullable.
    let struct_element_nullable = match list_array.data_type() {
        DataType::List(item) | DataType::LargeList(item) => item.is_nullable(),
        _ => true,
    };
    let item_nullable = child_field_ref.is_nullable() || struct_element_nullable;
    let item_field = Arc::new(
        child_field_ref
            .as_ref()
            .clone()
            .with_name("item")
            .with_nullable(item_nullable),
    );

    let result_list = GenericListArray::<O>::try_new(
        item_field,
        list_array.offsets().clone(),
        field_values,
        list_array.nulls().cloned(),
    )?;

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

/// Combine a struct's null bitmap with a child column's, so the child array
/// has nulls at every index where either the parent struct or the child was
/// null. Zero-copy — rebuilds the ArrayData with a merged null buffer.
fn apply_parent_nulls(child: ArrayRef, struct_array: &StructArray) -> Result<ArrayRef> {
    let Some(parent_nulls) = struct_array.nulls() else {
        return Ok(child);
    };
    let combined = match child.nulls() {
        Some(child_nulls) => NullBuffer::union(Some(parent_nulls), Some(child_nulls)),
        None => Some(parent_nulls.clone()),
    };
    let data = child
        .to_data()
        .into_builder()
        .nulls(combined)
        .build()
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok(make_array(data))
}

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

    // Route empty + null-element scalars through the same kernel so the
    // result carries the output type the planner declared. Short-circuiting
    // with scalar.clone() would return List<Struct{...}> where the planned
    // schema says List<child_field>, and downstream operators that inspect
    // the scalar's data type would reject the mismatch.
    let result = if is_large {
        apply_get_field_to_list(arr.as_list::<i64>(), field_name)?
    } else {
        apply_get_field_to_list(arr.as_list::<i32>(), field_name)?
    };

    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 array_get_field_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinArrayGetFieldUdf::default())
}
