//! JSON to Variant UDF for DataFusion.
//!
//! Implements `json_to_variant(json_string)` which parses a JSON string into
//! a native Arrow Variant value.
//!
//! Adapted from datafusion-variant crate.

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

use datafusion::arrow::array::{Array, ArrayRef, StringArray, StringViewArray};
use datafusion::arrow::datatypes::{DataType, Field};
use datafusion::common::{exec_datafusion_err, exec_err, Result, ScalarValue};

use super::{normalize_variant_struct, scalar_to_string, variant_data_type};
use datafusion::logical_expr::{
    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
    TypeSignature, Volatility,
};
use parquet_variant_compute::{VariantArrayBuilder, VariantType};
use rjiter::jiter::Jiter;

use crate::variant_walker::{write_subtree_to_array, WalkError};

/// UDF that converts a JSON string to a Variant value.
#[derive(Debug, Hash, PartialEq, Eq)]
pub struct JsonToVariantUdf {
    signature: Signature,
}

impl Default for JsonToVariantUdf {
    fn default() -> Self {
        Self {
            signature: Signature::new(
                TypeSignature::Uniform(
                    1,
                    vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View],
                ),
                Volatility::Immutable,
            ),
        }
    }
}

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

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

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

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

    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Arc<Field>> {
        Ok(Arc::new(
            Field::new(self.name(), variant_data_type(), true).with_extension_type(VariantType),
        ))
    }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let arg = args
            .args
            .first()
            .ok_or_else(|| exec_datafusion_err!("json_to_variant expects 1 argument"))?;

        match arg {
            ColumnarValue::Scalar(scalar_value) => {
                let json_str = scalar_to_string(scalar_value)?;

                let mut builder = VariantArrayBuilder::new(1);

                match json_str {
                    Some(json_str) => match write_one(json_str.as_str(), &mut builder) {
                        Ok(()) => {}
                        Err(_) => builder.append_null(),
                    },
                    None => builder.append_null(),
                }

                let struct_array = normalize_variant_struct(builder.build().into());
                Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(
                    struct_array,
                ))))
            }
            ColumnarValue::Array(arr) => match arr.data_type() {
                DataType::Utf8 => Ok(ColumnarValue::Array(from_string_array::<StringArray>(arr)?)),
                DataType::Utf8View => Ok(ColumnarValue::Array(
                    from_string_array::<StringViewArray>(arr)?,
                )),
                DataType::LargeUtf8 => Ok(ColumnarValue::Array(from_string_array::<
                    datafusion::arrow::array::LargeStringArray,
                >(arr)?)),
                _ => exec_err!(
                    "json_to_variant expects string input, got {}",
                    arr.data_type()
                ),
            },
        }
    }
}

/// Parse one JSON document straight into the builder via `Jiter`,
/// bypassing `serde_json::Value`. The previous implementation went
/// through `parquet_variant_json::JsonToVariant::append_json`, which
/// internally builds a `serde_json::Value` tree and then re-walks it
/// — flamegraphs of real ingest workloads showed ~78% of CPU in that
/// two-pass path. The shared walker emits Variant tokens directly
/// from parser events, with a single-pass and no `Value` allocation.
fn write_one(s: &str, out: &mut VariantArrayBuilder) -> std::result::Result<(), WalkError> {
    let mut j = Jiter::new(s.as_bytes());
    let peek = j.peek().map_err(WalkError::from)?;
    write_subtree_to_array(&mut j, out, peek)
}

/// Convert a string array to a VariantArray
fn from_string_array<T>(arr: &ArrayRef) -> Result<ArrayRef>
where
    T: Array + 'static,
    for<'a> &'a T: IntoIterator<Item = Option<&'a str>>,
{
    let arr = arr
        .as_any()
        .downcast_ref::<T>()
        .ok_or_else(|| exec_datafusion_err!("Unable to downcast array as expected string type"))?;

    let mut builder = VariantArrayBuilder::new(arr.len());

    for v in arr {
        match v {
            Some(json_str) => match write_one(json_str, &mut builder) {
                Ok(()) => {}
                Err(_) => builder.append_null(),
            },
            None => builder.append_null(),
        }
    }

    let variant_array = normalize_variant_struct(builder.build().into());
    Ok(Arc::new(variant_array) as ArrayRef)
}

/// Create the json_to_variant UDF
pub fn json_to_variant_udf() -> ScalarUDF {
    ScalarUDF::new_from_impl(JsonToVariantUdf::default())
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::array::Array;
    use parquet_variant::Variant;
    use parquet_variant_compute::VariantArray;

    #[test]
    fn test_json_to_variant_scalar_object() {
        let udf = JsonToVariantUdf::default();
        let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
        let return_field = udf
            .return_field_from_args(ReturnFieldArgs {
                arg_fields: &[arg_field.clone()],
                scalar_arguments: &[],
            })
            .unwrap();

        let json_input = ScalarValue::Utf8(Some(r#"{"name": "test", "value": 42}"#.to_string()));

        let args = ScalarFunctionArgs {
            args: vec![ColumnarValue::Scalar(json_input)],
            return_field,
            arg_fields: vec![arg_field],
            number_rows: 1,
            config_options: Default::default(),
        };

        let result = udf.invoke_with_args(args).unwrap();

        match result {
            ColumnarValue::Scalar(ScalarValue::Struct(struct_arr)) => {
                let variant_array = VariantArray::try_new(struct_arr.as_ref()).unwrap();
                let variant = variant_array.value(0);

                // Verify it's an object with the expected fields
                if let Variant::Object(obj) = variant {
                    assert!(obj.get("name").is_some());
                    assert!(obj.get("value").is_some());
                } else {
                    panic!("Expected Variant::Object, got {:?}", variant);
                }
            }
            _ => panic!("Expected scalar struct result"),
        }
    }

    #[test]
    fn test_json_to_variant_scalar_null() {
        let udf = JsonToVariantUdf::default();
        let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
        let return_field = udf
            .return_field_from_args(ReturnFieldArgs {
                arg_fields: &[arg_field.clone()],
                scalar_arguments: &[],
            })
            .unwrap();

        let json_input = ScalarValue::Utf8(None);

        let args = ScalarFunctionArgs {
            args: vec![ColumnarValue::Scalar(json_input)],
            return_field,
            arg_fields: vec![arg_field],
            number_rows: 1,
            config_options: Default::default(),
        };

        let result = udf.invoke_with_args(args).unwrap();

        match result {
            ColumnarValue::Scalar(ScalarValue::Struct(struct_arr)) => {
                assert!(struct_arr.is_null(0), "Expected null variant");
            }
            _ => panic!("Expected scalar struct result"),
        }
    }

    #[test]
    fn test_json_to_variant_array() {
        let udf = JsonToVariantUdf::default();
        let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
        let return_field = udf
            .return_field_from_args(ReturnFieldArgs {
                arg_fields: &[arg_field.clone()],
                scalar_arguments: &[],
            })
            .unwrap();

        let json_array: ArrayRef = Arc::new(StringArray::from(vec![
            Some(r#"{"a": 1}"#),
            Some(r#"{"a": 2}"#),
            None,
        ]));

        let args = ScalarFunctionArgs {
            args: vec![ColumnarValue::Array(json_array)],
            return_field,
            arg_fields: vec![arg_field],
            number_rows: 3,
            config_options: Default::default(),
        };

        let result = udf.invoke_with_args(args).unwrap();

        match result {
            ColumnarValue::Array(arr) => {
                assert_eq!(arr.len(), 3);
                assert!(!arr.is_null(0));
                assert!(!arr.is_null(1));
                assert!(arr.is_null(2));
            }
            _ => panic!("Expected array result"),
        }
    }

    /// Parse a realistic ~1KB Suricata-shaped payload and assert
    /// specific fields read back through the resulting Variant. This
    /// guards against regressions where the streaming Jiter walker
    /// produces a different binary shape (e.g. wrong number type,
    /// dropped key, mis-ordered object) than the prior
    /// `parquet_variant_json` path. Going through the public
    /// invoke_with_args ensures we exercise the same call site
    /// production hits.
    #[test]
    fn test_json_to_variant_realistic_payload_field_lookup() {
        let udf = JsonToVariantUdf::default();
        let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
        let return_field = udf
            .return_field_from_args(ReturnFieldArgs {
                arg_fields: &[arg_field.clone()],
                scalar_arguments: &[],
            })
            .unwrap();

        let json = r#"{
            "timestamp": "2024-09-19T12:34:56.789Z",
            "flow_id": 1234567890123456,
            "event_type": "alert",
            "src_ip": "10.0.0.1",
            "src_port": 49152,
            "dest_ip": "10.0.0.2",
            "dest_port": 443,
            "proto": "TCP",
            "alert": {
                "action": "allowed",
                "signature_id": 2027865,
                "rev": 3,
                "signature": "ET MALWARE Possible Malicious Activity",
                "severity": 1
            },
            "tags": ["malware", "stream", "tcp"],
            "stream_tcp": {
                "tcp_flags": "ap",
                "syn": false,
                "ack": true,
                "state": "established"
            }
        }"#;

        let json_input = ScalarValue::Utf8(Some(json.to_string()));
        let args = ScalarFunctionArgs {
            args: vec![ColumnarValue::Scalar(json_input)],
            return_field,
            arg_fields: vec![arg_field],
            number_rows: 1,
            config_options: Default::default(),
        };

        let result = udf.invoke_with_args(args).unwrap();
        let struct_arr = match result {
            ColumnarValue::Scalar(ScalarValue::Struct(s)) => s,
            other => panic!("expected scalar struct, got {:?}", other),
        };

        let variant_array = VariantArray::try_new(struct_arr.as_ref()).unwrap();
        let v = variant_array.value(0);
        let obj = v.as_object().expect("expected top-level object");

        // Top-level scalars present and correctly typed. Strings can
        // come back as either Variant::String or Variant::ShortString
        // depending on length; .as_string() handles both.
        assert_eq!(
            obj.get("event_type")
                .and_then(|v| v.as_string().map(String::from)),
            Some("alert".to_string()),
        );
        assert_eq!(obj.get("dest_port").and_then(|v| v.as_int64()), Some(443));
        assert_eq!(
            obj.get("flow_id").and_then(|v| v.as_int64()),
            Some(1_234_567_890_123_456),
        );

        // Nested object preserved.
        let alert = obj.get("alert").expect("alert present");
        let alert = alert.as_object().expect("alert is object");
        assert_eq!(
            alert
                .get("signature")
                .and_then(|v| v.as_string().map(String::from)),
            Some("ET MALWARE Possible Malicious Activity".to_string()),
        );
        assert_eq!(
            alert.get("signature_id").and_then(|v| v.as_int64()),
            Some(2_027_865),
        );

        // Array preserved with the right element count.
        let tags = obj.get("tags").expect("tags present");
        let tags = tags.as_list().expect("tags is list");
        assert_eq!(tags.len(), 3);

        // Nested object with booleans round-trips correctly.
        let stream_tcp = obj.get("stream_tcp").expect("stream_tcp present");
        let stream_tcp = stream_tcp.as_object().expect("stream_tcp is object");
        assert_eq!(
            stream_tcp.get("syn").and_then(|v| v.as_boolean()),
            Some(false)
        );
        assert_eq!(
            stream_tcp.get("ack").and_then(|v| v.as_boolean()),
            Some(true)
        );
    }
}
