wp-arrow 0.4.2

use arrow::array::RecordBatch;
use arrow::ipc::reader::StreamReader;
use arrow::ipc::writer::StreamWriter;

use crate::error::WpArrowError;

/// Decoded IPC data frame.
pub struct IpcFrame {
    pub tag: String,
    pub batch: RecordBatch,
}

/// Encode a RecordBatch into an IPC frame: `[4B tag_len BE][tag bytes][Arrow IPC stream]`.
pub fn encode_ipc(tag: &str, batch: &RecordBatch) -> Result<Vec<u8>, WpArrowError> {
    let tag_bytes = tag.as_bytes();
    let tag_len = tag_bytes.len() as u32;

    let mut buf = Vec::new();
    buf.extend_from_slice(&tag_len.to_be_bytes());
    buf.extend_from_slice(tag_bytes);

    let mut writer = StreamWriter::try_new(&mut buf, batch.schema().as_ref())
        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;
    writer
        .write(batch)
        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;
    writer
        .finish()
        .map_err(|e| WpArrowError::IpcEncodeError(e.to_string()))?;

    Ok(buf)
}

/// Decode a complete IPC frame from bytes.
pub fn decode_ipc(data: &[u8]) -> Result<IpcFrame, WpArrowError> {
    if data.len() < 4 {
        return Err(WpArrowError::IpcDecodeError(format!(
            "frame too short: {} bytes, minimum 4",
            data.len()
        )));
    }

    let tag_len = u32::from_be_bytes(data[0..4].try_into().unwrap()) as usize;
    let tag_end = 4 + tag_len;
    if data.len() < tag_end {
        return Err(WpArrowError::IpcDecodeError(format!(
            "frame truncated: tag_len={tag_len} but only {} bytes remain after header",
            data.len() - 4
        )));
    }

    let tag = String::from_utf8(data[4..tag_end].to_vec())
        .map_err(|e| WpArrowError::IpcDecodeError(format!("invalid UTF-8 in tag: {e}")))?;

    let ipc_payload = &data[tag_end..];
    let mut reader = StreamReader::try_new(ipc_payload, None)
        .map_err(|e| WpArrowError::IpcDecodeError(e.to_string()))?;

    let batch = reader
        .next()
        .ok_or_else(|| WpArrowError::IpcDecodeError("no RecordBatch in IPC payload".to_string()))?
        .map_err(|e| WpArrowError::IpcDecodeError(e.to_string()))?;

    Ok(IpcFrame { tag, batch })
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Int32Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use std::sync::Arc;

    fn make_batch(num_rows: usize) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("name", DataType::Utf8, true),
        ]));
        let ids: Vec<i32> = (0..num_rows as i32).collect();
        let names: Vec<Option<&str>> = (0..num_rows)
            .map(|i| if i % 2 == 0 { Some("even") } else { None })
            .collect();
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Int32Array::from(ids)),
                Arc::new(StringArray::from(names)),
            ],
        )
        .unwrap()
    }

    #[test]
    fn ipc_roundtrip_basic() {
        let batch = make_batch(5);
        let encoded = encode_ipc("test-tag", &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.tag, "test-tag");
        assert_eq!(frame.batch.num_rows(), batch.num_rows());
        assert_eq!(frame.batch.num_columns(), batch.num_columns());
        assert_eq!(frame.batch.schema(), batch.schema());
        assert_eq!(frame.batch, batch);
    }

    #[test]
    fn ipc_roundtrip_empty_batch() {
        let batch = make_batch(0);
        let encoded = encode_ipc("empty", &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.batch.num_rows(), 0);
        assert_eq!(frame.batch, batch);
    }

    #[test]
    fn ipc_roundtrip_large_batch() {
        let batch = make_batch(1000);
        let encoded = encode_ipc("large", &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.batch.num_rows(), 1000);
        assert_eq!(frame.batch, batch);
    }

    #[test]
    fn ipc_tag_preserved() {
        let batch = make_batch(1);
        let encoded = encode_ipc("my-tag", &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.tag, "my-tag");
    }

    #[test]
    fn ipc_utf8_tag() {
        let batch = make_batch(1);
        let tag = "数据标签-🚀";
        let encoded = encode_ipc(tag, &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.tag, tag);
    }

    #[test]
    fn ipc_empty_tag() {
        let batch = make_batch(1);
        let encoded = encode_ipc("", &batch).unwrap();
        let frame = decode_ipc(&encoded).unwrap();
        assert_eq!(frame.tag, "");
    }

    #[test]
    fn decode_ipc_too_short() {
        let result = decode_ipc(&[0x00; 2]);
        assert!(matches!(result, Err(WpArrowError::IpcDecodeError(_))));
    }

    #[test]
    fn decode_ipc_truncated_tag() {
        let mut data = Vec::new();
        data.extend_from_slice(&100u32.to_be_bytes()); // tag_len = 100 but no data
        let result = decode_ipc(&data);
        assert!(matches!(result, Err(WpArrowError::IpcDecodeError(_))));
    }
}

#[test]
fn decode_roundtrip_preserves_memory() {
    use crate::ipc::{decode_ipc, encode_ipc};
    use arrow::array::{Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use std::sync::Arc;

    let n = 75000;
    let vals: Vec<Option<String>> = (0..n)
        .map(|i| Some(format!("10.0.{}.{}", i % 4, i % 250 + 1)))
        .collect();
    let schema = Arc::new(Schema::new(vec![Field::new("sip", DataType::Utf8, true)]));
    let batch = RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(vals))]).unwrap();
    eprintln!(
        "original col bytes = {}",
        batch.column(0).get_array_memory_size()
    );

    let enc = encode_ipc("conn_events", &batch).unwrap();
    eprintln!("ipc payload bytes = {}", enc.len());
    let frame = decode_ipc(&enc).unwrap();
    eprintln!(
        "decoded col bytes = {}",
        frame.batch.column(0).get_array_memory_size()
    );
}

// 12 列批次（模拟 wfgen conn 批次）往返内存
#[test]
fn decode_roundtrip_12col_memory() {
    use crate::ipc::{decode_ipc, encode_ipc};
    use arrow::array::{Array, Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use std::sync::Arc;

    let n = 75000;
    let str_vals: Vec<Option<String>> = (0..n)
        .map(|i| Some(format!("10.0.{}.{}", i % 4, i % 250 + 1)))
        .collect();
    let int_vals: Vec<i64> = (0..n).map(|i| (i % 4096) as i64).collect();
    let schema = Arc::new(Schema::new(vec![
        Field::new("sip", DataType::Utf8, true),
        Field::new("dip", DataType::Utf8, true),
        Field::new("roles_obj", DataType::Utf8, true),
        Field::new("dport", DataType::Int64, true),
        Field::new("bytes", DataType::Int64, true),
        Field::new("event_time", DataType::Int64, true),
    ]));
    let batch = RecordBatch::try_new(
        schema,
        vec![
            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
            Arc::new(StringArray::from(str_vals.clone())) as arrow::array::ArrayRef,
            Arc::new(Int64Array::from(int_vals.clone())) as arrow::array::ArrayRef,
            Arc::new(Int64Array::from(int_vals.clone())) as arrow::array::ArrayRef,
            Arc::new(Int64Array::from(int_vals)) as arrow::array::ArrayRef,
        ],
    )
    .unwrap();
    eprintln!("orig batch bytes = {}", batch.get_array_memory_size());
    eprintln!("orig sip col = {}", batch.column(0).get_array_memory_size());
    let enc = encode_ipc("conn_events", &batch).unwrap();
    eprintln!("ipc payload = {}", enc.len());
    let frame = decode_ipc(&enc).unwrap();
    eprintln!(
        "decoded batch bytes = {}",
        frame.batch.get_array_memory_size()
    );
    eprintln!(
        "decoded sip col = {}",
        frame.batch.column(0).get_array_memory_size()
    );
}