//! Vectorized variant field extraction over list arrays.
//!
//! # Why this exists
//!
//! When Hamelin encounters `array_of_variants.field_name`, the broadcast lowering
//! generates `transform(arr, x -> x.field)`. Without this UDF, DataFusion (which
//! doesn't support lambdas) lowers that to:
//!
//!   sequence(0, len(arr)-1) → EXPLODE → variant_get per row → AGG array_agg BY __row_id
//!
//! The AGG BY __row_id forces a hash repartition (29% of execution time in profiling),
//! even though __row_id is partition-local. The explode/reaggregate also has inherent
//! overhead from UUID generation, sorting, and sliding_array_agg.
//!
//! This UDF avoids all of that. A `ListArray<Variant>` stores its child variant
//! elements as a single contiguous flat array. We apply `variant_get` once to the
//! entire flat child array, then re-wrap the result with the original list offsets
//! and null bitmap. Zero explode, zero reaggregate, zero repartition.

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

use datafusion::arrow::array::{Array, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait};
use datafusion::arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_err, Result, ScalarValue};
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::{scalar_to_string, variant_fields};

/// The variant element field for list wrapping.
fn variant_item_field() -> Arc<Field> {
    Arc::new(Field::new("item", DataType::Struct(variant_fields()), true))
}

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

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

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

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

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

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        // Preserve the container type: List stays List, LargeList stays LargeList.
        match arg_types.first() {
            Some(DataType::LargeList(_)) => Ok(DataType::LargeList(variant_item_field())),
            _ => Ok(DataType::List(variant_item_field())),
        }
    }

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

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

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

/// Apply variant_get to the flat child values of a ListArray, then re-wrap.
///
/// A `ListArray<Struct<metadata, value>>` stores all variant elements contiguously.
/// We call the `parquet_variant_compute` kernel once on the entire flat child array,
/// then create a new `ListArray` with the same offsets and null bitmap wrapping
/// the result.
fn apply_variant_get_to_list<O: OffsetSizeTrait>(
    list_array: &GenericListArray<O>,
    path: &str,
) -> Result<ArrayRef> {
    let child_values = list_array.values();

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

    let result_values = parquet_variant_get(child_values, options)?;

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

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

/// Handle scalar list inputs by extracting the inner array, processing, and re-wrapping.
fn scalar_array_variant_get(scalar: &ScalarValue, path: &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),
        _ => {
            return exec_err!(
                "hamelin_array_variant_get expects List type, got {:?}",
                scalar
            );
        }
    };

    if arr.is_empty() || arr.is_null(0) {
        return Ok(ColumnarValue::Scalar(scalar.clone()));
    }

    let result = if is_large {
        apply_variant_get_to_list(arr.as_list::<i64>(), path)?
    } else {
        apply_variant_get_to_list(arr.as_list::<i32>(), path)?
    };

    ScalarValue::try_from_array(&result, 0).map(ColumnarValue::Scalar)
}

pub fn array_variant_get_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinArrayVariantGetUdf::default())
}
