//! zip UDF for DataFusion: zips two arrays into an array of tuples (struct with `f0`, `f1` fields).

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

use datafusion::arrow::array::{
    make_array, new_null_array, Array, ArrayRef, AsArray, GenericListArray, ListArray,
    MutableArrayData, OffsetSizeTrait, StructArray,
};
use datafusion::arrow::buffer::{NullBuffer, OffsetBuffer};
use datafusion::arrow::compute::concat;
use datafusion::arrow::datatypes::{DataType, Field, Fields};
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
    Volatility,
};

use super::array_helpers::{list_element_type, scalar_list_values};

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

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

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

/// Tuple layout matches Hamelin tuple field names (`f0`, `f1`, …) used by `get_field`.
fn make_tuple_struct_type(left_elem: &DataType, right_elem: &DataType) -> DataType {
    DataType::Struct(Fields::from(vec![
        Field::new("f0", left_elem.clone(), true),
        Field::new("f1", right_elem.clone(), true),
    ]))
}

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

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        let left_elem = list_element_type(&arg_types[0], "zip")?;
        let right_elem = list_element_type(&arg_types[1], "zip")?;
        Ok(DataType::List(Arc::new(Field::new_list_field(
            make_tuple_struct_type(left_elem, right_elem),
            true,
        ))))
    }

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

        match (&args[0], &args[1]) {
            (ColumnarValue::Scalar(left), ColumnarValue::Scalar(right)) => zip_scalars(left, right),
            (ColumnarValue::Array(left), ColumnarValue::Array(right)) => zip_columns(left, right),
            (ColumnarValue::Scalar(left), ColumnarValue::Array(right)) => {
                let left = left.to_array_of_size(right.len())?;
                zip_columns(&left, right)
            }
            (ColumnarValue::Array(left), ColumnarValue::Scalar(right)) => {
                let right = right.to_array_of_size(left.len())?;
                zip_columns(left, &right)
            }
        }
    }
}

fn zip_scalars(left: &ScalarValue, right: &ScalarValue) -> Result<ColumnarValue> {
    let left_dt = left.data_type();
    let right_dt = right.data_type();
    let left_elem_type = list_element_type(&left_dt, "zip")?.clone();
    let right_elem_type = list_element_type(&right_dt, "zip")?.clone();
    let struct_type = make_tuple_struct_type(&left_elem_type, &right_elem_type);

    let left_values = scalar_list_values(left, "zip")?;
    let right_values = scalar_list_values(right, "zip")?;

    match (left_values, right_values) {
        (None, _) | (_, None) => {
            let list_type = DataType::List(Arc::new(Field::new_list_field(struct_type, true)));
            Ok(ColumnarValue::Scalar(ScalarValue::try_from(&list_type)?))
        }
        (Some(left_arr), Some(right_arr)) => {
            let zipped = zip_single_pair(&left_arr, &right_arr, &left_elem_type, &right_elem_type)?;
            let struct_field = Field::new_list_field(struct_type, true);
            let list_arr = ListArray::new(
                Arc::new(struct_field),
                OffsetBuffer::from_lengths([zipped.len()]),
                Arc::new(zipped) as ArrayRef,
                None,
            );
            Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new(list_arr))))
        }
    }
}

fn zip_columns(left: &ArrayRef, right: &ArrayRef) -> Result<ColumnarValue> {
    match (left.data_type(), right.data_type()) {
        (DataType::List(_), DataType::List(_)) => {
            zip_list_columns::<i32, i32>(left.as_list::<i32>(), right.as_list::<i32>())
        }
        (DataType::LargeList(_), DataType::LargeList(_)) => {
            zip_list_columns::<i64, i64>(left.as_list::<i64>(), right.as_list::<i64>())
        }
        (DataType::List(_), DataType::LargeList(_)) => {
            zip_list_columns::<i32, i64>(left.as_list::<i32>(), right.as_list::<i64>())
        }
        (DataType::LargeList(_), DataType::List(_)) => {
            zip_list_columns::<i64, i32>(left.as_list::<i64>(), right.as_list::<i32>())
        }
        (l, r) => exec_err!("zip expects array types, got {:?} and {:?}", l, r),
    }
}

fn zip_list_columns<L: OffsetSizeTrait, R: OffsetSizeTrait>(
    left: &GenericListArray<L>,
    right: &GenericListArray<R>,
) -> Result<ColumnarValue> {
    let left_elem_type = left.value_type();
    let right_elem_type = right.value_type();
    let struct_type = make_tuple_struct_type(&left_elem_type, &right_elem_type);
    let struct_field = Arc::new(Field::new_list_field(struct_type, true));

    let num_rows = left.len();
    let left_off = left.value_offsets();
    let right_off = right.value_offsets();
    let left_vals = left.values();
    let right_vals = right.values();

    let mut offsets: Vec<i32> = Vec::with_capacity(num_rows + 1);
    offsets.push(0);
    let mut end_offset: i32 = 0;
    let mut nulls = Vec::with_capacity(num_rows);

    for i in 0..num_rows {
        if left.is_null(i) || right.is_null(i) {
            nulls.push(false);
            offsets.push(end_offset);
            continue;
        }
        nulls.push(true);

        let left_start = left_off[i].as_usize();
        let left_end = left_off[i + 1].as_usize();
        let right_start = right_off[i].as_usize();
        let right_end = right_off[i + 1].as_usize();

        let left_len = left_end - left_start;
        let right_len = right_end - right_start;
        let max_len = left_len.max(right_len);
        let delta = match i32::try_from(max_len) {
            Ok(d) => d,
            Err(_) => {
                return exec_err!("zip: row element count {} does not fit in i32", max_len);
            }
        };
        end_offset = match end_offset.checked_add(delta) {
            Some(next) => next,
            None => {
                return exec_err!(
                    "zip: total zipped element count overflowed i32 (List array limit)"
                );
            }
        };
        offsets.push(end_offset);
    }

    let fields = Fields::from(vec![
        Field::new("f0", left_elem_type.clone(), true),
        Field::new("f1", right_elem_type.clone(), true),
    ]);

    let total_elements: usize = end_offset as usize;
    let all_values: ArrayRef = if total_elements == 0 {
        Arc::new(StructArray::new(
            fields,
            vec![
                new_null_array(&left_elem_type, 0),
                new_null_array(&right_elem_type, 0),
            ],
            None,
        ))
    } else {
        let left_data = left_vals.to_data();
        let right_data = right_vals.to_data();
        let mut left_mutable = MutableArrayData::new(vec![&left_data], true, total_elements);
        let mut right_mutable = MutableArrayData::new(vec![&right_data], true, total_elements);

        for i in 0..num_rows {
            if left.is_null(i) || right.is_null(i) {
                continue;
            }
            let left_start = left_off[i].as_usize();
            let left_end = left_off[i + 1].as_usize();
            let right_start = right_off[i].as_usize();
            let right_end = right_off[i + 1].as_usize();
            let left_len = left_end - left_start;
            let right_len = right_end - right_start;
            let max_len = left_len.max(right_len);

            left_mutable.extend(0, left_start, left_end);
            if left_len < max_len {
                left_mutable.extend_nulls(max_len - left_len);
            }
            right_mutable.extend(0, right_start, right_end);
            if right_len < max_len {
                right_mutable.extend_nulls(max_len - right_len);
            }
        }

        debug_assert_eq!(left_mutable.len(), total_elements);
        debug_assert_eq!(right_mutable.len(), total_elements);

        let left_all = make_array(left_mutable.freeze());
        let right_all = make_array(right_mutable.freeze());
        Arc::new(StructArray::new(fields, vec![left_all, right_all], None))
    };

    let result = ListArray::new(
        struct_field,
        OffsetBuffer::new(offsets.into()),
        all_values,
        Some(NullBuffer::from(nulls)),
    );

    Ok(ColumnarValue::Array(Arc::new(result)))
}

/// Zip a single pair of arrays into a StructArray.
fn zip_single_pair(
    left: &dyn Array,
    right: &dyn Array,
    left_type: &DataType,
    right_type: &DataType,
) -> Result<StructArray> {
    let max_len = left.len().max(right.len());

    let left_padded = pad_array(left, max_len, left_type)?;
    let right_padded = pad_array(right, max_len, right_type)?;

    let fields = Fields::from(vec![
        Field::new("f0", left_type.clone(), true),
        Field::new("f1", right_type.clone(), true),
    ]);

    Ok(StructArray::new(
        fields,
        vec![left_padded, right_padded],
        None,
    ))
}

/// Pad an array to the target length by appending nulls.
fn pad_array(arr: &dyn Array, target_len: usize, data_type: &DataType) -> Result<ArrayRef> {
    if arr.len() == target_len {
        return Ok(make_array(arr.to_data()));
    }

    let padding_len = target_len - arr.len();
    let null_padding = new_null_array(data_type, padding_len);

    Ok(concat(&[arr, null_padding.as_ref()])?)
}

pub fn array_zip_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(ArrayZipUdf::new())
}
