//! Vectorized variant-to-JSON conversion over list arrays.
//!
//! # Why this exists
//!
//! When Hamelin encounters `to_json_string(array_of_variants)` (broadcast
//! element-wise), the default transform lowering generates:
//!
//!   sequence(0, len(arr)-1) → EXPLODE → to_json_string per row → AGG array_agg BY __row_id
//!
//! On high-fan-out inputs this materializes the intermediate strings into per-row
//! offsets that can overflow Arrow's i32 list offsets. It also forces a hash
//! repartition on `__row_id`, plus the usual explode/reaggregate overhead.
//!
//! This UDF avoids all of that. A `ListArray<Variant>` stores its child variant
//! elements as a single contiguous flat array. We serialize each variant into a
//! single `StringArray` once, then re-wrap with the original list offsets and
//! null bitmap. Zero explode, zero reaggregate, zero repartition.
//!
//! Result container type mirrors the input: `List` → `List<Utf8View>`,
//! `LargeList` → `LargeList<Utf8View>`. Inner strings use `Utf8View` so the
//! value buffer is addressed per-string rather than by cumulative offsets —
//! a single batch can hold arbitrarily large total string bytes without
//! hitting the i32 offset cap that a `Utf8` StringArray would.

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

use datafusion::arrow::array::{
    Array, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait, StringViewArray, StringViewBuilder,
};
use datafusion::arrow::datatypes::{DataType, Field};
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::write_variant;

fn utf8_item_field() -> Arc<Field> {
    Arc::new(Field::new("item", DataType::Utf8View, true))
}

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

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

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

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

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

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

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

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

/// Serialize each variant in the flat child values of a ListArray, then re-wrap.
fn apply_to_list<O: OffsetSizeTrait>(list_array: &GenericListArray<O>) -> Result<ArrayRef> {
    let child_values = list_array.values();
    let variant_array = VariantArray::try_new(child_values).map_err(|e| {
        DataFusionError::Execution(format!(
            "hamelin_array_variant_to_json expects List of Variants, 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 result_values: StringViewArray = builder.finish();

    let result_list = GenericListArray::<O>::try_new(
        utf8_item_field(),
        list_array.offsets().clone(),
        Arc::new(result_values) as ArrayRef,
        list_array.nulls().cloned(),
    )?;

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

/// Handle scalar list inputs by extracting the inner array, processing, and re-wrapping.
///
/// Always runs through `apply_to_list` — even for null/empty inputs — so the
/// returned scalar's type matches the UDF's declared return type
/// (`List<Utf8View>` / `LargeList<Utf8View>`) rather than leaking the input
/// variant type.
fn scalar_array_variant_to_json(scalar: &ScalarValue) -> Result<ColumnarValue> {
    let result = match scalar {
        ScalarValue::List(arr) => apply_to_list(arr.as_ref())?,
        ScalarValue::LargeList(arr) => apply_to_list(arr.as_ref())?,
        _ => {
            return exec_err!(
                "hamelin_array_variant_to_json expects List type, got {:?}",
                scalar
            );
        }
    };

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

pub fn array_variant_to_json_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(HamelinArrayVariantToJsonUdf::default())
}
