//! Tests that string-accepting UDFs work with all three string types:
//! Utf8, LargeUtf8, and Utf8View.
//!
//! These tests invoke UDFs directly with different string array types
//! to ensure Parquet-originated data (which may use LargeUtf8 or Utf8View)
//! is handled correctly.

use std::sync::Arc;

use datafusion::arrow::array::{
    Array, ArrayRef, AsArray, BooleanArray, GenericStringBuilder, StringViewBuilder,
};
use datafusion::arrow::datatypes::{DataType, Field};
use datafusion::common::ScalarValue;
use datafusion::config::ConfigOptions;
use datafusion::logical_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl};

use super::ip_address::{CidrContainsUdf, IsIpv4Udf, IsIpv6Udf};
use super::parse_timestamp::ParseTimestampUdf;
use super::regexp_extract_all::RegexpExtractAllUdf;
use super::regexp_split::RegexpSplitUdf;
use super::variant_get::HamelinVariantGetUdf;

// ============================================================================
// Helpers
// ============================================================================

/// Build a Utf8 (StringArray) from string slices.
fn make_utf8(values: &[Option<&str>]) -> ArrayRef {
    let mut builder = GenericStringBuilder::<i32>::new();
    for v in values {
        match v {
            Some(s) => builder.append_value(s),
            None => builder.append_null(),
        }
    }
    Arc::new(builder.finish())
}

/// Build a LargeUtf8 (LargeStringArray) from string slices.
fn make_large_utf8(values: &[Option<&str>]) -> ArrayRef {
    let mut builder = GenericStringBuilder::<i64>::new();
    for v in values {
        match v {
            Some(s) => builder.append_value(s),
            None => builder.append_null(),
        }
    }
    Arc::new(builder.finish())
}

/// Build a Utf8View (StringViewArray) from string slices.
fn make_utf8_view(values: &[Option<&str>]) -> ArrayRef {
    let mut builder = StringViewBuilder::new();
    for v in values {
        match v {
            Some(s) => builder.append_value(s),
            None => builder.append_null(),
        }
    }
    Arc::new(builder.finish())
}

fn invoke_scalar(udf: &dyn ScalarUDFImpl, args: Vec<ColumnarValue>) -> ColumnarValue {
    let arg_fields: Vec<Arc<Field>> = args
        .iter()
        .map(|a| Arc::new(Field::new("arg", a.data_type(), true)))
        .collect();
    udf.invoke_with_args(ScalarFunctionArgs {
        args,
        arg_fields,
        number_rows: 1,
        return_field: Arc::new(Field::new("result", DataType::Null, true)),
        config_options: Arc::new(ConfigOptions::default()),
    })
    .expect("UDF invoke failed")
}

/// Like invoke_scalar but uses return_field_from_args (needed for variant_get).
fn invoke_variant_get(udf: &HamelinVariantGetUdf, args: Vec<ColumnarValue>) -> ColumnarValue {
    let arg_fields: Vec<Arc<Field>> = args
        .iter()
        .map(|a| Arc::new(Field::new("arg", a.data_type(), true)))
        .collect();
    let return_field = udf
        .return_field_from_args(ReturnFieldArgs {
            arg_fields: &[],
            scalar_arguments: &[],
        })
        .expect("return_field_from_args failed");
    udf.invoke_with_args(ScalarFunctionArgs {
        args,
        arg_fields,
        number_rows: 1,
        return_field,
        config_options: Arc::new(ConfigOptions::default()),
    })
    .expect("UDF invoke failed")
}

/// All three string array constructors, with labels for test output.
const STRING_TYPE_NAMES: [&str; 3] = ["Utf8", "LargeUtf8", "Utf8View"];

fn make_string_array(type_idx: usize, values: &[Option<&str>]) -> ArrayRef {
    match type_idx {
        0 => make_utf8(values),
        1 => make_large_utf8(values),
        2 => make_utf8_view(values),
        _ => unreachable!(),
    }
}

fn make_string_scalar(type_idx: usize, value: Option<&str>) -> ScalarValue {
    let s = value.map(|v| v.to_string());
    match type_idx {
        0 => ScalarValue::Utf8(s),
        1 => ScalarValue::LargeUtf8(s),
        2 => ScalarValue::Utf8View(s),
        _ => unreachable!(),
    }
}

// ============================================================================
// regexp_extract_all — all string types
// ============================================================================

#[test]
fn test_regexp_extract_all_scalar_string_types() {
    let udf = RegexpExtractAllUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        for (j, pat_type_name) in STRING_TYPE_NAMES.iter().enumerate() {
            let string_scalar = make_string_scalar(i, Some("abc123def456"));
            let pattern_scalar = make_string_scalar(j, Some("[0-9]+"));

            let result = invoke_scalar(
                &udf,
                vec![
                    ColumnarValue::Scalar(string_scalar),
                    ColumnarValue::Scalar(pattern_scalar),
                ],
            );

            let ColumnarValue::Scalar(ScalarValue::List(list)) = result else {
                panic!("{}/{}: expected List scalar", type_name, pat_type_name);
            };
            let arr = list.values();
            assert_eq!(
                arr.len(),
                2,
                "{}/{}: expected 2 matches",
                type_name,
                pat_type_name
            );
        }
    }
}

#[test]
fn test_regexp_extract_all_array_string_types() {
    let udf = RegexpExtractAllUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let string_arr = make_string_array(i, &[Some("cat bat rat"), None, Some("hello")]);
        let pattern_scalar = ScalarValue::Utf8(Some("[cbr]at".to_string()));

        let result = invoke_scalar(
            &udf,
            vec![
                ColumnarValue::Array(string_arr),
                ColumnarValue::Scalar(pattern_scalar),
            ],
        );

        let ColumnarValue::Array(arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        let list_arr = arr.as_list::<i32>();
        // Row 0: ["cat", "bat", "rat"] — 3 matches
        assert_eq!(
            list_arr.value(0).len(),
            3,
            "{}: row 0 should have 3 matches",
            type_name
        );
        // Row 1: null
        assert!(list_arr.is_null(1), "{}: row 1 should be null", type_name);
        // Row 2: [] — no matches
        assert_eq!(
            list_arr.value(2).len(),
            0,
            "{}: row 2 should be empty",
            type_name
        );
    }
}

// ============================================================================
// regexp_split — all string types
// ============================================================================

#[test]
fn test_regexp_split_scalar_string_types() {
    let udf = RegexpSplitUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        for (j, pat_type_name) in STRING_TYPE_NAMES.iter().enumerate() {
            let string_scalar = make_string_scalar(i, Some("a,b,c"));
            let pattern_scalar = make_string_scalar(j, Some(","));

            let result = invoke_scalar(
                &udf,
                vec![
                    ColumnarValue::Scalar(string_scalar),
                    ColumnarValue::Scalar(pattern_scalar),
                ],
            );

            let ColumnarValue::Scalar(ScalarValue::List(list)) = result else {
                panic!("{}/{}: expected List scalar", type_name, pat_type_name);
            };
            let arr = list.values();
            assert_eq!(
                arr.len(),
                3,
                "{}/{}: expected 3 parts",
                type_name,
                pat_type_name
            );
        }
    }
}

#[test]
fn test_regexp_split_array_string_types() {
    let udf = RegexpSplitUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let string_arr = make_string_array(i, &[Some("a-b-c"), None, Some("hello")]);
        let pattern_scalar = ScalarValue::Utf8(Some("-".to_string()));

        let result = invoke_scalar(
            &udf,
            vec![
                ColumnarValue::Array(string_arr),
                ColumnarValue::Scalar(pattern_scalar),
            ],
        );

        let ColumnarValue::Array(arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        let list_arr = arr.as_list::<i32>();
        // Row 0: ["a", "b", "c"]
        assert_eq!(
            list_arr.value(0).len(),
            3,
            "{}: row 0 should have 3 parts",
            type_name
        );
        // Row 1: null
        assert!(list_arr.is_null(1), "{}: row 1 should be null", type_name);
        // Row 2: ["hello"] — no split
        assert_eq!(
            list_arr.value(2).len(),
            1,
            "{}: row 2 should have 1 part",
            type_name
        );
    }
}

#[test]
fn test_regexp_split_scalar_null_returns_null() {
    let udf = RegexpSplitUdf::new();

    // Null string
    let result = invoke_scalar(
        &udf,
        vec![
            ColumnarValue::Scalar(ScalarValue::Utf8(None)),
            ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))),
        ],
    );
    let ColumnarValue::Scalar(sv) = result else {
        panic!("expected scalar result");
    };
    assert!(sv.is_null(), "null string should produce null result");

    // Null pattern
    let result = invoke_scalar(
        &udf,
        vec![
            ColumnarValue::Scalar(ScalarValue::Utf8(Some("a,b".to_string()))),
            ColumnarValue::Scalar(ScalarValue::Utf8(None)),
        ],
    );
    let ColumnarValue::Scalar(sv) = result else {
        panic!("expected scalar result");
    };
    assert!(sv.is_null(), "null pattern should produce null result");
}

#[test]
fn test_regexp_split_invalid_pattern_errors() {
    let udf = RegexpSplitUdf::new();

    // Invalid regex in scalar path
    let result = udf.invoke_with_args(ScalarFunctionArgs {
        args: vec![
            ColumnarValue::Scalar(ScalarValue::Utf8(Some("hello".to_string()))),
            ColumnarValue::Scalar(ScalarValue::Utf8(Some("[invalid".to_string()))),
        ],
        arg_fields: vec![
            Arc::new(Field::new("a", DataType::Utf8, true)),
            Arc::new(Field::new("b", DataType::Utf8, true)),
        ],
        number_rows: 1,
        return_field: Arc::new(Field::new("result", DataType::Null, true)),
        config_options: Arc::new(ConfigOptions::default()),
    });
    assert!(result.is_err(), "invalid regex should error");

    // Invalid regex in array path (scalar pattern)
    let result = udf.invoke_with_args(ScalarFunctionArgs {
        args: vec![
            ColumnarValue::Array(make_utf8(&[Some("hello")])),
            ColumnarValue::Scalar(ScalarValue::Utf8(Some("[invalid".to_string()))),
        ],
        arg_fields: vec![
            Arc::new(Field::new("a", DataType::Utf8, true)),
            Arc::new(Field::new("b", DataType::Utf8, true)),
        ],
        number_rows: 1,
        return_field: Arc::new(Field::new("result", DataType::Null, true)),
        config_options: Arc::new(ConfigOptions::default()),
    });
    assert!(result.is_err(), "invalid regex in array path should error");
}

// ============================================================================
// parse_timestamp — all string types
// ============================================================================

#[test]
fn test_parse_timestamp_scalar_string_types() {
    let udf = ParseTimestampUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let scalar = make_string_scalar(i, Some("2024-01-15T10:30:00Z"));

        let result = invoke_scalar(&udf, vec![ColumnarValue::Scalar(scalar)]);

        let ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(Some(_micros), Some(tz))) =
            result
        else {
            panic!("{}: expected TimestampMicrosecond scalar", type_name);
        };
        assert_eq!(
            tz.as_ref(),
            "+00:00",
            "{}: timezone should be +00:00",
            type_name
        );
    }
}

#[test]
fn test_parse_timestamp_array_string_types() {
    let udf = ParseTimestampUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let arr = make_string_array(i, &[Some("2024-01-15T10:30:00Z"), None, Some("2024-06-01")]);

        let result = invoke_scalar(&udf, vec![ColumnarValue::Array(arr)]);

        let ColumnarValue::Array(result_arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        assert_eq!(result_arr.len(), 3, "{}: expected 3 rows", type_name);
        assert!(
            !result_arr.is_null(0),
            "{}: row 0 should not be null",
            type_name
        );
        assert!(result_arr.is_null(1), "{}: row 1 should be null", type_name);
        assert!(
            !result_arr.is_null(2),
            "{}: row 2 should not be null",
            type_name
        );
    }
}

// ============================================================================
// is_ipv4 — all string types
// ============================================================================

#[test]
fn test_is_ipv4_scalar_string_types() {
    let udf = IsIpv4Udf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let scalar = make_string_scalar(i, Some("192.168.1.1"));

        let result = invoke_scalar(&udf, vec![ColumnarValue::Scalar(scalar)]);

        let ColumnarValue::Scalar(ScalarValue::Boolean(Some(val))) = result else {
            panic!("{}: expected Boolean scalar", type_name);
        };
        assert!(val, "{}: 192.168.1.1 should be IPv4", type_name);
    }
}

#[test]
fn test_is_ipv4_array_string_types() {
    let udf = IsIpv4Udf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let arr = make_string_array(i, &[Some("192.168.1.1"), None, Some("::1")]);

        let result = invoke_scalar(&udf, vec![ColumnarValue::Array(arr)]);

        let ColumnarValue::Array(result_arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        let bool_arr = result_arr
            .as_any()
            .downcast_ref::<BooleanArray>()
            .expect("Expected BooleanArray");
        assert_eq!(
            bool_arr.value(0),
            true,
            "{}: row 0 should be true",
            type_name
        );
        assert!(bool_arr.is_null(1), "{}: row 1 should be null", type_name);
        assert_eq!(
            bool_arr.value(2),
            false,
            "{}: row 2 should be false",
            type_name
        );
    }
}

// ============================================================================
// is_ipv6 — all string types
// ============================================================================

#[test]
fn test_is_ipv6_scalar_string_types() {
    let udf = IsIpv6Udf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let scalar = make_string_scalar(i, Some("::1"));

        let result = invoke_scalar(&udf, vec![ColumnarValue::Scalar(scalar)]);

        let ColumnarValue::Scalar(ScalarValue::Boolean(Some(val))) = result else {
            panic!("{}: expected Boolean scalar", type_name);
        };
        assert!(val, "{}: ::1 should be IPv6", type_name);
    }
}

#[test]
fn test_is_ipv6_array_string_types() {
    let udf = IsIpv6Udf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let arr = make_string_array(i, &[Some("::1"), None, Some("192.168.1.1")]);

        let result = invoke_scalar(&udf, vec![ColumnarValue::Array(arr)]);

        let ColumnarValue::Array(result_arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        let bool_arr = result_arr
            .as_any()
            .downcast_ref::<BooleanArray>()
            .expect("Expected BooleanArray");
        assert_eq!(
            bool_arr.value(0),
            true,
            "{}: row 0 should be true",
            type_name
        );
        assert!(bool_arr.is_null(1), "{}: row 1 should be null", type_name);
        assert_eq!(
            bool_arr.value(2),
            false,
            "{}: row 2 should be false",
            type_name
        );
    }
}

// ============================================================================
// cidr_contains — all string types
// ============================================================================

#[test]
fn test_cidr_contains_scalar_string_types() {
    let udf = CidrContainsUdf::new();

    for (i, cidr_type) in STRING_TYPE_NAMES.iter().enumerate() {
        for (j, ip_type) in STRING_TYPE_NAMES.iter().enumerate() {
            let cidr_scalar = make_string_scalar(i, Some("10.0.0.0/8"));
            let ip_scalar = make_string_scalar(j, Some("10.1.2.3"));

            let result = invoke_scalar(
                &udf,
                vec![
                    ColumnarValue::Scalar(cidr_scalar),
                    ColumnarValue::Scalar(ip_scalar),
                ],
            );

            let ColumnarValue::Scalar(ScalarValue::Boolean(Some(val))) = result else {
                panic!("{}/{}: expected Boolean scalar", cidr_type, ip_type);
            };
            assert!(
                val,
                "{}/{}: 10.1.2.3 should be in 10.0.0.0/8",
                cidr_type, ip_type
            );
        }
    }
}

#[test]
fn test_cidr_contains_array_string_types() {
    let udf = CidrContainsUdf::new();

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let cidr_scalar = ScalarValue::Utf8(Some("10.0.0.0/8".to_string()));
        let ip_arr = make_string_array(i, &[Some("10.1.2.3"), None, Some("192.168.1.1")]);

        let result = invoke_scalar(
            &udf,
            vec![
                ColumnarValue::Scalar(cidr_scalar),
                ColumnarValue::Array(ip_arr),
            ],
        );

        let ColumnarValue::Array(result_arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        let bool_arr = result_arr
            .as_any()
            .downcast_ref::<BooleanArray>()
            .expect("Expected BooleanArray");
        assert_eq!(
            bool_arr.value(0),
            true,
            "{}: row 0 should be true",
            type_name
        );
        assert!(bool_arr.is_null(1), "{}: row 1 should be null", type_name);
        assert_eq!(
            bool_arr.value(2),
            false,
            "{}: row 2 should be false",
            type_name
        );
    }
}

// ============================================================================
// variant_get — all string types for path argument
// ============================================================================

#[test]
fn test_variant_get_scalar_path_string_types() {
    use parquet_variant_compute::VariantArrayBuilder;
    use parquet_variant_json::JsonToVariant;

    let udf = HamelinVariantGetUdf::default();

    // Build a variant scalar
    let mut builder = VariantArrayBuilder::new(1);
    builder
        .append_json(r#"{"foo": 42}"#)
        .expect("Failed to build variant");
    let struct_array: datafusion::arrow::array::StructArray = builder.build().into();
    let variant_scalar = ScalarValue::Struct(Arc::new(struct_array));

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let path_scalar = make_string_scalar(i, Some("foo"));

        let result = invoke_variant_get(
            &udf,
            vec![
                ColumnarValue::Scalar(variant_scalar.clone()),
                ColumnarValue::Scalar(path_scalar),
            ],
        );

        assert!(
            matches!(result, ColumnarValue::Scalar(ScalarValue::Struct(_))),
            "{}: expected Struct scalar for variant result",
            type_name
        );
    }
}

#[test]
fn test_variant_get_array_path_string_types() {
    use parquet_variant_compute::VariantArrayBuilder;
    use parquet_variant_json::JsonToVariant;

    let udf = HamelinVariantGetUdf::default();

    // Build a variant scalar
    let mut builder = VariantArrayBuilder::new(1);
    builder
        .append_json(r#"{"foo": 1, "bar": 2}"#)
        .expect("Failed to build variant");
    let struct_array: datafusion::arrow::array::StructArray = builder.build().into();
    let variant_scalar = ScalarValue::Struct(Arc::new(struct_array));

    for (i, type_name) in STRING_TYPE_NAMES.iter().enumerate() {
        let path_arr = make_string_array(i, &[Some("foo"), Some("bar"), None]);

        let result = invoke_variant_get(
            &udf,
            vec![
                ColumnarValue::Scalar(variant_scalar.clone()),
                ColumnarValue::Array(path_arr),
            ],
        );

        let ColumnarValue::Array(result_arr) = result else {
            panic!("{}: expected array result", type_name);
        };
        assert_eq!(result_arr.len(), 3, "{}: expected 3 rows", type_name);
        // Rows 0 and 1 should be non-null (valid paths), row 2 should be null (null path)
        assert!(
            !result_arr.is_null(0),
            "{}: row 0 should not be null",
            type_name
        );
        assert!(
            !result_arr.is_null(1),
            "{}: row 1 should not be null",
            type_name
        );
        assert!(result_arr.is_null(2), "{}: row 2 should be null", type_name);
    }
}
