//! Vectorized variant field extraction over map values.
//!
//! When Hamelin encounters `transform_values(map(K, variant), x -> x.field)`,
//! the broadcast lowering would otherwise generate:
//!
//!   map_values → sequence → EXPLODE → variant_get → AGG array_agg BY __row_id
//!   → map(map_keys, ...)
//!
//! which forces a hash repartition and materializes an intermediate
//! `array(variant)` per row. This UDF operates directly on the map's value
//! child, applies `variant_get` once to the flat values, and re-wraps the
//! result with the original entries struct / map offsets.

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

use datafusion::arrow::array::{ArrayRef, AsArray, MapArray};
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::{VariantPath, VariantPathElement};
use parquet_variant_compute::{variant_get as parquet_variant_get, GetOptions};

use super::{rebuild_map, scalar_to_string, variant_fields};

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

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

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

    fn name(&self) -> &str {
        "hamelin_map_variant_get_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_get_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::Struct(variant_fields()),
            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() != 2 {
            return exec_err!(
                "hamelin_map_variant_get_values expects 2 arguments, got {}",
                args.len()
            );
        }

        let path = match &args[1] {
            ColumnarValue::Scalar(s) => scalar_to_string(s)?.ok_or_else(|| {
                DataFusionError::Execution(
                    "hamelin_map_variant_get_values: path must be a non-null string literal".into(),
                )
            })?,
            _ => return exec_err!("hamelin_map_variant_get_values: path must be a scalar string"),
        };

        match &args[0] {
            ColumnarValue::Array(array) => {
                let map_array = array.as_map();
                let result = apply_variant_get_to_map(map_array, &path)?;
                Ok(ColumnarValue::Array(result))
            }
            ColumnarValue::Scalar(ScalarValue::Map(arr)) => {
                // Always run through apply_variant_get_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_variant_get_to_map(arr.as_ref(), &path)?;
                ScalarValue::try_from_array(&result, 0).map(ColumnarValue::Scalar)
            }
            ColumnarValue::Scalar(other) => exec_err!(
                "hamelin_map_variant_get_values expects Map, got {:?}",
                other
            ),
        }
    }
}

fn apply_variant_get_to_map(map_array: &MapArray, path: &str) -> Result<ArrayRef> {
    let entries = map_array.entries();
    let keys = entries.column(0);
    let values = entries.column(1);

    let variant_path = VariantPath::new(vec![VariantPathElement::field(path)]);
    let options = GetOptions {
        path: variant_path,
        as_type: None,
        ..Default::default()
    };
    let new_values = parquet_variant_get(values, options)?;

    rebuild_map(map_array, keys.clone(), new_values)
}

pub fn map_variant_get_values_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinMapVariantGetValuesUdf::default())
}
