//! In-process distributed cluster for testing and benchmarking.
//!
//! `DistributedCluster<G>` is a `GraphBackend` implementation that manages N
//! [`PartitionedBackend<G>`] instances in a single process, routing every
//! operation to the correct partition(s).  No Docker, no etcd, no network
//! required — all partitions live in the same memory space and share a single
//! [`PartitionMap`].
//!
//! ## Design
//!
//! All edges (same-partition and cross-partition) are owned by the cluster via
//! `ClusterEdgeStore`.  The underlying `PartitionedBackend<G>` instances handle
//! **node-only** operations.  Edge traversal (outgoing/incoming/all) is served
//! from the cluster's own adjacency index, which uses stable cluster-level IDs
//! and is independent of the per-backend internal ID namespace.
//!
//! This design is necessary because `NetworKitRustBackend` (and similar backends)
//! maintain their own internal sequential edge IDs that conflict with the
//! partition-encoded stable IDs used by `PartitionedBackend`.

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

use indexmap::IndexMap;

use crate::distributed::graph::partitioned_backend::PartitionedBackend;
use crate::distributed::partition::{PartitionedId, PartitionIdAllocator, PartitionMap};
use crate::error::{ExecutionError, ExecutionResult};
use crate::graph::{EdgeLike, GraphBackend, GraphEdge, GraphNode, NodeLike, Path, PropertyValue};

// ─── Owned node/edge wrappers ──────────────────────────────────────────────

/// Owned node returned by [`DistributedCluster`].
#[derive(Debug, Clone)]
pub struct ClusterNode(pub GraphNode);

impl NodeLike for ClusterNode {
    fn id(&self) -> u64 { self.0.id }
    fn labels(&self) -> &HashSet<String> { &self.0.labels }
    fn properties(&self) -> &IndexMap<String, PropertyValue> { &self.0.properties }
}

/// Owned edge returned by [`DistributedCluster`].
#[derive(Debug, Clone)]
pub struct ClusterEdge(pub GraphEdge);

impl EdgeLike for ClusterEdge {
    fn id(&self) -> u64 { self.0.id }
    fn rel_type(&self) -> &str { &self.0.rel_type }
    fn properties(&self) -> &IndexMap<String, PropertyValue> { &self.0.properties }
}

// ─── Cluster-level edge store ──────────────────────────────────────────────

/// A relationship record owned entirely by the cluster.
#[derive(Debug, Clone)]
struct ClusterEdgeRecord {
    id: u64,
    start_id: u64,
    end_id: u64,
    rel_type: String,
    properties: IndexMap<String, PropertyValue>,
}

/// Owns all edges in the cluster with O(1) lookup and adjacency iteration.
#[derive(Debug, Default)]
struct ClusterEdgeStore {
    /// All edges keyed by edge_id.
    edges: HashMap<u64, ClusterEdgeRecord>,
    /// Adjacency: node_id → list of edge_ids where the node is the *source*.
    out_adj: HashMap<u64, Vec<u64>>,
    /// Adjacency: node_id → list of edge_ids where the node is the *target*.
    in_adj: HashMap<u64, Vec<u64>>,
}

impl ClusterEdgeStore {
    fn insert(&mut self, rec: ClusterEdgeRecord) {
        self.out_adj.entry(rec.start_id).or_default().push(rec.id);
        self.in_adj.entry(rec.end_id).or_default().push(rec.id);
        self.edges.insert(rec.id, rec);
    }

    fn remove(&mut self, edge_id: u64) -> Option<ClusterEdgeRecord> {
        let rec = self.edges.remove(&edge_id)?;
        if let Some(v) = self.out_adj.get_mut(&rec.start_id) {
            v.retain(|&e| e != edge_id);
        }
        if let Some(v) = self.in_adj.get_mut(&rec.end_id) {
            v.retain(|&e| e != edge_id);
        }
        Some(rec)
    }

    fn remove_for_node(&mut self, node_id: u64) -> Vec<ClusterEdgeRecord> {
        // Collect all edge_ids touching this node.
        let mut to_remove: HashSet<u64> = HashSet::new();
        if let Some(ids) = self.out_adj.get(&node_id) {
            to_remove.extend(ids);
        }
        if let Some(ids) = self.in_adj.get(&node_id) {
            to_remove.extend(ids);
        }
        to_remove.into_iter()
            .filter_map(|eid| self.remove(eid))
            .collect()
    }

    fn get(&self, edge_id: u64) -> Option<&ClusterEdgeRecord> {
        self.edges.get(&edge_id)
    }

    fn get_mut(&mut self, edge_id: u64) -> Option<&mut ClusterEdgeRecord> {
        self.edges.get_mut(&edge_id)
    }

    fn outgoing(&self, node_id: u64) -> Vec<&ClusterEdgeRecord> {
        self.out_adj.get(&node_id)
            .map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
            .unwrap_or_default()
    }

    fn incoming(&self, node_id: u64) -> Vec<&ClusterEdgeRecord> {
        self.in_adj.get(&node_id)
            .map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
            .unwrap_or_default()
    }

    fn all_touching(&self, node_id: u64) -> Vec<&ClusterEdgeRecord> {
        let mut seen = HashSet::new();
        let mut out = Vec::new();
        for rec in self.outgoing(node_id) {
            if seen.insert(rec.id) { out.push(rec); }
        }
        for rec in self.incoming(node_id) {
            if seen.insert(rec.id) { out.push(rec); }
        }
        out
    }

    fn count_touching(&self, node_id: u64) -> usize {
        let mut seen = HashSet::new();
        for rec in self.outgoing(node_id) { seen.insert(rec.id); }
        for rec in self.incoming(node_id) { seen.insert(rec.id); }
        seen.len()
    }

    fn len(&self) -> usize {
        self.edges.len()
    }

    fn iter_all(&self) -> impl Iterator<Item = &ClusterEdgeRecord> {
        self.edges.values()
    }
}

// ─── DistributedCluster ────────────────────────────────────────────────────

/// In-process N-partition cluster implementing [`GraphBackend`].
pub struct DistributedCluster<G: GraphBackend + Default> {
    /// One backend per partition (node storage only).
    partitions: Vec<PartitionedBackend<G>>,

    /// Shared partition map.
    partition_map: Arc<RwLock<PartitionMap>>,

    /// Round-robin counter for assigning new nodes to partitions.
    next_partition: usize,

    /// All edges owned at the cluster level.
    edge_store: ClusterEdgeStore,

    /// Cluster-level edge ID allocator (uses num_partitions as partition sentinel).
    edge_allocator: PartitionIdAllocator,
}

impl<G: GraphBackend + Default> DistributedCluster<G> {
    /// Create a new in-process cluster with `num_partitions` partitions.
    pub fn new(num_partitions: u32) -> Self {
        assert!(num_partitions > 0, "num_partitions must be >= 1");

        let partition_map = Arc::new(RwLock::new(PartitionMap::new(num_partitions)));
        let partitions = (0..num_partitions)
            .map(|pid| PartitionedBackend::<G>::new(pid, Arc::clone(&partition_map)))
            .collect();

        // Use num_partitions as sentinel partition id for cluster-level edge IDs.
        let edge_allocator = PartitionIdAllocator::new(num_partitions, 1);

        Self {
            partitions,
            partition_map,
            next_partition: 0,
            edge_store: ClusterEdgeStore::default(),
            edge_allocator,
        }
    }

    /// Number of partitions in this cluster.
    pub fn num_partitions(&self) -> usize {
        self.partitions.len()
    }

    /// Total node count summed across all partitions.
    pub fn total_node_count(&self) -> usize {
        self.partitions.iter().map(|p| p.node_count()).sum()
    }

    /// Node count for a specific partition.
    pub fn partition_node_count(&self, pid: u32) -> usize {
        self.partitions.get(pid as usize).map_or(0, |p| p.node_count())
    }

    /// Total edge count (cluster-level store).
    pub fn total_edge_count(&self) -> usize {
        self.edge_store.len()
    }

    /// Simulate a partition crash.
    pub fn kill_partition(&mut self, pid: u32) {
        let idx = pid as usize;
        assert!(idx < self.partitions.len(), "pid out of range");
        self.partitions[idx] =
            PartitionedBackend::<G>::new(pid, Arc::clone(&self.partition_map));
    }

    /// Restore a partition by replaying node entries.
    pub fn recover_partition_nodes(
        &mut self,
        pid: u32,
        entries: Vec<(u64, Vec<String>, IndexMap<String, PropertyValue>)>,
    ) {
        let idx = pid as usize;
        assert!(idx < self.partitions.len(), "pid out of range");
        for (id, labels, props) in entries {
            self.partitions[idx].replay_create_node(id, labels, props);
        }
    }

    /// Read-only access to the raw partition backends.
    pub fn partitions(&self) -> &[PartitionedBackend<G>] {
        &self.partitions
    }

    // ── Internal helpers ──────────────────────────────────────────────────

    fn next_partition_id(&mut self) -> u32 {
        let pid = self.next_partition as u32;
        self.next_partition = (self.next_partition + 1) % self.partitions.len();
        pid
    }

    fn partition_idx_for_node(&self, node_id: u64) -> Option<usize> {
        let pid = PartitionedId::from_raw(node_id).partition_id() as usize;
        if pid < self.partitions.len() { Some(pid) } else { None }
    }

    fn make_cluster_edge(rec: &ClusterEdgeRecord) -> GraphEdge {
        let mut edge = GraphEdge::new(rec.id, rec.rel_type.clone());
        edge.properties = rec.properties.clone();
        edge
    }
}

// ─── GraphBackend impl ─────────────────────────────────────────────────────

impl<G: GraphBackend + Default> GraphBackend for DistributedCluster<G> {
    type NodeRef<'a> = ClusterNode where G: 'a;
    type EdgeRef<'a> = ClusterEdge where G: 'a;
    type NodesIter<'a> = std::vec::IntoIter<ClusterNode> where G: 'a;
    type EdgesIter<'a> = std::vec::IntoIter<ClusterEdge> where G: 'a;

    // === Core Read Operations ===

    fn get_node(&self, id: u64) -> Option<Self::NodeRef<'_>> {
        let idx = self.partition_idx_for_node(id)?;
        let node_ref = self.partitions[idx].get_node(id)?;
        let mut node = GraphNode::new(node_ref.id());
        node.labels = node_ref.labels().clone();
        node.properties = node_ref.properties().clone();
        Some(ClusterNode(node))
    }

    fn get_node_mut(&mut self, id: u64) -> Option<&mut GraphNode> {
        let idx = self.partition_idx_for_node(id)?;
        self.partitions[idx].get_node_mut(id)
    }

    fn all_nodes(&self) -> Self::NodesIter<'_> {
        let nodes: Vec<ClusterNode> = self.partitions.iter()
            .flat_map(|p| {
                p.all_nodes().map(|n| {
                    let mut node = GraphNode::new(n.id());
                    node.labels = n.labels().clone();
                    node.properties = n.properties().clone();
                    ClusterNode(node)
                }).collect::<Vec<_>>()
            })
            .collect();
        nodes.into_iter()
    }

    fn nodes_with_label(&self, label: &str) -> Vec<Self::NodeRef<'_>> {
        self.partitions.iter()
            .flat_map(|p| {
                p.nodes_with_label(label).into_iter().map(|n| {
                    let mut node = GraphNode::new(n.id());
                    node.labels = n.labels().clone();
                    node.properties = n.properties().clone();
                    ClusterNode(node)
                }).collect::<Vec<_>>()
            })
            .collect()
    }

    fn get_edge(&self, id: u64) -> Option<Self::EdgeRef<'_>> {
        let rec = self.edge_store.get(id)?;
        Some(ClusterEdge(Self::make_cluster_edge(rec)))
    }

    fn get_edge_mut(&mut self, _id: u64) -> Option<&mut GraphEdge> {
        // Cluster-level edges are not directly mutable via this interface.
        None
    }

    fn all_edges(&self) -> Self::EdgesIter<'_> {
        let edges: Vec<ClusterEdge> = self.edge_store.iter_all()
            .map(|rec| ClusterEdge(Self::make_cluster_edge(rec)))
            .collect();
        edges.into_iter()
    }

    fn edges_with_type(&self, rel_type: &str) -> Vec<Self::EdgeRef<'_>> {
        self.edge_store.iter_all()
            .filter(|rec| rec.rel_type == rel_type)
            .map(|rec| ClusterEdge(Self::make_cluster_edge(rec)))
            .collect()
    }

    fn get_edge_endpoints(&self, edge_id: u64) -> Option<(u64, u64)> {
        let rec = self.edge_store.get(edge_id)?;
        Some((rec.start_id, rec.end_id))
    }

    fn outgoing_edges(&self, node_id: u64) -> Vec<(u64, Self::EdgeRef<'_>)> {
        self.edge_store.outgoing(node_id)
            .into_iter()
            .map(|rec| (rec.end_id, ClusterEdge(Self::make_cluster_edge(rec))))
            .collect()
    }

    fn incoming_edges(&self, node_id: u64) -> Vec<(u64, Self::EdgeRef<'_>)> {
        self.edge_store.incoming(node_id)
            .into_iter()
            .map(|rec| (rec.start_id, ClusterEdge(Self::make_cluster_edge(rec))))
            .collect()
    }

    fn all_edges_for_node(&self, node_id: u64) -> Vec<(u64, u64, Self::EdgeRef<'_>)> {
        self.edge_store.all_touching(node_id)
            .into_iter()
            .map(|rec| (rec.start_id, rec.end_id, ClusterEdge(Self::make_cluster_edge(rec))))
            .collect()
    }

    fn count_node_relationships(&self, node_id: u64) -> usize {
        self.edge_store.count_touching(node_id)
    }

    // === Path Finding ===

    fn shortest_path(&self, start_id: u64, end_id: u64) -> Option<Path> {
        // Delegate to the source partition's path finder (intra-partition only).
        let idx = self.partition_idx_for_node(start_id)?;
        self.partitions[idx].shortest_path(start_id, end_id)
    }

    fn all_shortest_paths(&self, start_id: u64, end_id: u64) -> Vec<Path> {
        let idx = match self.partition_idx_for_node(start_id) {
            Some(i) => i,
            None => return vec![],
        };
        self.partitions[idx].all_shortest_paths(start_id, end_id)
    }

    // === Core Write Operations ===

    fn add_existing_node(&mut self, node: GraphNode) {
        let pid = PartitionedId::from_raw(node.id).partition_id() as usize;
        let idx = pid.min(self.partitions.len() - 1);
        self.partitions[idx].add_existing_node(node);
    }

    fn create_relationship_with_id(
        &mut self,
        edge_id: u64,
        start_id: u64,
        end_id: u64,
        rel_type: impl Into<String>,
        properties: IndexMap<String, PropertyValue>,
    ) -> ExecutionResult<u64> {
        // Validate both endpoints exist.
        if self.get_node(start_id).is_none() {
            return Err(ExecutionError::NodeNotFound(start_id));
        }
        if self.get_node(end_id).is_none() {
            return Err(ExecutionError::NodeNotFound(end_id));
        }
        self.edge_store.insert(ClusterEdgeRecord {
            id: edge_id,
            start_id,
            end_id,
            rel_type: rel_type.into(),
            properties,
        });
        Ok(edge_id)
    }

    fn create_node(
        &mut self,
        labels: impl IntoIterator<Item = impl Into<String>>,
        properties: IndexMap<String, PropertyValue>,
    ) -> u64 {
        let pid = self.next_partition_id();
        self.partitions[pid as usize].create_node(labels, properties)
    }

    fn create_relationship(
        &mut self,
        start_id: u64,
        end_id: u64,
        rel_type: impl Into<String>,
        properties: IndexMap<String, PropertyValue>,
    ) -> ExecutionResult<u64> {
        // Validate both endpoints exist.
        if self.get_node(start_id).is_none() {
            return Err(ExecutionError::NodeNotFound(start_id));
        }
        if self.get_node(end_id).is_none() {
            return Err(ExecutionError::NodeNotFound(end_id));
        }
        let edge_id = self.edge_allocator.allocate_edge_id().as_raw();
        self.edge_store.insert(ClusterEdgeRecord {
            id: edge_id,
            start_id,
            end_id,
            rel_type: rel_type.into(),
            properties,
        });
        Ok(edge_id)
    }

    fn delete_node(&mut self, id: u64) -> ExecutionResult<GraphNode> {
        // Check for remaining edges in the cluster store.
        if self.edge_store.count_touching(id) > 0 {
            return Err(ExecutionError::DeleteError(format!(
                "Cannot delete node {} because it still has relationships. Use DETACH DELETE.",
                id
            )));
        }
        let idx = self
            .partition_idx_for_node(id)
            .ok_or_else(|| ExecutionError::NodeNotFound(id))?;
        // Delegate node deletion to the partition (which also checks for local edges,
        // but the local graph has none since we manage all edges at cluster level).
        // Use detach_delete_node to avoid double-check issues.
        self.partitions[idx].detach_delete_node(id)
    }

    fn delete_relationship(&mut self, id: u64) -> ExecutionResult<GraphEdge> {
        let rec = self.edge_store.remove(id)
            .ok_or_else(|| ExecutionError::RelationshipNotFound(id))?;
        Ok(Self::make_cluster_edge(&rec))
    }

    fn detach_delete_node(&mut self, id: u64) -> ExecutionResult<GraphNode> {
        // Remove all edges touching this node.
        self.edge_store.remove_for_node(id);
        let idx = self
            .partition_idx_for_node(id)
            .ok_or_else(|| ExecutionError::NodeNotFound(id))?;
        self.partitions[idx].detach_delete_node(id)
    }

    // === Property and Label Operations ===

    fn set_node_property(
        &mut self,
        node_id: u64,
        key: impl Into<String>,
        value: PropertyValue,
    ) -> ExecutionResult<()> {
        let idx = self
            .partition_idx_for_node(node_id)
            .ok_or_else(|| ExecutionError::NodeNotFound(node_id))?;
        self.partitions[idx].set_node_property(node_id, key, value)
    }

    fn set_edge_property(
        &mut self,
        edge_id: u64,
        key: impl Into<String>,
        value: PropertyValue,
    ) -> ExecutionResult<()> {
        let rec = self.edge_store.get_mut(edge_id)
            .ok_or_else(|| ExecutionError::RelationshipNotFound(edge_id))?;
        rec.properties.insert(key.into(), value);
        Ok(())
    }

    fn remove_node_property(
        &mut self,
        node_id: u64,
        key: &str,
    ) -> ExecutionResult<Option<PropertyValue>> {
        let idx = self
            .partition_idx_for_node(node_id)
            .ok_or_else(|| ExecutionError::NodeNotFound(node_id))?;
        self.partitions[idx].remove_node_property(node_id, key)
    }

    fn remove_edge_property(
        &mut self,
        edge_id: u64,
        key: &str,
    ) -> ExecutionResult<Option<PropertyValue>> {
        let rec = self.edge_store.get_mut(edge_id)
            .ok_or_else(|| ExecutionError::RelationshipNotFound(edge_id))?;
        Ok(rec.properties.swap_remove(key))
    }

    fn add_label(&mut self, node_id: u64, label: impl Into<String>) -> ExecutionResult<()> {
        let idx = self
            .partition_idx_for_node(node_id)
            .ok_or_else(|| ExecutionError::NodeNotFound(node_id))?;
        self.partitions[idx].add_label(node_id, label)
    }

    fn remove_label(&mut self, node_id: u64, label: &str) -> ExecutionResult<bool> {
        let idx = self
            .partition_idx_for_node(node_id)
            .ok_or_else(|| ExecutionError::NodeNotFound(node_id))?;
        self.partitions[idx].remove_label(node_id, label)
    }

    // === Metadata ===

    fn node_count(&self) -> usize {
        self.partitions.iter().map(|p| p.node_count()).sum()
    }

    fn edge_count(&self) -> usize {
        self.edge_store.len()
    }
}

// ─── Tests ─────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::NetworKitRustBackend;

    type Cluster = DistributedCluster<NetworKitRustBackend>;

    #[test]
    fn test_cluster_creation() {
        let cluster = Cluster::new(3);
        assert_eq!(cluster.num_partitions(), 3);
        assert_eq!(cluster.node_count(), 0);
        assert_eq!(cluster.edge_count(), 0);
    }

    #[test]
    fn test_create_nodes_round_robin() {
        let mut cluster = Cluster::new(3);
        for _ in 0..6 {
            cluster.create_node(vec!["Person"], IndexMap::new());
        }
        assert_eq!(cluster.node_count(), 6);
        assert_eq!(cluster.partition_node_count(0), 2);
        assert_eq!(cluster.partition_node_count(1), 2);
        assert_eq!(cluster.partition_node_count(2), 2);
    }

    #[test]
    fn test_get_node_routes_to_correct_partition() {
        let mut cluster = Cluster::new(3);
        let mut props = IndexMap::new();
        props.insert("name".to_string(), PropertyValue::String("Alice".to_string()));
        let id = cluster.create_node(vec!["Person"], props);

        let node = cluster.get_node(id).expect("node must exist");
        assert_eq!(node.id(), id);
        assert!(node.labels().contains("Person"));
        assert_eq!(
            node.properties().get("name"),
            Some(&PropertyValue::String("Alice".to_string()))
        );
    }

    #[test]
    fn test_all_nodes_collects_from_all_partitions() {
        let mut cluster = Cluster::new(3);
        for i in 0..9u64 {
            let mut props = IndexMap::new();
            props.insert("i".to_string(), PropertyValue::Integer(i as i64));
            cluster.create_node(vec!["N"], props);
        }
        let all: Vec<_> = cluster.all_nodes().collect();
        assert_eq!(all.len(), 9);
    }

    #[test]
    fn test_create_relationship_within_partition() {
        let mut cluster = Cluster::new(2);
        let a = cluster.create_node(vec!["A"], IndexMap::new()); // p0
        let _b = cluster.create_node(vec!["B"], IndexMap::new()); // p1
        let c = cluster.create_node(vec!["C"], IndexMap::new()); // p0

        let edge_id = cluster.create_relationship(a, c, "KNOWS", IndexMap::new()).unwrap();
        assert!(cluster.get_edge(edge_id).is_some());
        assert_eq!(cluster.edge_count(), 1);
    }

    #[test]
    fn test_create_cross_partition_relationship() {
        let mut cluster = Cluster::new(2);
        let a = cluster.create_node(vec!["A"], IndexMap::new()); // p0
        let b = cluster.create_node(vec!["B"], IndexMap::new()); // p1

        let edge_id = cluster.create_relationship(a, b, "LINKS", IndexMap::new()).unwrap();
        assert!(cluster.get_edge(edge_id).is_some());
        assert_eq!(cluster.edge_count(), 1);
    }

    #[test]
    fn test_outgoing_incoming_edges_visible() {
        let mut cluster = Cluster::new(2);
        let a = cluster.create_node(vec!["A"], IndexMap::new()); // p0
        let b = cluster.create_node(vec!["B"], IndexMap::new()); // p1

        cluster.create_relationship(a, b, "LINKS", IndexMap::new()).unwrap();

        let out = cluster.outgoing_edges(a);
        assert_eq!(out.len(), 1);
        assert_eq!(out[0].0, b);

        let inc = cluster.incoming_edges(b);
        assert_eq!(inc.len(), 1);
        assert_eq!(inc[0].0, a);
    }

    #[test]
    fn test_delete_node_with_no_edges() {
        let mut cluster = Cluster::new(3);
        let id = cluster.create_node(vec!["X"], IndexMap::new());
        cluster.delete_node(id).unwrap();
        assert!(cluster.get_node(id).is_none());
        assert_eq!(cluster.node_count(), 0);
    }

    #[test]
    fn test_nodes_with_label() {
        let mut cluster = Cluster::new(3);
        for _ in 0..6 { cluster.create_node(vec!["Person"], IndexMap::new()); }
        for _ in 0..3 { cluster.create_node(vec!["Company"], IndexMap::new()); }
        assert_eq!(cluster.nodes_with_label("Person").len(), 6);
        assert_eq!(cluster.nodes_with_label("Company").len(), 3);
        assert_eq!(cluster.nodes_with_label("Unknown").len(), 0);
    }

    #[test]
    fn test_kill_and_recover_partition() {
        let mut cluster = Cluster::new(3);
        let mut p0_nodes: Vec<(u64, Vec<String>, IndexMap<String, PropertyValue>)> = Vec::new();
        for i in 0..3u64 {
            let mut props = IndexMap::new();
            props.insert("i".to_string(), PropertyValue::Integer(i as i64));
            let id = cluster.create_node(vec!["X"], props.clone());
            if PartitionedId::from_raw(id).partition_id() == 0 {
                p0_nodes.push((id, vec!["X".to_string()], props));
            }
        }
        let before = cluster.total_node_count();
        cluster.kill_partition(0);
        assert!(cluster.total_node_count() < before);
        cluster.recover_partition_nodes(0, p0_nodes);
        assert_eq!(cluster.total_node_count(), before);
    }

    #[test]
    fn test_set_and_remove_node_property() {
        let mut cluster = Cluster::new(2);
        let id = cluster.create_node(vec!["N"], IndexMap::new());
        cluster.set_node_property(id, "key", PropertyValue::Integer(42)).unwrap();
        let node = cluster.get_node(id).unwrap();
        assert_eq!(node.properties().get("key"), Some(&PropertyValue::Integer(42)));
        cluster.remove_node_property(id, "key").unwrap();
        let node2 = cluster.get_node(id).unwrap();
        assert!(node2.properties().get("key").is_none());
    }

    #[test]
    fn test_set_and_remove_edge_property() {
        let mut cluster = Cluster::new(2);
        let a = cluster.create_node(vec!["A"], IndexMap::new());
        let b = cluster.create_node(vec!["B"], IndexMap::new());
        let eid = cluster.create_relationship(a, b, "R", IndexMap::new()).unwrap();
        cluster.set_edge_property(eid, "w", PropertyValue::Integer(5)).unwrap();
        let edge = cluster.get_edge(eid).unwrap();
        assert_eq!(edge.properties().get("w"), Some(&PropertyValue::Integer(5)));
        cluster.remove_edge_property(eid, "w").unwrap();
        let edge2 = cluster.get_edge(eid).unwrap();
        assert!(edge2.properties().get("w").is_none());
    }

    #[test]
    fn test_add_and_remove_label() {
        let mut cluster = Cluster::new(2);
        let id = cluster.create_node(vec!["A"], IndexMap::new());
        cluster.add_label(id, "B").unwrap();
        let node = cluster.get_node(id).unwrap();
        assert!(node.labels().contains("B"));
        cluster.remove_label(id, "B").unwrap();
        let node2 = cluster.get_node(id).unwrap();
        assert!(!node2.labels().contains("B"));
    }

    #[test]
    fn test_execute_create_and_match_relationship() {
        use crate::execute;
        let mut cluster = Cluster::new(3);
        execute(&mut cluster,
            "CREATE (a:Person {name: 'Alice'})-[:KNOWS]->(b:Person {name: 'Bob'})").unwrap();
        let r = execute(&mut cluster,
            "MATCH (a:Person)-[:KNOWS]->(b:Person) RETURN a.name, b.name").unwrap();
        assert_eq!(r.row_count(), 1);
    }

    #[test]
    fn test_detach_delete_removes_cross_partition_edges() {
        use crate::execute;
        let mut cluster = Cluster::new(2);
        execute(&mut cluster, "CREATE (:A {n:1})-[:R]->(:B {n:2})").unwrap();
        assert_eq!(cluster.edge_count(), 1);
        execute(&mut cluster, "MATCH (a:A) DETACH DELETE a").unwrap();
        assert_eq!(cluster.node_count(), 1);
        assert_eq!(cluster.edge_count(), 0);
    }
}
