shared.rs

use std::sync::Arc;

#[cfg(test)]
use arrow::array::StructArray;
use arrow::array::{Array, ArrayRef, cast::AsArray};
#[cfg(test)]
use arrow_schema::Fields;
use arrow_schema::extension::ExtensionType;
use arrow_schema::{DataType, Field};
use datafusion::common::exec_datafusion_err;
use datafusion::error::{DataFusionError, Result};
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs};
use datafusion::{common::exec_err, scalar::ScalarValue};
use parquet_variant::Variant;
use parquet_variant_compute::{VariantArray, VariantType};

use crate::variant_get::PathMode;

#[cfg(test)]
use parquet_variant_compute::VariantArrayBuilder;

#[cfg(test)]
use parquet_variant_json::JsonToVariant;

pub fn try_field_as_variant_array(field: &Field) -> Result<()> {
    ensure(
        matches!(field.extension_type(), VariantType),
        "field does not have extension type VariantType",
    )?;

    let variant_type = VariantType;
    variant_type.supports_data_type(field.data_type())?;

    Ok(())
}

pub fn _try_field_as_binary(field: &Field) -> Result<()> {
    match field.data_type() {
        DataType::Binary | DataType::BinaryView | DataType::LargeBinary => {}
        unsupported => return exec_err!("expected binary field, got {unsupported} field"),
    }

    Ok(())
}

pub fn try_field_as_string(field: &Field) -> Result<()> {
    match field.data_type() {
        DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => {}
        unsupported => return exec_err!("expected string field, got {unsupported} field"),
    }

    Ok(())
}

pub fn try_parse_variant_scalar(scalar: &ScalarValue) -> Result<VariantArray> {
    let v = match scalar {
        ScalarValue::Struct(v) => v,
        unsupported => {
            return exec_err!(
                "expected variant scalar value, got data type: {}",
                unsupported.data_type()
            );
        }
    };

    VariantArray::try_new(v.as_ref()).map_err(Into::into)
}

pub fn try_parse_binary_scalar(scalar: &ScalarValue) -> Result<Option<&Vec<u8>>> {
    let b = match scalar {
        ScalarValue::Binary(b) | ScalarValue::BinaryView(b) | ScalarValue::LargeBinary(b) => b,
        unsupported => {
            return exec_err!(
                "expected binary scalar value, got data type: {}",
                unsupported.data_type()
            );
        }
    };

    Ok(b.as_ref())
}

pub fn try_parse_binary_columnar(array: &Arc<dyn Array>) -> Result<Vec<Option<&[u8]>>> {
    if let Some(binary_array) = array.as_binary_opt::<i32>() {
        return Ok(binary_array.into_iter().collect::<Vec<_>>());
    }

    if let Some(binary_view_array) = array.as_binary_view_opt() {
        return Ok(binary_view_array.into_iter().collect::<Vec<_>>());
    }

    if let Some(large_binary_array) = array.as_binary_opt::<i64>() {
        return Ok(large_binary_array.into_iter().collect::<Vec<_>>());
    }

    Err(exec_datafusion_err!("expected binary array"))
}

pub fn try_parse_string_scalar(scalar: &ScalarValue) -> Result<Option<&String>> {
    let b = match scalar {
        ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => s,
        unsupported => {
            return exec_err!(
                "expected binary scalar value, got data type: {}",
                unsupported.data_type()
            );
        }
    };

    Ok(b.as_ref())
}

pub fn try_parse_string_columnar(array: &Arc<dyn Array>) -> Result<Vec<Option<&str>>> {
    if let Some(string_array) = array.as_string_opt::<i32>() {
        return Ok(string_array.into_iter().collect::<Vec<_>>());
    }

    if let Some(string_view_array) = array.as_string_view_opt() {
        return Ok(string_view_array.into_iter().collect::<Vec<_>>());
    }

    if let Some(large_string_array) = array.as_string_opt::<i64>() {
        return Ok(large_string_array.into_iter().collect::<Vec<_>>());
    }

    Err(exec_datafusion_err!("expected string array"))
}

pub fn variant_get_single_value<T>(
    variant_array: &VariantArray,
    index: usize,
    path: &str,
    path_mode: PathMode,
    extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
) -> Result<Option<T>> {
    let Some(variant) = variant_array.iter().nth(index).flatten() else {
        return Ok(None);
    };

    let variant_path = path_mode.try_build_path(path)?;
    let Some(value) = variant.get_path(&variant_path) else {
        return Ok(None);
    };

    extract(value)
}

pub fn variant_get_array_values<T>(
    variant_array: &VariantArray,
    path: &str,
    path_mode: PathMode,
    extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
) -> Result<Vec<Option<T>>> {
    let variant_path = path_mode.try_build_path(path)?;

    variant_array
        .iter()
        .map(|maybe_variant| {
            let Some(variant) = maybe_variant else {
                return Ok(None);
            };

            let Some(value) = variant.get_path(&variant_path) else {
                return Ok(None);
            };

            extract(value)
        })
        .collect()
}

pub fn invoke_variant_get_typed<T>(
    args: ScalarFunctionArgs,
    scalar_from_option: fn(Option<T>) -> ScalarValue,
    array_from_values: fn(Vec<Option<T>>) -> ArrayRef,
    extract: for<'m, 'v> fn(Variant<'m, 'v>) -> Result<Option<T>>,
    path_mode: PathMode,
) -> Result<ColumnarValue> {
    let (variant_arg, path_arg) = match args.args.as_slice() {
        [variant_arg, path_arg] => (variant_arg, path_arg),
        _ => return exec_err!("expected 2 arguments"),
    };

    let variant_field = args
        .arg_fields
        .first()
        .ok_or_else(|| exec_datafusion_err!("expected argument field"))?;

    try_field_as_variant_array(variant_field.as_ref())?;

    let out = match (variant_arg, path_arg) {
        (ColumnarValue::Array(variant_array), ColumnarValue::Scalar(path_scalar)) => {
            let path = try_parse_string_scalar(path_scalar)?
                .map(|s| s.as_str())
                .unwrap_or_default();

            let variant_array = VariantArray::try_new(variant_array.as_ref())?;
            let values = variant_get_array_values(&variant_array, path, path_mode, extract)?;
            ColumnarValue::Array(array_from_values(values))
        }
        (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(path_scalar)) => {
            let ScalarValue::Struct(variant_array) = scalar_variant else {
                return exec_err!("expected struct array");
            };

            let path = try_parse_string_scalar(path_scalar)?
                .map(|s| s.as_str())
                .unwrap_or_default();

            let variant_array = VariantArray::try_new(variant_array.as_ref())?;
            let value = variant_get_single_value(&variant_array, 0, path, path_mode, extract)?;

            ColumnarValue::Scalar(scalar_from_option(value))
        }
        (ColumnarValue::Array(variant_array), ColumnarValue::Array(paths)) => {
            if variant_array.len() != paths.len() {
                return exec_err!("expected variant array and paths to be of same length");
            }

            let paths = try_parse_string_columnar(paths)?;
            let variant_array = VariantArray::try_new(variant_array.as_ref())?;

            let values: Vec<Option<T>> = (0..variant_array.len())
                .map(|i| {
                    let path = paths[i].unwrap_or_default();
                    variant_get_single_value(&variant_array, i, path, path_mode, extract)
                })
                .collect::<Result<_>>()?;

            ColumnarValue::Array(array_from_values(values))
        }
        (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Array(paths)) => {
            let ScalarValue::Struct(variant_array) = scalar_variant else {
                return exec_err!("expected struct array");
            };

            let variant_array = VariantArray::try_new(variant_array.as_ref())?;
            let paths = try_parse_string_columnar(paths)?;

            let values: Vec<Option<T>> = paths
                .iter()
                .map(|path| {
                    let path = path.unwrap_or_default();
                    variant_get_single_value(&variant_array, 0, path, path_mode, extract)
                })
                .collect::<Result<_>>()?;

            ColumnarValue::Array(array_from_values(values))
        }
    };

    Ok(out)
}

/// This is similar to anyhow's ensure! macro
/// If the `pred` fails, it will return a DataFusionError
pub fn ensure(pred: bool, err_msg: &str) -> Result<()> {
    if !pred {
        return exec_err!("{}", err_msg);
    }

    Ok(())
}

/// Helper for argument count errors.
pub fn args_count_err(udf: &str, expected: &'static str, actual: usize) -> DataFusionError {
    DataFusionError::Execution(format!(
        "{udf}: expected {expected} argument(s), got {actual}"
    ))
}

/// Helper for argument type errors.
pub fn arg_type_err<T>(
    udf: &str,
    arg_index: u8,
    expected: &str,
    actual: &DataType,
) -> Result<T, DataFusionError> {
    Err(DataFusionError::Execution(format!(
        "{udf} arg #{arg_index}: expected {expected}, got {actual}"
    )))
}

/// Helper for unexpected NULL argument values.
pub fn arg_null_err<T>(udf: &str, arg_index: u8, expected: &str) -> Result<T, DataFusionError> {
    Err(arg_null_error(udf, arg_index, expected))
}

/// Helper for unexpected NULL argument values as a plain DataFusionError.
pub fn arg_null_error(udf: &str, arg_index: u8, expected: &str) -> DataFusionError {
    DataFusionError::Execution(format!(
        "{udf} arg #{arg_index}: expected {expected}, got NULL"
    ))
}

/// Helper for scalar/array shape mismatches.
pub fn arg_shape_err(udf: &str, arg_index: u8, expected: &str, actual: &str) -> DataFusionError {
    DataFusionError::Execution(format!(
        "{udf} arg #{arg_index}: expected {expected}, got {actual}"
    ))
}

/// Helper for invalid Variant kind in an argument.
pub fn arg_variant_kind_err(
    udf: &str,
    arg_index: u8,
    expected_variant_kind: &str,
) -> DataFusionError {
    DataFusionError::Execution(format!(
        "{udf} arg #{arg_index}: expected variant {expected_variant_kind}"
    ))
}

/// Helper for missing argument field metadata.
pub fn arg_field_meta_missing_err(udf: &str, arg_index: u8) -> DataFusionError {
    DataFusionError::Execution(format!("{udf} arg #{arg_index} field metadata is missing"))
}

// test related methods

#[cfg(test)]
pub fn build_variant_array_from_json(value: &serde_json::Value) -> VariantArray {
    let json_str = value.to_string();
    let mut builder = VariantArrayBuilder::new(1);
    builder.append_json(json_str.as_str()).unwrap();

    builder.build()
}

#[cfg(test)]
pub fn variant_scalar_from_json(json: serde_json::Value) -> ScalarValue {
    let mut builder = VariantArrayBuilder::new(1);
    builder.append_json(json.to_string().as_str()).unwrap();
    ScalarValue::Struct(Arc::new(builder.build().into()))
}

#[cfg(test)]
pub fn variant_array_from_json_rows(json_rows: &[serde_json::Value]) -> ArrayRef {
    let mut builder = VariantArrayBuilder::new(json_rows.len());
    for value in json_rows {
        builder.append_json(value.to_string().as_str()).unwrap();
    }
    let variant_array: StructArray = builder.build().into();
    Arc::new(variant_array) as ArrayRef
}

#[cfg(test)]
pub fn standard_variant_get_arg_fields() -> Vec<Arc<Field>> {
    vec![
        Arc::new(
            Field::new("input", DataType::Struct(Fields::empty()), true)
                .with_extension_type(VariantType),
        ),
        Arc::new(Field::new("path", DataType::Utf8, true)),
    ]
}

#[cfg(test)]
pub fn build_variant_get_args(
    variant_input: ColumnarValue,
    path: ColumnarValue,
    return_data_type: DataType,
    arg_fields: Vec<Arc<Field>>,
) -> ScalarFunctionArgs {
    ScalarFunctionArgs {
        args: vec![variant_input, path],
        return_field: Arc::new(Field::new("result", return_data_type, true)),
        arg_fields,
        number_rows: Default::default(),
        config_options: Default::default(),
    }
}

#[cfg(test)]
#[allow(unused)]
pub fn build_variant_array_from_json_array(jsons: &[Option<serde_json::Value>]) -> VariantArray {
    let mut builder = VariantArrayBuilder::new(jsons.len());

    jsons.iter().for_each(|v| match v.as_ref() {
        Some(json) => builder.append_json(json.to_string().as_str()).unwrap(),
        None => builder.append_null(),
    });

    builder.build()
}