//! Vectorized variant-to-JSON conversion over map values.
//!
//! Counterpart to `hamelin_array_variant_to_json`, but for map values. Operates
//! directly on the `MapArray`'s value column, serializes each variant with the
//! shared `write_variant` writer, and re-wraps the result with the original
//! keys, entries struct, and map offsets.

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

use datafusion::arrow::array::{ArrayRef, AsArray, MapArray, StringViewArray, StringViewBuilder};
use datafusion::arrow::datatypes::{DataType, Field, Fields};
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature,
    Volatility,
};
use parquet_variant_compute::VariantArray;

use super::{rebuild_map, write_variant};

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

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

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

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        let Some(DataType::Map(field, sorted)) = arg_types.first() else {
            return exec_err!(
                "hamelin_map_variant_to_json_values expects Map, got {:?}",
                arg_types.first()
            );
        };
        let DataType::Struct(entry_fields) = field.data_type() else {
            return exec_err!("expected Struct inside Map");
        };
        let new_value_field = Arc::new(Field::new(
            entry_fields[1].name(),
            DataType::Utf8View,
            entry_fields[1].is_nullable(),
        ));
        let new_struct_fields = Fields::from(vec![entry_fields[0].clone(), new_value_field]);
        Ok(DataType::Map(
            Arc::new(Field::new(
                field.name(),
                DataType::Struct(new_struct_fields),
                field.is_nullable(),
            )),
            *sorted,
        ))
    }

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

        match &args[0] {
            ColumnarValue::Array(array) => {
                let map_array = array.as_map();
                let result = apply_to_map(map_array)?;
                Ok(ColumnarValue::Array(result))
            }
            ColumnarValue::Scalar(ScalarValue::Map(arr)) => {
                // Always run through apply_to_map so the returned scalar's
                // type matches the UDF's declared return type rather than
                // leaking the input variant type on the null/empty path.
                let result = apply_to_map(arr.as_ref())?;
                ScalarValue::try_from_array(&result, 0).map(ColumnarValue::Scalar)
            }
            ColumnarValue::Scalar(other) => exec_err!(
                "hamelin_map_variant_to_json_values expects Map, got {:?}",
                other
            ),
        }
    }
}

fn apply_to_map(map_array: &MapArray) -> Result<ArrayRef> {
    let entries = map_array.entries();
    let keys = entries.column(0).clone();
    let values = entries.column(1);

    let variant_array = VariantArray::try_new(values).map_err(|e| {
        DataFusionError::Execution(format!(
            "hamelin_map_variant_to_json_values expects Map of Variant values, got error: {e}"
        ))
    })?;

    let mut builder = StringViewBuilder::with_capacity(variant_array.len());
    let mut scratch = String::new();
    for idx in 0..variant_array.len() {
        if !variant_array.is_valid(idx) {
            builder.append_null();
        } else {
            scratch.clear();
            write_variant(&mut scratch, &variant_array.value(idx))
                .map_err(|e| DataFusionError::Execution(e.to_string()))?;
            builder.append_value(&scratch);
        }
    }
    let new_values: StringViewArray = builder.finish();

    rebuild_map(map_array, keys, Arc::new(new_values) as ArrayRef)
}

pub fn map_variant_to_json_values_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinMapVariantToJsonValuesUdf::default())
}
