//! Parquet serialization for graph partitions.
//!
//! This module provides serialization and deserialization of graph partitions
//! to/from Apache Parquet format with ZSTD compression for efficient storage
//! and recovery.

#[cfg(feature = "distributed")]
use std::fs::{self, File};
#[cfg(feature = "distributed")]
use std::path::Path;
#[cfg(feature = "distributed")]
use std::sync::Arc;

#[cfg(feature = "distributed")]
use arrow::array::{
    ArrayRef, ListArray, StringArray, UInt64Array, Array,
};
#[cfg(feature = "distributed")]
use arrow::datatypes::{DataType, Field, Schema};
#[cfg(feature = "distributed")]
use arrow::record_batch::RecordBatch;
#[cfg(feature = "distributed")]
use parquet::arrow::ArrowWriter;
#[cfg(feature = "distributed")]
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
#[cfg(feature = "distributed")]
use parquet::basic::{Compression, ZstdLevel};
#[cfg(feature = "distributed")]
use parquet::file::properties::WriterProperties;

#[cfg(feature = "distributed")]
use serde::{Deserialize, Serialize};

#[cfg(feature = "distributed")]
use crate::distributed::graph::PartitionedGraphBackend;
#[cfg(feature = "distributed")]
use crate::graph::{GraphBackend, NodeLike, EdgeLike};
#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

/// Helper macro to convert error messages to CypherError
#[cfg(feature = "distributed")]
macro_rules! cypher_err {
    ($msg:expr) => {
        CypherError::Execution(ExecutionError::Internal($msg.to_string()))
    };
}

/// Metadata about a serialized partition.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionMetadata {
    /// Partition ID
    pub partition_id: u32,

    /// Number of nodes in this partition
    pub node_count: usize,

    /// Number of edges in this partition
    pub edge_count: usize,

    /// Labels present in this partition
    pub labels: Vec<String>,

    /// Relationship types present in this partition
    pub relationship_types: Vec<String>,

    /// Schema version for compatibility
    pub schema_version: u32,

    /// Timestamp of serialization
    pub timestamp: i64,

    /// Compression statistics
    pub compression_stats: CompressionStats,
}

/// Compression statistics for the partition.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressionStats {
    /// Uncompressed size in bytes
    pub uncompressed_bytes: u64,

    /// Compressed size in bytes
    pub compressed_bytes: u64,

    /// Compression ratio
    pub compression_ratio: f64,

    /// Compression codec used
    pub codec: String,
}

/// Save a partition to Parquet files.
///
/// Creates the following structure:
/// ```text
/// {partition_dir}/
///   nodes.parquet
///   edges.parquet
///   label_index.parquet
///   metadata.json
/// ```
#[cfg(feature = "distributed")]
pub fn save_partition_to_parquet(
    backend: &PartitionedGraphBackend,
    partition_dir: &Path,
) -> CypherResult<PartitionMetadata> {
    // Create partition directory
    fs::create_dir_all(partition_dir)
        .map_err(|e| CypherError::Execution(ExecutionError::Internal(format!("Failed to create partition directory: {}", e))))?;

    // Save nodes
    let nodes_path = partition_dir.join("nodes.parquet");
    let node_count = save_nodes_to_parquet(backend, &nodes_path)?;

    // Save edges
    let edges_path = partition_dir.join("edges.parquet");
    let edge_count = save_edges_to_parquet(backend, &edges_path)?;

    // Save label index
    let label_index_path = partition_dir.join("label_index.parquet");
    save_label_index_to_parquet(backend, &label_index_path)?;

    // Collect metadata
    let labels: Vec<String> = backend.all_nodes()
        .flat_map(|node| node.labels().iter().map(|s| (*s).to_owned()).collect::<Vec<_>>())
        .collect::<std::collections::HashSet<_>>()
        .into_iter()
        .collect();

    let relationship_types: Vec<String> = backend.all_edges()
        .map(|edge| edge.rel_type().to_owned())
        .collect::<std::collections::HashSet<_>>()
        .into_iter()
        .collect();

    // Calculate compression stats (simplified for now)
    let compression_stats = CompressionStats {
        uncompressed_bytes: 0, // Would need to track during serialization
        compressed_bytes: 0,
        compression_ratio: 0.0,
        codec: "ZSTD".to_string(),
    };

    let metadata = PartitionMetadata {
        partition_id: backend.partition_id(),
        node_count,
        edge_count,
        labels,
        relationship_types,
        schema_version: 1,
        timestamp: chrono::Utc::now().timestamp(),
        compression_stats,
    };

    // Save metadata
    let metadata_path = partition_dir.join("metadata.json");
    let metadata_json = serde_json::to_string_pretty(&metadata)
        .map_err(|e| cypher_err!(format!("Failed to serialize metadata: {}", e)))?;

    fs::write(metadata_path, metadata_json)
        .map_err(|e| cypher_err!(format!("Failed to write metadata: {}", e)))?;

    Ok(metadata)
}

/// Save nodes to a Parquet file.
#[cfg(feature = "distributed")]
fn save_nodes_to_parquet(
    backend: &PartitionedGraphBackend,
    path: &Path,
) -> CypherResult<usize> {
    let nodes: Vec<_> = backend.all_nodes().collect();

    if nodes.is_empty() {
        // Create empty file
        File::create(path)
            .map_err(|e| cypher_err!(format!("Failed to create empty nodes file: {}", e)))?;
        return Ok(0);
    }

    // Define schema
    let schema = Arc::new(Schema::new(vec![
        Field::new("vertex_id", DataType::UInt64, false),
        Field::new(
            "labels",
            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
            true,
        ),
        // Properties are stored as a simplified JSON string for now
        // In a production system, we'd use a proper Struct type
        Field::new("properties_json", DataType::Utf8, true),
    ]));

    // Build arrays
    let vertex_ids: Vec<u64> = nodes.iter().map(|n| n.id()).collect();
    let vertex_id_array = UInt64Array::from(vertex_ids);

    // Build labels array (List<String>)
    let mut labels_builder = arrow::array::ListBuilder::new(arrow::array::StringBuilder::new());
    for node in &nodes {
        for label in node.labels() {
            labels_builder.values().append_value(label);
        }
        labels_builder.append(true);
    }
    let labels_array = labels_builder.finish();

    // Build properties JSON array
    let properties_json: Vec<String> = nodes.iter()
        .map(|n| serde_json::to_string(n.properties()).unwrap_or_default())
        .collect();
    let properties_array = StringArray::from(properties_json);

    // Create record batch
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![
            Arc::new(vertex_id_array) as ArrayRef,
            Arc::new(labels_array) as ArrayRef,
            Arc::new(properties_array) as ArrayRef,
        ],
    )
    .map_err(|e| cypher_err!(format!("Failed to create record batch: {}", e)))?;

    // Write to Parquet with ZSTD compression
    let file = File::create(path)
        .map_err(|e| cypher_err!(format!("Failed to create nodes file: {}", e)))?;

    let props = WriterProperties::builder()
        .set_compression(Compression::ZSTD(ZstdLevel::try_new(3).unwrap()))
        .set_dictionary_enabled(true)
        .build();

    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
        .map_err(|e| cypher_err!(format!("Failed to create Arrow writer: {}", e)))?;

    writer.write(&batch)
        .map_err(|e| cypher_err!(format!("Failed to write batch: {}", e)))?;

    writer.close()
        .map_err(|e| cypher_err!(format!("Failed to close writer: {}", e)))?;

    Ok(nodes.len())
}

/// Save edges to a Parquet file.
#[cfg(feature = "distributed")]
fn save_edges_to_parquet(
    backend: &PartitionedGraphBackend,
    path: &Path,
) -> CypherResult<usize> {
    let edges: Vec<_> = backend.all_edges().collect();

    if edges.is_empty() {
        // Create empty file
        File::create(path)
            .map_err(|e| cypher_err!(format!("Failed to create empty edges file: {}", e)))?;
        return Ok(0);
    }

    // Get edge endpoints
    let edge_endpoints: Vec<(u64, u64)> = edges.iter()
        .filter_map(|e| backend.get_edge_endpoints(e.id()))
        .collect();

    // Define schema
    let schema = Arc::new(Schema::new(vec![
        Field::new("edge_id", DataType::UInt64, false),
        Field::new("src_id", DataType::UInt64, false),
        Field::new("dst_id", DataType::UInt64, false),
        Field::new("rel_type", DataType::Utf8, false),
        Field::new("properties_json", DataType::Utf8, true),
    ]));

    // Build arrays
    let edge_ids: Vec<u64> = edges.iter().map(|e| e.id()).collect();
    let edge_id_array = UInt64Array::from(edge_ids);

    let src_ids: Vec<u64> = edge_endpoints.iter().map(|(src, _)| *src).collect();
    let src_id_array = UInt64Array::from(src_ids);

    let dst_ids: Vec<u64> = edge_endpoints.iter().map(|(_, dst)| *dst).collect();
    let dst_id_array = UInt64Array::from(dst_ids);

    let rel_types: Vec<&str> = edges.iter().map(|e| e.rel_type()).collect();
    let rel_type_array = StringArray::from(rel_types);

    let properties_json: Vec<String> = edges.iter()
        .map(|e| serde_json::to_string(e.properties()).unwrap_or_default())
        .collect();
    let properties_array = StringArray::from(properties_json);

    // Create record batch
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![
            Arc::new(edge_id_array) as ArrayRef,
            Arc::new(src_id_array) as ArrayRef,
            Arc::new(dst_id_array) as ArrayRef,
            Arc::new(rel_type_array) as ArrayRef,
            Arc::new(properties_array) as ArrayRef,
        ],
    )
    .map_err(|e| cypher_err!(format!("Failed to create edge record batch: {}", e)))?;

    // Write to Parquet with ZSTD compression
    let file = File::create(path)
        .map_err(|e| cypher_err!(format!("Failed to create edges file: {}", e)))?;

    let props = WriterProperties::builder()
        .set_compression(Compression::ZSTD(ZstdLevel::try_new(3).unwrap()))
        .set_dictionary_enabled(true)
        .build();

    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
        .map_err(|e| cypher_err!(format!("Failed to create Arrow writer: {}", e)))?;

    writer.write(&batch)
        .map_err(|e| cypher_err!(format!("Failed to write edge batch: {}", e)))?;

    writer.close()
        .map_err(|e| cypher_err!(format!("Failed to close edge writer: {}", e)))?;

    Ok(edges.len())
}

/// Save label index to a Parquet file.
#[cfg(feature = "distributed")]
fn save_label_index_to_parquet(
    backend: &PartitionedGraphBackend,
    path: &Path,
) -> CypherResult<()> {
    // Build label index
    let mut label_to_nodes: std::collections::HashMap<String, Vec<u64>> = std::collections::HashMap::new();

    for node in backend.all_nodes() {
        let node_id = node.id();
        for label in node.labels() {
            label_to_nodes
                .entry(label.to_owned())
                .or_default()
                .push(node_id);
        }
    }

    if label_to_nodes.is_empty() {
        // Create empty file
        File::create(path)
            .map_err(|e| cypher_err!(format!("Failed to create empty label index: {}", e)))?;
        return Ok(());
    }

    // Define schema
    let schema = Arc::new(Schema::new(vec![
        Field::new("label", DataType::Utf8, false),
        Field::new(
            "vertex_ids",
            DataType::List(Arc::new(Field::new("item", DataType::UInt64, false))),
            false,
        ),
    ]));

    // Build arrays
    let labels: Vec<String> = label_to_nodes.keys().cloned().collect();
    let label_array = StringArray::from(labels.clone());

    // Build vertex IDs list with explicit field configuration
    let vertex_ids_field = Arc::new(Field::new("item", DataType::UInt64, false));
    let mut vertex_ids_builder = arrow::array::ListBuilder::new(arrow::array::UInt64Builder::new())
        .with_field(vertex_ids_field);
    for label in &labels {
        if let Some(node_ids) = label_to_nodes.get(label) {
            for &node_id in node_ids {
                vertex_ids_builder.values().append_value(node_id);
            }
        }
        vertex_ids_builder.append(true);
    }
    let vertex_ids_array = vertex_ids_builder.finish();

    // Create record batch
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![
            Arc::new(label_array) as ArrayRef,
            Arc::new(vertex_ids_array) as ArrayRef,
        ],
    )
    .map_err(|e| cypher_err!(format!("Failed to create label index batch: {}", e)))?;

    // Write to Parquet
    let file = File::create(path)
        .map_err(|e| cypher_err!(format!("Failed to create label index file: {}", e)))?;

    let props = WriterProperties::builder()
        .set_compression(Compression::ZSTD(ZstdLevel::try_new(3).unwrap()))
        .build();

    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
        .map_err(|e| cypher_err!(format!("Failed to create label index writer: {}", e)))?;

    writer.write(&batch)
        .map_err(|e| cypher_err!(format!("Failed to write label index batch: {}", e)))?;

    writer.close()
        .map_err(|e| cypher_err!(format!("Failed to close label index writer: {}", e)))?;

    Ok(())
}

/// Restore a partition's nodes and edges from Parquet files.
///
/// Reads `nodes.parquet` and `edges.parquet` from `partition_dir` and
/// recreates every node and relationship in `backend` via the `GraphBackend`
/// trait.  Node IDs are remapped so the stable graph-internal IDs produced by
/// `create_node` are used; edges reference the new IDs through the remap table.
///
/// Returns `(node_count, edge_count)` on success.
#[cfg(feature = "distributed")]
pub fn load_partition_from_parquet(
    backend: &mut PartitionedGraphBackend,
    partition_dir: &Path,
) -> CypherResult<(usize, usize)> {
    use std::collections::HashMap;

    // ── Nodes ──────────────────────────────────────────────────────────────
    let nodes_path = partition_dir.join("nodes.parquet");
    let mut id_remap: HashMap<u64, u64> = HashMap::new();
    let mut node_count = 0usize;

    if nodes_path.exists() && nodes_path.metadata().map(|m| m.len() > 0).unwrap_or(false) {
        let file = File::open(&nodes_path)
            .map_err(|e| cypher_err!(format!("Failed to open nodes.parquet: {}", e)))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
            .map_err(|e| cypher_err!(format!("Failed to read nodes.parquet schema: {}", e)))?;
        let mut reader = builder.build()
            .map_err(|e| cypher_err!(format!("Failed to build nodes reader: {}", e)))?;

        while let Some(batch) = reader.next() {
            let batch = batch
                .map_err(|e| cypher_err!(format!("Failed to read node batch: {}", e)))?;

            let vertex_ids = batch.column(0)
                .as_any().downcast_ref::<UInt64Array>()
                .ok_or_else(|| cypher_err!("vertex_id column is not UInt64"))?;

            let labels_col = batch.column(1)
                .as_any().downcast_ref::<ListArray>()
                .ok_or_else(|| cypher_err!("labels column is not ListArray"))?;

            let props_col = batch.column(2)
                .as_any().downcast_ref::<StringArray>()
                .ok_or_else(|| cypher_err!("properties_json column is not StringArray"))?;

            for row in 0..batch.num_rows() {
                let old_id = vertex_ids.value(row);

                // Reconstruct labels
                let label_slice = labels_col.value(row);
                let label_arr = label_slice.as_any()
                    .downcast_ref::<StringArray>()
                    .ok_or_else(|| cypher_err!("labels inner array is not StringArray"))?;
                let labels: Vec<&str> = (0..label_arr.len())
                    .filter(|&i| !label_arr.is_null(i))
                    .map(|i| label_arr.value(i))
                    .collect();

                // Reconstruct properties
                let props_json = if props_col.is_null(row) { "{}" } else { props_col.value(row) };
                let properties: indexmap::IndexMap<String, crate::graph::PropertyValue> =
                    serde_json::from_str(props_json)
                        .unwrap_or_default();

                let new_id = backend.create_node(labels, properties);
                id_remap.insert(old_id, new_id);
                node_count += 1;
            }
        }
    }

    // ── Edges ──────────────────────────────────────────────────────────────
    let edges_path = partition_dir.join("edges.parquet");
    let mut edge_count = 0usize;

    if edges_path.exists() && edges_path.metadata().map(|m| m.len() > 0).unwrap_or(false) {
        let file = File::open(&edges_path)
            .map_err(|e| cypher_err!(format!("Failed to open edges.parquet: {}", e)))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
            .map_err(|e| cypher_err!(format!("Failed to read edges.parquet schema: {}", e)))?;
        let mut reader = builder.build()
            .map_err(|e| cypher_err!(format!("Failed to build edges reader: {}", e)))?;

        while let Some(batch) = reader.next() {
            let batch = batch
                .map_err(|e| cypher_err!(format!("Failed to read edge batch: {}", e)))?;

            // schema: edge_id, src_id, dst_id, rel_type, properties_json
            let src_ids = batch.column(1)
                .as_any().downcast_ref::<UInt64Array>()
                .ok_or_else(|| cypher_err!("src_id column is not UInt64"))?;
            let dst_ids = batch.column(2)
                .as_any().downcast_ref::<UInt64Array>()
                .ok_or_else(|| cypher_err!("dst_id column is not UInt64"))?;
            let rel_types = batch.column(3)
                .as_any().downcast_ref::<StringArray>()
                .ok_or_else(|| cypher_err!("rel_type column is not StringArray"))?;
            let props_col = batch.column(4)
                .as_any().downcast_ref::<StringArray>()
                .ok_or_else(|| cypher_err!("edge properties_json column is not StringArray"))?;

            for row in 0..batch.num_rows() {
                let old_src = src_ids.value(row);
                let old_dst = dst_ids.value(row);
                let rel_type = rel_types.value(row);
                let props_json = if props_col.is_null(row) { "{}" } else { props_col.value(row) };

                let new_src = *id_remap.get(&old_src)
                    .ok_or_else(|| cypher_err!(format!("src node {} not found in remap", old_src)))?;
                let new_dst = *id_remap.get(&old_dst)
                    .ok_or_else(|| cypher_err!(format!("dst node {} not found in remap", old_dst)))?;

                let properties: indexmap::IndexMap<String, crate::graph::PropertyValue> =
                    serde_json::from_str(props_json).unwrap_or_default();

                backend.create_relationship(new_src, new_dst, rel_type, properties)
                    .map_err(|e| cypher_err!(format!("Failed to restore edge: {}", e)))?;
                edge_count += 1;
            }
        }
    }

    tracing::info!(node_count, edge_count, "partition restored from Parquet snapshot");
    Ok((node_count, edge_count))
}

/// Load metadata from a partition directory.
#[cfg(feature = "distributed")]
pub fn load_partition_metadata(partition_dir: &Path) -> CypherResult<PartitionMetadata> {
    let metadata_path = partition_dir.join("metadata.json");
    let metadata_json = fs::read_to_string(metadata_path)
        .map_err(|e| cypher_err!(format!("Failed to read metadata: {}", e)))?;

    let metadata: PartitionMetadata = serde_json::from_str(&metadata_json)
        .map_err(|e| cypher_err!(format!("Failed to deserialize metadata: {}", e)))?;

    Ok(metadata)
}

// Stub implementations for non-distributed builds
#[cfg(not(feature = "distributed"))]
pub fn save_partition_to_parquet(_backend: &(), _partition_dir: &std::path::Path) -> Result<(), String> {
    Err("Distributed features not enabled. Enable the 'distributed' feature.".to_string())
}

#[cfg(not(feature = "distributed"))]
pub fn load_partition_metadata(_partition_dir: &std::path::Path) -> Result<(), String> {
    Err("Distributed features not enabled. Enable the 'distributed' feature.".to_string())
}

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use crate::distributed::partition::PartitionMap;
    use std::sync::{Arc, RwLock};
    use tempfile::TempDir;

    #[test]
    fn test_save_empty_partition() {
        let partition_map = Arc::new(RwLock::new(PartitionMap::new(1)));
        let backend = PartitionedGraphBackend::new(0, partition_map);

        let temp_dir = TempDir::new().unwrap();
        let partition_dir = temp_dir.path().join("partition_0");

        let metadata = save_partition_to_parquet(&backend, &partition_dir).unwrap();

        assert_eq!(metadata.partition_id, 0);
        assert_eq!(metadata.node_count, 0);
        assert_eq!(metadata.edge_count, 0);
        assert!(partition_dir.join("metadata.json").exists());
    }

    #[test]
    fn test_save_partition_with_nodes() {
        use indexmap::IndexMap;

        let partition_map = Arc::new(RwLock::new(PartitionMap::new(1)));
        let mut backend = PartitionedGraphBackend::new(0, partition_map);

        // Create some nodes
        backend.create_node(vec!["Person"], IndexMap::new());
        backend.create_node(vec!["Person", "Engineer"], IndexMap::new());

        let temp_dir = TempDir::new().unwrap();
        let partition_dir = temp_dir.path().join("partition_0");

        let metadata = save_partition_to_parquet(&backend, &partition_dir).unwrap();

        assert_eq!(metadata.partition_id, 0);
        assert_eq!(metadata.node_count, 2);
        assert_eq!(metadata.edge_count, 0);
        assert!(metadata.labels.contains(&"Person".to_string()));
        assert!(partition_dir.join("nodes.parquet").exists());
        assert!(partition_dir.join("label_index.parquet").exists());
    }

    #[test]
    fn test_save_partition_with_edges() {
        use indexmap::IndexMap;

        let partition_map = Arc::new(RwLock::new(PartitionMap::new(1)));
        let mut backend = PartitionedGraphBackend::new(0, partition_map);

        // Create nodes and edge
        let node1 = backend.create_node(vec!["Person"], IndexMap::new());
        let node2 = backend.create_node(vec!["Person"], IndexMap::new());
        backend.create_relationship(node1, node2, "KNOWS", IndexMap::new()).unwrap();

        let temp_dir = TempDir::new().unwrap();
        let partition_dir = temp_dir.path().join("partition_0");

        let metadata = save_partition_to_parquet(&backend, &partition_dir).unwrap();

        assert_eq!(metadata.partition_id, 0);
        assert_eq!(metadata.node_count, 2);
        assert_eq!(metadata.edge_count, 1);
        assert!(metadata.relationship_types.contains(&"KNOWS".to_string()));
        assert!(partition_dir.join("edges.parquet").exists());
    }

    #[test]
    fn test_load_metadata() {
        use indexmap::IndexMap;

        let partition_map = Arc::new(RwLock::new(PartitionMap::new(1)));
        let mut backend = PartitionedGraphBackend::new(0, partition_map);

        backend.create_node(vec!["Person"], IndexMap::new());

        let temp_dir = TempDir::new().unwrap();
        let partition_dir = temp_dir.path().join("partition_0");

        save_partition_to_parquet(&backend, &partition_dir).unwrap();

        let loaded_metadata = load_partition_metadata(&partition_dir).unwrap();

        assert_eq!(loaded_metadata.partition_id, 0);
        assert_eq!(loaded_metadata.node_count, 1);
        assert_eq!(loaded_metadata.schema_version, 1);
    }

    #[test]
    fn test_round_trip_nodes_and_edges() {
        use crate::graph::GraphBackend;
        use indexmap::IndexMap;

        let partition_map = Arc::new(RwLock::new(PartitionMap::new(1)));
        let mut backend = PartitionedGraphBackend::new(0, partition_map.clone());

        // Create nodes with labels and properties
        let mut props = IndexMap::new();
        props.insert("name".to_string(), crate::graph::PropertyValue::String("Alice".to_string()));
        let n1 = backend.create_node(vec!["Person"], props);

        let mut props2 = IndexMap::new();
        props2.insert("name".to_string(), crate::graph::PropertyValue::String("Bob".to_string()));
        let n2 = backend.create_node(vec!["Person", "Engineer"], props2);

        backend.create_relationship(n1, n2, "KNOWS", IndexMap::new()).unwrap();

        let temp_dir = TempDir::new().unwrap();
        let partition_dir = temp_dir.path().join("partition_0");

        // Save
        let meta = save_partition_to_parquet(&backend, &partition_dir).unwrap();
        assert_eq!(meta.node_count, 2);
        assert_eq!(meta.edge_count, 1);

        // Restore into a fresh backend
        let partition_map2 = Arc::new(RwLock::new(PartitionMap::new(1)));
        let mut backend2 = PartitionedGraphBackend::new(0, partition_map2);
        let (nodes, edges) = load_partition_from_parquet(&mut backend2, &partition_dir).unwrap();

        assert_eq!(nodes, 2);
        assert_eq!(edges, 1);

        // Verify node count and labels
        let restored_nodes: Vec<_> = backend2.all_nodes().collect();
        assert_eq!(restored_nodes.len(), 2);

        let has_engineer = restored_nodes.iter().any(|n| n.labels().iter().any(|l| *l == "Engineer"));
        assert!(has_engineer, "Engineer label should be restored");

        // Verify edge count
        let restored_edges: Vec<_> = backend2.all_edges().collect();
        assert_eq!(restored_edges.len(), 1);
        assert_eq!(restored_edges[0].rel_type(), "KNOWS");
    }
}
