//! MemoryNode ↔ Arrow RecordBatch bidirectional conversion
//!
//! Defines the Arrow schema for `MemoryNode` and provides conversion functions
//! between `Vec<MemoryNode>` and `arrow::RecordBatch`. This is the bridge
//! between AiFastDb's domain model and the Lance columnar storage format.
//!
//! # Schema Design
//!
//! | MemoryNode field    | Arrow type                        | Notes                          |
//! |---------------------|-----------------------------------|--------------------------------|
//! | id                  | Utf8                              | Primary key                    |
//! | content             | Utf8                              |                                |
//! | embedding           | FixedSizeList(Float32, dim)       | Nullable; dim from first node  |
//! | content_type        | Utf8                              | Enum serialized as string      |
//! | parent_id           | Utf8                              | Nullable                       |
//! | children            | LargeUtf8 (JSON array)            | Serialized as JSON string      |
//! | depth               | UInt32                            |                                |
//! | metadata            | LargeUtf8 (JSON object)           | Serialized as JSON string      |
//! | tags                | LargeUtf8 (JSON array)            | Serialized as JSON string      |
//! | created_at          | Int64                             | Unix ms                        |
//! | last_accessed_at    | Int64                             | Unix ms                        |
//! | access_count        | UInt64                            |                                |
//! | importance          | Float64                           |                                |
//! | decay_rate          | Float64                           |                                |
//! | entity_ids          | LargeUtf8 (JSON array)            | Serialized as JSON string      |
//! | relation_ids        | LargeUtf8 (JSON array)            | Serialized as JSON string      |
//! | owner_id            | Utf8                              | Nullable                       |
//! | visibility          | LargeUtf8 (JSON)                  | Enum serialized as JSON        |
//! | protection          | LargeUtf8 (JSON)                  | Enum serialized as JSON        |
//! | collaborators       | LargeUtf8 (JSON array)            | Serialized as JSON string      |

use std::collections::HashMap;
use std::sync::Arc;

use arrow_array::{
    Array, ArrayRef, Float32Array, Float64Array, Int64Array, RecordBatch,
    StringArray, UInt32Array, UInt64Array, FixedSizeListArray,
    builder::{ArrayBuilder, Float32Builder, FixedSizeListBuilder},
};
use arrow_schema::{DataType, Field, Schema};

use crate::types::memory::MemoryNode;
use crate::types::ownership::{Visibility, ProtectionLevel, Collaborator};
use crate::storage::backend::StorageError;

/// Default embedding dimension when none is provided.
const DEFAULT_EMBEDDING_DIM: i32 = 1536;

/// Build the Arrow schema for MemoryNode storage.
///
/// The `embedding_dim` parameter controls the FixedSizeList width.
/// If `None`, defaults to 1536 (OpenAI ada-002).
pub fn memory_node_arrow_schema(embedding_dim: Option<i32>) -> Schema {
    let dim = embedding_dim.unwrap_or(DEFAULT_EMBEDDING_DIM);
    Schema::new(vec![
        Field::new("id", DataType::Utf8, false),
        Field::new("content", DataType::Utf8, false),
        Field::new(
            "embedding",
            DataType::FixedSizeList(
                Arc::new(Field::new("item", DataType::Float32, true)),
                dim,
            ),
            true, // nullable — some nodes may not have embeddings
        ),
        Field::new("content_type", DataType::Utf8, false),
        Field::new("parent_id", DataType::Utf8, true),
        Field::new("children", DataType::LargeUtf8, false),
        Field::new("depth", DataType::UInt32, false),
        Field::new("metadata", DataType::LargeUtf8, false),
        Field::new("tags", DataType::LargeUtf8, false),
        Field::new("created_at", DataType::Int64, false),
        Field::new("last_accessed_at", DataType::Int64, false),
        Field::new("access_count", DataType::UInt64, false),
        Field::new("importance", DataType::Float64, false),
        Field::new("decay_rate", DataType::Float64, false),
        Field::new("entity_ids", DataType::LargeUtf8, false),
        Field::new("relation_ids", DataType::LargeUtf8, false),
        Field::new("owner_id", DataType::Utf8, true),
        Field::new("visibility", DataType::LargeUtf8, false),
        Field::new("protection", DataType::LargeUtf8, false),
        Field::new("collaborators", DataType::LargeUtf8, false),
        // Semantic Axes (BriLLM SiFu) — 3 nullable Float32 scalar columns
        Field::new("sem_axis_0", DataType::Float32, true),
        Field::new("sem_axis_1", DataType::Float32, true),
        Field::new("sem_axis_2", DataType::Float32, true),
    ])
}

/// Detect embedding dimension from a slice of MemoryNodes.
///
/// Returns the dimension of the first non-None embedding found,
/// or `None` if no embeddings are present.
fn detect_embedding_dim(nodes: &[MemoryNode]) -> Option<i32> {
    nodes
        .iter()
        .find_map(|n| n.embedding.as_ref().map(|e| e.len() as i32))
}

/// Serialize a ContentType to its string representation.
fn content_type_to_str(ct: &crate::types::ContentType) -> &'static str {
    match ct {
        crate::types::ContentType::Text => "text",
        crate::types::ContentType::Image => "image",
        crate::types::ContentType::Code => "code",
        crate::types::ContentType::Conversation => "conversation",
        crate::types::ContentType::Document => "document",
        crate::types::ContentType::Audio => "audio",
        crate::types::ContentType::Video => "video",
        crate::types::ContentType::ImageBinary => "image_binary",
        crate::types::ContentType::Other => "other",
    }
}

/// Deserialize a string to ContentType.
fn str_to_content_type(s: &str) -> crate::types::ContentType {
    match s {
        "text" => crate::types::ContentType::Text,
        "image" => crate::types::ContentType::Image,
        "code" => crate::types::ContentType::Code,
        "conversation" => crate::types::ContentType::Conversation,
        "document" => crate::types::ContentType::Document,
        "audio" => crate::types::ContentType::Audio,
        "video" => crate::types::ContentType::Video,
        "image_binary" => crate::types::ContentType::ImageBinary,
        _ => crate::types::ContentType::Other,
    }
}

/// Convert a batch of MemoryNodes into an Arrow RecordBatch.
///
/// # Arguments
///
/// * `nodes` — Slice of MemoryNode to convert
/// * `embedding_dim` — If `None`, auto-detected from the first embedding found;
///                      if no embeddings exist, defaults to 1536.
///
/// # Returns
///
/// A `RecordBatch` conforming to the schema from `memory_node_arrow_schema`.
pub fn memory_node_to_record_batch(
    nodes: &[MemoryNode],
    embedding_dim: Option<i32>,
) -> Result<RecordBatch, StorageError> {
    if nodes.is_empty() {
        let schema = Arc::new(memory_node_arrow_schema(embedding_dim));
        return Ok(RecordBatch::new_empty(schema));
    }

    let dim = embedding_dim
        .or_else(|| detect_embedding_dim(nodes))
        .unwrap_or(DEFAULT_EMBEDDING_DIM);

    // -- Scalar columns --
    let ids: Vec<&str> = nodes.iter().map(|n| n.id.as_str()).collect();
    let contents: Vec<&str> = nodes.iter().map(|n| n.content.as_str()).collect();
    let content_types: Vec<&str> = nodes.iter().map(|n| content_type_to_str(&n.content_type)).collect();
    let parent_ids: Vec<Option<&str>> = nodes.iter().map(|n| n.parent_id.as_deref()).collect();
    let depths: Vec<u32> = nodes.iter().map(|n| n.depth).collect();
    let created_ats: Vec<i64> = nodes.iter().map(|n| n.created_at).collect();
    let last_accessed_ats: Vec<i64> = nodes.iter().map(|n| n.last_accessed_at).collect();
    let access_counts: Vec<u64> = nodes.iter().map(|n| n.access_count).collect();
    let importances: Vec<f64> = nodes.iter().map(|n| n.importance).collect();
    let decay_rates: Vec<f64> = nodes.iter().map(|n| n.decay_rate).collect();
    let owner_ids: Vec<Option<&str>> = nodes.iter().map(|n| n.owner_id.as_deref()).collect();

    // -- JSON-serialized columns --
    let children_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.children).unwrap_or_else(|_| "[]".to_string()))
        .collect();
    let metadata_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.metadata).unwrap_or_else(|_| "{}".to_string()))
        .collect();
    let tags_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.tags).unwrap_or_else(|_| "[]".to_string()))
        .collect();
    let entity_ids_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.entity_ids).unwrap_or_else(|_| "[]".to_string()))
        .collect();
    let relation_ids_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.relation_ids).unwrap_or_else(|_| "[]".to_string()))
        .collect();
    let visibility_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.visibility).unwrap_or_else(|_| "\"Public\"".to_string()))
        .collect();
    let protection_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.protection).unwrap_or_else(|_| "\"None\"".to_string()))
        .collect();
    let collaborators_json: Vec<String> = nodes
        .iter()
        .map(|n| serde_json::to_string(&n.collaborators).unwrap_or_else(|_| "[]".to_string()))
        .collect();

    // -- Embedding (FixedSizeList<Float32>) --
    let embedding_array = build_embedding_array(nodes, dim)?;

    // -- Semantic Axes (3 nullable Float32 columns) --
    let sem_axis_0: Vec<Option<f32>> = nodes.iter().map(|n| n.semantic_axes.map(|a| a.axis_0)).collect();
    let sem_axis_1: Vec<Option<f32>> = nodes.iter().map(|n| n.semantic_axes.map(|a| a.axis_1)).collect();
    let sem_axis_2: Vec<Option<f32>> = nodes.iter().map(|n| n.semantic_axes.map(|a| a.axis_2)).collect();

    // Build schema
    let schema = Arc::new(memory_node_arrow_schema(Some(dim)));

    // Build arrays
    let columns: Vec<ArrayRef> = vec![
        Arc::new(StringArray::from(ids)),
        Arc::new(StringArray::from(contents)),
        embedding_array,
        Arc::new(StringArray::from(content_types)),
        Arc::new(StringArray::from(parent_ids)),
        Arc::new(arrow_array::LargeStringArray::from(children_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(UInt32Array::from(depths)),
        Arc::new(arrow_array::LargeStringArray::from(metadata_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(arrow_array::LargeStringArray::from(tags_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(Int64Array::from(created_ats)),
        Arc::new(Int64Array::from(last_accessed_ats)),
        Arc::new(UInt64Array::from(access_counts)),
        Arc::new(Float64Array::from(importances)),
        Arc::new(Float64Array::from(decay_rates)),
        Arc::new(arrow_array::LargeStringArray::from(entity_ids_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(arrow_array::LargeStringArray::from(relation_ids_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(StringArray::from(owner_ids)),
        Arc::new(arrow_array::LargeStringArray::from(visibility_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(arrow_array::LargeStringArray::from(protection_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        Arc::new(arrow_array::LargeStringArray::from(collaborators_json.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
        // Semantic Axes
        Arc::new(Float32Array::from(sem_axis_0)),
        Arc::new(Float32Array::from(sem_axis_1)),
        Arc::new(Float32Array::from(sem_axis_2)),
    ];

    RecordBatch::try_new(schema, columns)
        .map_err(|e| StorageError::Serialization(format!("Failed to build RecordBatch: {}", e)))
}

/// Build the embedding FixedSizeList array.
///
/// Nodes without an embedding get a null entry.
/// Nodes with an embedding that doesn't match `dim` are padded/truncated.
fn build_embedding_array(
    nodes: &[MemoryNode],
    dim: i32,
) -> Result<ArrayRef, StorageError> {
    let dim_usize = dim as usize;
    let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), dim);

    for node in nodes {
        match &node.embedding {
            Some(emb) => {
                let values = builder.values().as_any_mut().downcast_mut::<Float32Builder>().unwrap();
                for i in 0..dim_usize {
                    if i < emb.len() {
                        values.append_value(emb[i]);
                    } else {
                        values.append_value(0.0); // pad
                    }
                }
                builder.append(true);
            }
            None => {
                let values = builder.values().as_any_mut().downcast_mut::<Float32Builder>().unwrap();
                for _ in 0..dim_usize {
                    values.append_null();
                }
                builder.append(false); // null entry
            }
        }
    }

    Ok(Arc::new(builder.finish()))
}

/// Convert an Arrow RecordBatch back to a vector of MemoryNodes.
///
/// The batch must conform to the schema from `memory_node_arrow_schema`.
pub fn record_batch_to_memory_nodes(batch: &RecordBatch) -> Result<Vec<MemoryNode>, StorageError> {
    let n = batch.num_rows();
    if n == 0 {
        return Ok(Vec::new());
    }

    // Helper macro to downcast a column
    macro_rules! col_as {
        ($batch:expr, $name:expr, $ty:ty) => {
            $batch
                .column_by_name($name)
                .ok_or_else(|| StorageError::SchemaMismatch(format!("Missing column: {}", $name)))?
                .as_any()
                .downcast_ref::<$ty>()
                .ok_or_else(|| StorageError::SchemaMismatch(format!("Type mismatch for column: {}", $name)))?
        };
    }

    let id_col = col_as!(batch, "id", StringArray);
    let content_col = col_as!(batch, "content", StringArray);
    let content_type_col = col_as!(batch, "content_type", StringArray);
    let parent_id_col = col_as!(batch, "parent_id", StringArray);
    let depth_col = col_as!(batch, "depth", UInt32Array);
    let created_at_col = col_as!(batch, "created_at", Int64Array);
    let last_accessed_at_col = col_as!(batch, "last_accessed_at", Int64Array);
    let access_count_col = col_as!(batch, "access_count", UInt64Array);
    let importance_col = col_as!(batch, "importance", Float64Array);
    let decay_rate_col = col_as!(batch, "decay_rate", Float64Array);
    let owner_id_col = col_as!(batch, "owner_id", StringArray);

    // JSON columns (LargeUtf8)
    let children_col = col_as!(batch, "children", arrow_array::LargeStringArray);
    let metadata_col = col_as!(batch, "metadata", arrow_array::LargeStringArray);
    let tags_col = col_as!(batch, "tags", arrow_array::LargeStringArray);
    let entity_ids_col = col_as!(batch, "entity_ids", arrow_array::LargeStringArray);
    let relation_ids_col = col_as!(batch, "relation_ids", arrow_array::LargeStringArray);
    let visibility_col = col_as!(batch, "visibility", arrow_array::LargeStringArray);
    let protection_col = col_as!(batch, "protection", arrow_array::LargeStringArray);
    let collaborators_col = col_as!(batch, "collaborators", arrow_array::LargeStringArray);

    // Embedding column
    let embedding_col = batch
        .column_by_name("embedding")
        .ok_or_else(|| StorageError::SchemaMismatch("Missing column: embedding".to_string()))?
        .as_any()
        .downcast_ref::<FixedSizeListArray>()
        .ok_or_else(|| StorageError::SchemaMismatch("Type mismatch for column: embedding".to_string()))?;

    // Semantic Axes columns (optional — backward compatible with old datasets)
    let sem_axis_0_col = batch.column_by_name("sem_axis_0")
        .and_then(|c| c.as_any().downcast_ref::<Float32Array>().map(|a| a.clone()));
    let sem_axis_1_col = batch.column_by_name("sem_axis_1")
        .and_then(|c| c.as_any().downcast_ref::<Float32Array>().map(|a| a.clone()));
    let sem_axis_2_col = batch.column_by_name("sem_axis_2")
        .and_then(|c| c.as_any().downcast_ref::<Float32Array>().map(|a| a.clone()));

    let mut nodes = Vec::with_capacity(n);

    for i in 0..n {
        // Extract embedding
        let embedding = if embedding_col.is_null(i) {
            None
        } else {
            let values = embedding_col
                .value(i)
                .as_any()
                .downcast_ref::<Float32Array>()
                .map(|arr| arr.values().to_vec());
            values
        };

        // Deserialize JSON fields
        let children: Vec<String> = serde_json::from_str(children_col.value(i))
            .unwrap_or_default();
        let metadata: HashMap<String, serde_json::Value> =
            serde_json::from_str(metadata_col.value(i)).unwrap_or_default();
        let tags: Vec<String> = serde_json::from_str(tags_col.value(i))
            .unwrap_or_default();
        let entity_ids: Vec<String> = serde_json::from_str(entity_ids_col.value(i))
            .unwrap_or_default();
        let relation_ids: Vec<String> = serde_json::from_str(relation_ids_col.value(i))
            .unwrap_or_default();
        let visibility: Visibility = serde_json::from_str(visibility_col.value(i))
            .unwrap_or_default();
        let protection: ProtectionLevel = serde_json::from_str(protection_col.value(i))
            .unwrap_or_default();
        let collaborators: Vec<Collaborator> = serde_json::from_str(collaborators_col.value(i))
            .unwrap_or_default();

        // Extract semantic axes (backward compatible: None if columns don't exist)
        let semantic_axes = match (&sem_axis_0_col, &sem_axis_1_col, &sem_axis_2_col) {
            (Some(a0), Some(a1), Some(a2)) if !a0.is_null(i) && !a1.is_null(i) && !a2.is_null(i) => {
                Some(crate::types::SemanticAxes::new(a0.value(i), a1.value(i), a2.value(i)))
            }
            _ => None,
        };

        let node = MemoryNode {
            id: id_col.value(i).to_string(),
            content: content_col.value(i).to_string(),
            embedding,
            content_type: str_to_content_type(content_type_col.value(i)),
            parent_id: if parent_id_col.is_null(i) {
                None
            } else {
                Some(parent_id_col.value(i).to_string())
            },
            children,
            depth: depth_col.value(i),
            metadata,
            tags,
            created_at: created_at_col.value(i),
            last_accessed_at: last_accessed_at_col.value(i),
            access_count: access_count_col.value(i),
            importance: importance_col.value(i),
            decay_rate: decay_rate_col.value(i),
            entity_ids,
            relation_ids,
            owner_id: if owner_id_col.is_null(i) {
                None
            } else {
                Some(owner_id_col.value(i).to_string())
            },
            visibility,
            protection,
            collaborators,
            semantic_axes,
        };

        nodes.push(node);
    }

    Ok(nodes)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ContentType, MemoryNode};

    #[test]
    fn test_roundtrip_basic() {
        let node = MemoryNode::new("Hello world")
            .with_tag("test")
            .with_importance(0.8)
            .with_embedding(vec![0.1, 0.2, 0.3, 0.4]);

        let batch = memory_node_to_record_batch(&[node.clone()], Some(4)).unwrap();
        assert_eq!(batch.num_rows(), 1);
        assert_eq!(batch.num_columns(), 23); // 20 original + 3 semantic axes

        let recovered = record_batch_to_memory_nodes(&batch).unwrap();
        assert_eq!(recovered.len(), 1);
        assert_eq!(recovered[0].id, node.id);
        assert_eq!(recovered[0].content, "Hello world");
        assert_eq!(recovered[0].importance, 0.8);
        assert_eq!(recovered[0].tags, vec!["test"]);
        assert_eq!(recovered[0].embedding, Some(vec![0.1, 0.2, 0.3, 0.4]));
    }

    #[test]
    fn test_roundtrip_no_embedding() {
        let node = MemoryNode::new("No embedding");

        let batch = memory_node_to_record_batch(&[node.clone()], Some(4)).unwrap();
        let recovered = record_batch_to_memory_nodes(&batch).unwrap();

        assert_eq!(recovered.len(), 1);
        assert!(recovered[0].embedding.is_none());
    }

    #[test]
    fn test_roundtrip_mixed_embeddings() {
        let node1 = MemoryNode::new("With embedding")
            .with_embedding(vec![1.0, 2.0, 3.0]);
        let node2 = MemoryNode::new("Without embedding");
        let node3 = MemoryNode::new("Also with embedding")
            .with_embedding(vec![4.0, 5.0, 6.0]);

        let batch = memory_node_to_record_batch(&[node1.clone(), node2.clone(), node3.clone()], Some(3)).unwrap();
        assert_eq!(batch.num_rows(), 3);

        let recovered = record_batch_to_memory_nodes(&batch).unwrap();
        assert!(recovered[0].embedding.is_some());
        assert!(recovered[1].embedding.is_none());
        assert!(recovered[2].embedding.is_some());
        assert_eq!(recovered[0].embedding.as_ref().unwrap(), &vec![1.0, 2.0, 3.0]);
        assert_eq!(recovered[2].embedding.as_ref().unwrap(), &vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_roundtrip_empty() {
        let batch = memory_node_to_record_batch(&[], Some(4)).unwrap();
        assert_eq!(batch.num_rows(), 0);

        let recovered = record_batch_to_memory_nodes(&batch).unwrap();
        assert!(recovered.is_empty());
    }

    #[test]
    fn test_roundtrip_with_metadata_and_parent() {
        let node = MemoryNode::new("Child node")
            .with_parent("parent-123")
            .with_metadata("source", serde_json::json!("test"))
            .with_metadata("score", serde_json::json!(0.95))
            .with_content_type(ContentType::Code);

        let batch = memory_node_to_record_batch(&[node.clone()], Some(4)).unwrap();
        let recovered = record_batch_to_memory_nodes(&batch).unwrap();

        assert_eq!(recovered[0].parent_id, Some("parent-123".to_string()));
        assert_eq!(recovered[0].content_type, ContentType::Code);
        assert_eq!(recovered[0].metadata.get("source"), Some(&serde_json::json!("test")));
        assert_eq!(recovered[0].metadata.get("score"), Some(&serde_json::json!(0.95)));
    }

    #[test]
    fn test_roundtrip_semantic_axes() {
        use crate::types::SemanticAxes;

        let node = MemoryNode::new("With semantic axes")
            .with_embedding(vec![0.1, 0.2, 0.3, 0.4])
            .with_semantic_axes(SemanticAxes::new(0.5, -0.3, 0.8));

        let batch = memory_node_to_record_batch(&[node.clone()], Some(4)).unwrap();
        let recovered = record_batch_to_memory_nodes(&batch).unwrap();

        assert_eq!(recovered.len(), 1);
        let axes = recovered[0].semantic_axes.expect("Should have semantic axes");
        assert!((axes.axis_0 - 0.5).abs() < 1e-6);
        assert!((axes.axis_1 - (-0.3)).abs() < 1e-6);
        assert!((axes.axis_2 - 0.8).abs() < 1e-6);
    }

    #[test]
    fn test_roundtrip_no_semantic_axes() {
        // Nodes without semantic axes should roundtrip with None
        let node = MemoryNode::new("No semantic axes")
            .with_embedding(vec![0.1, 0.2, 0.3, 0.4]);

        let batch = memory_node_to_record_batch(&[node.clone()], Some(4)).unwrap();
        let recovered = record_batch_to_memory_nodes(&batch).unwrap();

        assert_eq!(recovered.len(), 1);
        assert!(recovered[0].semantic_axes.is_none());
    }

    #[test]
    fn test_roundtrip_mixed_semantic_axes() {
        use crate::types::SemanticAxes;

        let node1 = MemoryNode::new("With axes")
            .with_semantic_axes(SemanticAxes::new(1.0, 2.0, 3.0));
        let node2 = MemoryNode::new("Without axes");
        let node3 = MemoryNode::new("Also with axes")
            .with_semantic_axes(SemanticAxes::new(-1.0, 0.0, 0.5));

        let batch = memory_node_to_record_batch(&[node1, node2, node3], Some(4)).unwrap();
        let recovered = record_batch_to_memory_nodes(&batch).unwrap();

        assert!(recovered[0].semantic_axes.is_some());
        assert!(recovered[1].semantic_axes.is_none());
        assert!(recovered[2].semantic_axes.is_some());
        assert!((recovered[2].semantic_axes.unwrap().axis_0 - (-1.0)).abs() < 1e-6);
    }

    #[test]
    fn test_schema_has_correct_fields() {
        let schema = memory_node_arrow_schema(Some(768));
        assert_eq!(schema.fields().len(), 23); // 20 original + 3 semantic axes
        assert!(schema.field_with_name("id").is_ok());
        assert!(schema.field_with_name("embedding").is_ok());
        assert!(schema.field_with_name("content").is_ok());

        // Check embedding type
        let emb_field = schema.field_with_name("embedding").unwrap();
        match emb_field.data_type() {
            DataType::FixedSizeList(_, size) => assert_eq!(*size, 768),
            _ => panic!("Expected FixedSizeList for embedding"),
        }
    }

    #[test]
    fn test_auto_detect_embedding_dim() {
        let nodes = vec![
            MemoryNode::new("no emb"),
            MemoryNode::new("has emb").with_embedding(vec![1.0; 384]),
        ];
        let dim = detect_embedding_dim(&nodes);
        assert_eq!(dim, Some(384));
    }
}
