//! Custom DataFusion UDFs for Hamelin language functions.
//!
//! These UDFs implement Hamelin functions that don't have direct equivalents in DataFusion.
//! They are used by translation functions in the function registry.

use std::fmt::Write;
use std::sync::Arc;

use datafusion::arrow::array::{Array, ArrayRef, MapArray, StructArray};
use datafusion::arrow::datatypes::{DataType, Field, Fields};
use datafusion::common::{exec_err, Result};
use parquet_variant::Variant;

mod any_value;
mod array_cast;
mod array_coalesce_default;
mod array_get_field;
mod array_helpers;
mod array_length;
mod array_variant_get;
mod array_variant_to_json;
mod array_zip;
mod filter_null;
mod from_variant;
mod interval_helper;
mod ip_address;
mod json_to_variant;
mod map_agg;
mod map_extract;
mod map_from_entries;
mod map_variant_get_values;
mod map_variant_to_json_values;
mod multimap_agg;
mod parse_timestamp;
mod regexp_extract_all;
mod regexp_split;
mod sliding_array_agg;
mod string_utils;
mod to_variant;
mod url;
mod uuid5;
mod variant_get;
mod variant_to_json;
mod width_bucket;

/// The canonical Variant struct schema.
///
/// The `value` field is declared nullable because `parquet_variant_compute`'s
/// `VariantArray::from_parts` (used by the `variant_get` kernel) marks it so,
/// while `VariantArrayBuilder::build` marks it non-nullable. We normalize to
/// nullable so DataFusion's return-type assertion is satisfied regardless of
/// which code path produced the result.
pub fn variant_fields() -> Fields {
    Fields::from(vec![
        Field::new("metadata", DataType::BinaryView, false),
        Field::new("value", DataType::BinaryView, true),
    ])
}

/// The canonical Variant struct `DataType`.
pub fn variant_data_type() -> DataType {
    DataType::Struct(variant_fields())
}

/// Check whether an Arrow `DataType` represents a Variant struct.
///
/// Matches both unshredded variants (`{metadata: BinaryView, value: BinaryView}`)
/// and shredded variants (`{metadata: BinaryView, value: BinaryView, typed_value: ...}`).
pub fn is_variant_data_type(dt: &DataType) -> bool {
    if let DataType::Struct(fields) = dt {
        let has_metadata = fields
            .iter()
            .any(|f| f.name() == "metadata" && f.data_type() == &DataType::BinaryView);
        let has_value = fields
            .iter()
            .any(|f| f.name() == "value" && f.data_type() == &DataType::BinaryView);
        has_metadata && has_value
    } else {
        false
    }
}

/// Re-wrap a StructArray so its field nullability matches [`variant_fields`].
///
/// This is a zero-copy operation — only the schema metadata changes, not the
/// underlying buffers.
pub fn normalize_variant_struct(struct_arr: StructArray) -> StructArray {
    let expected = variant_fields();
    if struct_arr.fields() == &expected {
        return struct_arr;
    }
    StructArray::new(
        expected,
        struct_arr.columns().to_vec(),
        struct_arr.nulls().cloned(),
    )
}

/// Extract the string value from a scalar, accepting all string types (Utf8, Utf8View, LargeUtf8).
///
/// Returns `Ok(None)` for null strings, `Ok(Some(s))` for non-null strings,
/// and `Err` for non-string types.
pub fn scalar_to_string(
    scalar: &datafusion::common::ScalarValue,
) -> datafusion::common::Result<Option<String>> {
    use datafusion::common::ScalarValue;
    match scalar {
        ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) | ScalarValue::LargeUtf8(s) => {
            Ok(s.clone())
        }
        other => datafusion::common::exec_err!("expected string scalar, got {}", other.data_type()),
    }
}

/// Rebuild a `MapArray` by replacing its values column, preserving keys,
/// offsets, and null bitmap.
///
/// Used by the vectorized map-value UDFs (`hamelin_map_variant_get_values`,
/// `hamelin_map_variant_to_json_values`) to re-wrap transformed values without
/// touching the original structure.
pub fn rebuild_map(map_array: &MapArray, keys: ArrayRef, new_values: ArrayRef) -> Result<ArrayRef> {
    let DataType::Map(field, sorted) = map_array.data_type() else {
        return exec_err!("map_array.data_type() must be DataType::Map");
    };
    let DataType::Struct(entry_fields) = field.data_type() else {
        return exec_err!("expected Struct inside Map");
    };
    let key_field = entry_fields[0].clone();
    let value_field = Arc::new(Field::new(
        entry_fields[1].name(),
        new_values.data_type().clone(),
        entry_fields[1].is_nullable(),
    ));
    let new_struct_fields = Fields::from(vec![key_field, value_field]);
    let new_entries = StructArray::new(
        new_struct_fields.clone(),
        vec![keys, new_values],
        map_array.entries().nulls().cloned(),
    );

    let new_map_field = Arc::new(Field::new(
        field.name(),
        DataType::Struct(new_struct_fields),
        field.is_nullable(),
    ));

    let result = MapArray::new(
        new_map_field,
        map_array.offsets().clone(),
        new_entries,
        map_array.nulls().cloned(),
        *sorted,
    );
    Ok(Arc::new(result) as ArrayRef)
}

/// Render a scalar Variant to its bare display form (no JSON quoting).
///
/// - Strings and short strings write their unwrapped characters (`hello`, not `"hello"`).
/// - Numeric, boolean, decimal, and temporal scalars write their plain textual form.
/// - Null / Binary / List / Object variants and non-finite floats (NaN / ±Inf)
///   leave `buf` unchanged. Callers must check for "no display form" (e.g. via
///   `append_variant_to_stringview`) and emit a SQL NULL in that case —
///   `buf` being empty is not itself a signal, since `Variant::String("")` is
///   a valid empty-string value.
///
/// Used by `FromVariantUdf` when the target is a native string, and as the
/// first leg of the variant-string double-cast path into numeric / boolean /
/// timestamp targets.
pub fn write_variant_display(buf: &mut String, variant: &Variant) -> std::fmt::Result {
    match variant {
        Variant::String(s) => buf.push_str(s),
        Variant::ShortString(s) => buf.push_str(s.as_ref()),
        Variant::BooleanTrue => buf.push_str("true"),
        Variant::BooleanFalse => buf.push_str("false"),
        Variant::Int8(n) => write!(buf, "{n}")?,
        Variant::Int16(n) => write!(buf, "{n}")?,
        Variant::Int32(n) => write!(buf, "{n}")?,
        Variant::Int64(n) => write!(buf, "{n}")?,
        Variant::Float(f) if f.is_finite() => write!(buf, "{f}")?,
        Variant::Double(f) if f.is_finite() => write!(buf, "{f}")?,
        Variant::Decimal4(d) => write_decimal(buf, d.integer() as i128, d.scale())?,
        Variant::Decimal8(d) => write_decimal(buf, d.integer() as i128, d.scale())?,
        Variant::Decimal16(d) => write_decimal(buf, d.integer(), d.scale())?,
        Variant::Date(d) => write!(buf, "{d}")?,
        Variant::TimestampMicros(ts) => write!(buf, "{}", ts.to_rfc3339())?,
        Variant::TimestampNtzMicros(ts) => write!(buf, "{ts}")?,
        Variant::TimestampNanos(ts) => write!(buf, "{}", ts.to_rfc3339())?,
        Variant::TimestampNtzNanos(ts) => write!(buf, "{ts}")?,
        Variant::Time(t) => write!(buf, "{t}")?,
        Variant::Uuid(u) => write!(buf, "{u}")?,
        // Binary / Null / Object / List have no defined Hamelin primitive-to-string cast.
        // Non-finite Float/Double fall here too (NaN / ±Inf have no display form).
        _ => {}
    }
    Ok(())
}

/// Write a Variant's JSON representation into an existing buffer, avoiding intermediate
/// allocations for nested structures. Shared by the scalar `hamelin_to_json_string` UDF
/// and the vectorized `hamelin_array_variant_to_json` UDF.
pub fn write_variant(buf: &mut String, variant: &Variant) -> std::fmt::Result {
    match variant {
        Variant::Null => buf.push_str("null"),
        Variant::BooleanTrue => buf.push_str("true"),
        Variant::BooleanFalse => buf.push_str("false"),
        Variant::Int8(n) => write!(buf, "{n}")?,
        Variant::Int16(n) => write!(buf, "{n}")?,
        Variant::Int32(n) => write!(buf, "{n}")?,
        Variant::Int64(n) => write!(buf, "{n}")?,
        Variant::Float(f) => write_float(buf, *f as f64)?,
        Variant::Double(f) => write_float(buf, *f)?,
        Variant::Decimal4(d) => write_decimal(buf, d.integer() as i128, d.scale())?,
        Variant::Decimal8(d) => write_decimal(buf, d.integer() as i128, d.scale())?,
        Variant::Decimal16(d) => write_decimal(buf, d.integer(), d.scale())?,
        Variant::String(s) => write_json_string(buf, s)?,
        Variant::ShortString(s) => write_json_string(buf, s.as_ref())?,
        Variant::Binary(b) => {
            buf.push('"');
            write_hex_bytes(buf, b.as_ref());
            buf.push('"');
        }
        Variant::Date(d) => write!(buf, "\"{d}\"")?,
        Variant::TimestampMicros(ts) => write!(buf, "\"{}\"", ts.to_rfc3339())?,
        Variant::TimestampNtzMicros(ts) => write!(buf, "\"{ts}\"")?,
        Variant::TimestampNanos(ts) => write!(buf, "\"{}\"", ts.to_rfc3339())?,
        Variant::TimestampNtzNanos(ts) => write!(buf, "\"{ts}\"")?,
        Variant::Time(t) => write!(buf, "\"{t}\"")?,
        Variant::Uuid(u) => write!(buf, "\"{u}\"")?,
        Variant::List(list) => {
            buf.push('[');
            for (i, v) in list.iter().enumerate() {
                if i > 0 {
                    buf.push(',');
                }
                write_variant(buf, &v)?;
            }
            buf.push(']');
        }
        Variant::Object(obj) => {
            buf.push('{');
            for (i, (k, v)) in obj.iter().enumerate() {
                if i > 0 {
                    buf.push(',');
                }
                write_json_string(buf, k)?;
                buf.push(':');
                write_variant(buf, &v)?;
            }
            buf.push('}');
        }
    }
    Ok(())
}

/// Format a decimal value preserving trailing zeros (matching Trino's JSON serialization).
fn write_decimal(buf: &mut String, integer: i128, scale: u8) -> std::fmt::Result {
    if scale == 0 {
        return write!(buf, "{integer}");
    }
    let divisor = 10_i128.pow(scale as u32);
    let sign = if integer < 0 { "-" } else { "" };
    let abs = integer.unsigned_abs();
    let quotient = abs / divisor as u128;
    let remainder = abs % divisor as u128;
    write!(
        buf,
        "{sign}{quotient}.{remainder:0>width$}",
        width = scale as usize
    )
}

fn write_hex_bytes(buf: &mut String, bytes: &[u8]) {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    buf.reserve(bytes.len().saturating_mul(2));
    for &byte in bytes {
        buf.push(HEX[(byte >> 4) as usize] as char);
        buf.push(HEX[(byte & 0x0f) as usize] as char);
    }
}

fn write_float(buf: &mut String, f: f64) -> std::fmt::Result {
    if f.is_nan() || f.is_infinite() {
        buf.push_str("null");
    } else {
        write!(buf, "{f}")?;
    }
    Ok(())
}

fn write_json_string(buf: &mut String, s: &str) -> std::fmt::Result {
    buf.push('"');
    // Fast path: no escaping needed for common ASCII/UTF-8 strings.
    if !s
        .as_bytes()
        .iter()
        .any(|&b| b == b'"' || b == b'\\' || b < 0x20)
    {
        buf.push_str(s);
        buf.push('"');
        return Ok(());
    }

    for c in s.chars() {
        match c {
            '"' => buf.push_str("\\\""),
            '\\' => buf.push_str("\\\\"),
            '\n' => buf.push_str("\\n"),
            '\r' => buf.push_str("\\r"),
            '\t' => buf.push_str("\\t"),
            c if c.is_control() => write!(buf, "\\u{:04x}", c as u32)?,
            c => buf.push(c),
        }
    }
    buf.push('"');
    Ok(())
}

// Any value aggregate.
pub use any_value::any_value_udaf;

// Array cast for complex element transformations.
pub use array_cast::{array_cast_udf, ArrayCastUdf, CastDescriptor};

// Element-wise coalesce over a list array with a scalar default.
pub use array_coalesce_default::array_coalesce_default_udf;

// Array helpers.
pub use array_helpers::array_avg_udf;
pub use array_helpers::array_sum_udf;

// Fast array length (offsets-based, replaces DataFusion's built-in).
pub use array_length::hamelin_array_length_udf;

// Vectorized struct field extraction over list arrays.
pub use array_get_field::array_get_field_udf;

// Drop null elements from each row of a list array.
pub use filter_null::filter_null_udf;

// Vectorized variant field extraction over list arrays.
pub use array_variant_get::array_variant_get_udf;

// Vectorized variant-to-JSON serialization over list arrays.
pub use array_variant_to_json::array_variant_to_json_udf;

// Array zip.
pub use array_zip::array_zip_udf;

// Vectorized map-value operators.
pub use map_variant_get_values::map_variant_get_values_udf;
pub use map_variant_to_json_values::map_variant_to_json_values_udf;

// Map aggregate functions.
pub use map_agg::map_agg_udaf;
pub use multimap_agg::multimap_agg_udaf;

// Map scalar functions.
pub use map_extract::map_extract_udf;
pub use map_from_entries::map_from_entries_udf;

// Sliding window aggregate functions.
pub use sliding_array_agg::sliding_array_agg_udaf;

// Interval helpers.
pub use interval_helper::from_millis_udf;
pub use interval_helper::from_months_udf;
pub use interval_helper::from_nanos_udf;
pub use interval_helper::to_millis_udf;
pub use interval_helper::to_months_udf;
pub use interval_helper::to_nanos_udf;

// IP Address functions.
pub use ip_address::cidr_contains_udf;
pub use ip_address::is_ipv4_udf;
pub use ip_address::is_ipv6_udf;

// Variant to JSON.
pub use variant_to_json::variant_to_json_udf;

// Variant/JSON functions.
pub use from_variant::{from_variant_udf, FromVariantUdf};
pub use json_to_variant::json_to_variant_udf;
pub use to_variant::cast_to_variant_udf;
pub use variant_get::variant_get_udf;

// Width bucket.
pub use width_bucket::width_bucket_array_udf;

// Timestamp parsing.
pub use parse_timestamp::{parse_timestamp_micros, parse_timestamp_udf};

// Regex functions.
pub use regexp_extract_all::regexp_extract_all_udf;
pub use regexp_split::regexp_split_udf;

// URL encode/decode (Trino-compatible).
pub use url::{url_decode_udf, url_encode_udf};

// UUID5 (RFC 4122).
pub use uuid5::uuid5_udf;

#[cfg(test)]
mod large_list_tests;
#[cfg(test)]
mod string_type_tests;
