//! Variant to JSON string UDF for DataFusion.
//!
//! Converts a Variant to its JSON string representation.

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

use datafusion::arrow::array::{ArrayRef, StringViewBuilder};
use datafusion::arrow::datatypes::DataType;
use datafusion::common::{exec_err, Result, ScalarValue};
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};
use parquet_variant_compute::VariantArray;

use super::write_variant;

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

impl Default for VariantToJsonUdf {
    fn default() -> Self {
        Self {
            signature: Signature::any(1, Volatility::Immutable),
        }
    }
}

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

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

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

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(DataType::Utf8View)
    }

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

        match &args[0] {
            ColumnarValue::Scalar(scalar) => {
                if scalar.is_null() {
                    return Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(None)));
                }
                let array = scalar.to_array_of_size(1)?;
                let variant_array = VariantArray::try_new(&array).map_err(|e| {
                    datafusion::error::DataFusionError::Execution(format!(
                        "hamelin_to_json_string expects a Variant, got error: {e}"
                    ))
                })?;
                let mut json = String::new();
                write_variant(&mut json, &variant_array.value(0))
                    .map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string()))?;
                Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(json))))
            }
            ColumnarValue::Array(array) => {
                let variant_array = VariantArray::try_new(array).map_err(|e| {
                    datafusion::error::DataFusionError::Execution(format!(
                        "hamelin_to_json_string expects a Variant, got error: {e}"
                    ))
                })?;

                let mut builder = StringViewBuilder::new();
                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| {
                            datafusion::error::DataFusionError::Execution(e.to_string())
                        })?;
                        builder.append_value(&scratch);
                    }
                }

                Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
            }
        }
    }
}

pub fn variant_to_json_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(VariantToJsonUdf::default())
}
