//! Two-list map constructor for Hamelin-created maps.
//!
//! DataFusion's built-in `map(keys, values)` is strict about the physical shape
//! of its two list arguments. After optimization/serialization a scalar list can
//! be paired with a per-row list array, and the built-in compares the scalar
//! key-list length with the batch row count of the values. This UDF makes the
//! row semantics explicit: scalar lists broadcast, maps are built row by row,
//! and null/mismatched rows become null maps.

use std::any::Any;
use std::collections::HashSet;
use std::sync::{Arc, OnceLock};

use datafusion::arrow::array::{Array, ArrayRef, AsArray, MapArray, OffsetSizeTrait, StructArray};
use datafusion::arrow::buffer::{NullBuffer, OffsetBuffer};
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 crate::struct_expansion::map_data_type;

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

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

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

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

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        if arg_types.len() != 2 {
            return exec_err!(
                "hamelin_map_from_arrays expects exactly 2 arguments, got {}",
                arg_types.len()
            );
        }
        Ok(map_data_type(
            list_element_type(&arg_types[0])?.clone(),
            list_element_type(&arg_types[1])?.clone(),
        ))
    }

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

        let keys = ListInput::try_new(&args[0])?;
        let values = ListInput::try_new(&args[1])?;
        let key_type = keys.element_type().clone();
        let value_type = values.element_type().clone();

        let output_len = match (keys.len(), values.len()) {
            (None, None) => 1,
            (Some(len), None) | (None, Some(len)) => len,
            (Some(left), Some(right)) if left == right => left,
            (Some(left), Some(right)) => {
                return exec_err!(
                    "hamelin_map_from_arrays array arguments have different row counts: {left} and {right}"
                );
            }
        };

        let result = build_map_array(&keys, &values, output_len, key_type, value_type)?;
        if keys.is_scalar() && values.is_scalar() {
            ScalarValue::try_from_array(result.as_ref(), 0).map(ColumnarValue::Scalar)
        } else {
            Ok(ColumnarValue::Array(result))
        }
    }
}

enum ListInput<'a> {
    Scalar {
        list: Option<ArrayRef>,
        element_type: DataType,
    },
    Array {
        array: &'a ArrayRef,
        element_type: DataType,
    },
}

impl<'a> ListInput<'a> {
    fn try_new(value: &'a ColumnarValue) -> Result<Self> {
        match value {
            ColumnarValue::Scalar(scalar) => match scalar {
                ScalarValue::List(array) => {
                    let list = (!array.is_empty() && !array.is_null(0)).then(|| array.value(0));
                    Ok(Self::Scalar {
                        list,
                        element_type: list_element_type(&scalar.data_type())?.clone(),
                    })
                }
                ScalarValue::LargeList(array) => {
                    let list = (!array.is_empty() && !array.is_null(0)).then(|| array.value(0));
                    Ok(Self::Scalar {
                        list,
                        element_type: list_element_type(&scalar.data_type())?.clone(),
                    })
                }
                other => exec_err!(
                    "hamelin_map_from_arrays expects List or LargeList arguments, got scalar {:?}",
                    other.data_type()
                ),
            },
            ColumnarValue::Array(array) => Ok(Self::Array {
                array,
                element_type: list_element_type(array.data_type())?.clone(),
            }),
        }
    }

    fn len(&self) -> Option<usize> {
        match self {
            Self::Scalar { .. } => None,
            Self::Array { array, .. } => Some(array.len()),
        }
    }

    fn is_scalar(&self) -> bool {
        matches!(self, Self::Scalar { .. })
    }

    fn element_type(&self) -> &DataType {
        match self {
            Self::Scalar { element_type, .. } | Self::Array { element_type, .. } => element_type,
        }
    }

    fn row(&self, row: usize) -> Result<Option<ArrayRef>> {
        match self {
            Self::Scalar { list, .. } => Ok(list.clone()),
            Self::Array { array, .. } => match array.data_type() {
                DataType::List(_) => Ok(list_row::<i32>(array.as_ref(), row)),
                DataType::LargeList(_) => Ok(list_row::<i64>(array.as_ref(), row)),
                other => exec_err!(
                    "hamelin_map_from_arrays expects List or LargeList arguments, got {:?}",
                    other
                ),
            },
        }
    }
}

fn list_row<O: OffsetSizeTrait>(array: &dyn Array, row: usize) -> Option<ArrayRef> {
    let list = array.as_list::<O>();
    (!list.is_null(row)).then(|| list.value(row))
}

fn list_element_type(data_type: &DataType) -> Result<&DataType> {
    match data_type {
        DataType::List(field) | DataType::LargeList(field) => Ok(field.data_type()),
        other => exec_err!(
            "hamelin_map_from_arrays expects List or LargeList arguments, got {:?}",
            other
        ),
    }
}

fn build_map_array(
    keys: &ListInput<'_>,
    values: &ListInput<'_>,
    output_len: usize,
    key_type: DataType,
    value_type: DataType,
) -> Result<ArrayRef> {
    let mut all_keys = Vec::new();
    let mut all_values = Vec::new();
    let mut offsets = Vec::with_capacity(output_len + 1);
    let mut valid = Vec::with_capacity(output_len);
    offsets.push(0i32);

    for row in 0..output_len {
        let Some(row_keys) = keys.row(row)? else {
            valid.push(false);
            offsets.push(*offsets.last().unwrap_or(&0));
            continue;
        };
        let Some(row_values) = values.row(row)? else {
            valid.push(false);
            offsets.push(*offsets.last().unwrap_or(&0));
            continue;
        };

        if row_keys.len() != row_values.len() {
            valid.push(false);
            offsets.push(*offsets.last().unwrap_or(&0));
            continue;
        }

        validate_map_keys(row_keys.as_ref())?;

        let new_len = all_keys.len().checked_add(row_keys.len()).ok_or_else(|| {
            datafusion::common::DataFusionError::Execution(
                "hamelin_map_from_arrays offset overflow".to_string(),
            )
        })?;
        let new_offset = i32::try_from(new_len).map_err(|_| {
            datafusion::common::DataFusionError::Execution(format!(
                "hamelin_map_from_arrays offset overflow: {new_len} entries exceeds i32::MAX"
            ))
        })?;

        for idx in 0..row_keys.len() {
            all_keys.push(ScalarValue::try_from_array(row_keys.as_ref(), idx)?);
            all_values.push(ScalarValue::try_from_array(row_values.as_ref(), idx)?);
        }

        valid.push(true);
        offsets.push(new_offset);
    }

    let keys_array = if all_keys.is_empty() {
        datafusion::arrow::array::new_empty_array(&key_type)
    } else {
        ScalarValue::iter_to_array(all_keys.into_iter())?
    };
    let values_array = if all_values.is_empty() {
        datafusion::arrow::array::new_empty_array(&value_type)
    } else {
        ScalarValue::iter_to_array(all_values.into_iter())?
    };

    let entry_fields = Fields::from(vec![
        Field::new("key", key_type, false),
        Field::new("value", value_type, true),
    ]);
    let entries = StructArray::new(entry_fields.clone(), vec![keys_array, values_array], None);
    let nulls = (!valid.iter().all(|is_valid| *is_valid)).then(|| NullBuffer::from(valid));
    let map_array = MapArray::new(
        Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)),
        OffsetBuffer::new(offsets.into()),
        entries,
        nulls,
        false,
    );

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

fn validate_map_keys(array: &dyn Array) -> Result<()> {
    let mut seen_keys = HashSet::with_capacity(array.len());
    for idx in 0..array.len() {
        let key = ScalarValue::try_from_array(array, idx)?;
        if key.is_null() {
            return exec_err!("map key cannot be null");
        }
        if seen_keys.contains(&key) {
            return exec_err!("map key must be unique, duplicate key found: {}", key);
        }
        seen_keys.insert(key);
    }
    Ok(())
}

static MAP_FROM_ARRAYS_UDF: OnceLock<ScalarUDF> = OnceLock::new();

pub fn map_from_arrays_udf() -> ScalarUDF {
    MAP_FROM_ARRAYS_UDF
        .get_or_init(|| ScalarUDF::new_from_impl(MapFromArraysUdf::new()))
        .clone()
}

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

    fn string_array(values: &[&str]) -> ArrayRef {
        Arc::new(
            values
                .iter()
                .map(|value| Some(*value))
                .collect::<StringViewArray>(),
        )
    }

    fn string_list_array(
        lengths: impl IntoIterator<Item = usize>,
        values: &[&str],
        validity: Option<Vec<bool>>,
    ) -> ListArray {
        ListArray::new(
            Arc::new(Field::new("item", DataType::Utf8View, true)),
            OffsetBuffer::from_lengths(lengths),
            string_array(values),
            validity.map(NullBuffer::from),
        )
    }

    fn scalar_string_list(values: &[&str]) -> ScalarValue {
        ScalarValue::List(Arc::new(string_list_array([values.len()], values, None)))
    }

    fn invoke(args: Vec<ColumnarValue>, number_rows: usize) -> ColumnarValue {
        let udf = MapFromArraysUdf::new();
        let arg_fields: Vec<_> = args
            .iter()
            .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true)))
            .collect();
        let arg_types: Vec<_> = args.iter().map(ColumnarValue::data_type).collect();
        let return_type = udf.return_type(&arg_types).expect("return type");
        udf.invoke_with_args(ScalarFunctionArgs {
            args,
            arg_fields,
            number_rows,
            return_field: Arc::new(Field::new("result", return_type, true)),
            config_options: Arc::new(ConfigOptions::default()),
        })
        .expect("invoke")
    }

    fn output_map(value: ColumnarValue) -> MapArray {
        match value {
            ColumnarValue::Array(array) => array.as_map().clone(),
            other => panic!("expected array output, got {other:?}"),
        }
    }

    fn string_values(array: &ArrayRef) -> Vec<Option<String>> {
        let strings = array
            .as_any()
            .downcast_ref::<StringViewArray>()
            .expect("StringViewArray");
        (0..strings.len())
            .map(|idx| {
                strings
                    .is_valid(idx)
                    .then(|| strings.value(idx).to_string())
            })
            .collect()
    }

    #[test]
    fn broadcasts_scalar_key_list_over_value_rows() {
        let values = string_list_array(
            [3, 3, 3, 3],
            &[
                "api",
                "get",
                "us-east-1",
                "worker",
                "put",
                "us-west-2",
                "web",
                "post",
                "eu-west-1",
                "cron",
                "delete",
                "ap-south-1",
            ],
            None,
        );

        let result = output_map(invoke(
            vec![
                ColumnarValue::Scalar(scalar_string_list(&["service", "action", "region"])),
                ColumnarValue::Array(Arc::new(values)),
            ],
            4,
        ));

        assert_eq!(result.len(), 4);
        assert_eq!(
            result.offsets().iter().copied().collect::<Vec<_>>(),
            vec![0, 3, 6, 9, 12]
        );
        assert_eq!(
            string_values(result.keys()),
            vec![
                Some("service".to_string()),
                Some("action".to_string()),
                Some("region".to_string()),
                Some("service".to_string()),
                Some("action".to_string()),
                Some("region".to_string()),
                Some("service".to_string()),
                Some("action".to_string()),
                Some("region".to_string()),
                Some("service".to_string()),
                Some("action".to_string()),
                Some("region".to_string()),
            ]
        );
        assert_eq!(
            string_values(result.values()),
            vec![
                Some("api".to_string()),
                Some("get".to_string()),
                Some("us-east-1".to_string()),
                Some("worker".to_string()),
                Some("put".to_string()),
                Some("us-west-2".to_string()),
                Some("web".to_string()),
                Some("post".to_string()),
                Some("eu-west-1".to_string()),
                Some("cron".to_string()),
                Some("delete".to_string()),
                Some("ap-south-1".to_string()),
            ]
        );
    }

    #[test]
    fn mismatched_value_rows_become_null_maps() {
        let values = string_list_array(
            [3, 2, 3],
            &[
                "api",
                "get",
                "us-east-1",
                "worker",
                "put",
                "web",
                "post",
                "eu-west-1",
            ],
            None,
        );

        let result = output_map(invoke(
            vec![
                ColumnarValue::Scalar(scalar_string_list(&["service", "action", "region"])),
                ColumnarValue::Array(Arc::new(values)),
            ],
            3,
        ));

        assert_eq!(result.len(), 3);
        assert!(result.is_valid(0));
        assert!(result.is_null(1));
        assert!(result.is_valid(2));
        assert_eq!(
            result.offsets().iter().copied().collect::<Vec<_>>(),
            vec![0, 3, 3, 6]
        );
    }
}
