//! Partitioned graph backend implementation.
//!
//! This module provides a GraphBackend implementation that operates on a single
//! partition of a distributed graph while tracking cross-partition edges.

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

use indexmap::IndexMap;

use crate::distributed::partition::{
    PartitionIdAllocator, PartitionMap, PartitionedId,
};
#[cfg(feature = "distributed")]
use crate::distributed::partition::fencing::FencingToken;
use crate::error::{ExecutionError, ExecutionResult};
use crate::graph::{GraphBackend, GraphEdge, GraphNode, NetworKitRustBackend, Path, PropertyValue};

#[cfg(feature = "distributed")]
use crate::distributed::storage::wal::{WalConfig, WalOperation, WriteAheadLog};
#[cfg(feature = "distributed")]
use crate::error::CypherResult;

/// Metadata about a cross-partition edge.
///
/// When an edge spans partitions, we track which partitions need to be
/// queried for complete edge information.
#[derive(Debug, Clone)]
pub struct CrossPartitionEdgeMetadata {
    /// ID of the edge
    pub edge_id: u64,

    /// Source node ID
    pub source_node_id: u64,

    /// Target node ID
    pub target_node_id: u64,

    /// Partition that owns the source node
    pub source_partition: u32,

    /// Partition that owns the target node
    pub target_partition: u32,

    /// Relationship type
    pub rel_type: String,
}

/// A partitioned graph backend that implements GraphBackend.
///
/// This backend manages a single partition of a larger distributed graph.
/// It maintains:
/// - Local graph data for nodes/edges owned by this partition
/// - Partition map for routing queries
/// - Cross-partition edge metadata for edges spanning partitions
///
/// # Type Parameter
///
/// `G` is the underlying local graph backend. Defaults to `NetworKitRustBackend`
/// for best performance. Any backend implementing `GraphBackend + Default` works.
///
/// # Ghost Nodes Strategy
///
/// For cross-partition edges, we use "ghost nodes":
/// - Edges are stored in the source node's partition
/// - If target node is in another partition, we track it as cross-partition
/// - Queries for complete edge lists may need to contact other partitions
pub struct PartitionedBackend<G = NetworKitRustBackend> {
    /// The partition ID this backend manages
    partition_id: u32,

    /// Local graph data (nodes and edges owned by this partition)
    local_graph: G,

    /// Global partition map (shared, read-mostly)
    partition_map: Arc<RwLock<PartitionMap>>,

    /// ID allocator for this partition
    id_allocator: PartitionIdAllocator,

    /// Metadata about edges that cross partition boundaries
    /// Key: edge ID, Value: cross-partition metadata
    cross_partition_edges: HashMap<u64, CrossPartitionEdgeMetadata>,

    /// Optional Write-Ahead Log; present only when WAL is enabled in config.
    #[cfg(feature = "distributed")]
    pub(crate) wal: Option<WriteAheadLog>,

    /// Optional fencing token for split-brain prevention.
    /// When set, every mutation checks that the caller's epoch is not stale.
    #[cfg(feature = "distributed")]
    fencing_token: Option<FencingToken>,
}

impl<G: GraphBackend + Default> PartitionedBackend<G> {
    /// Create a new partitioned graph backend.
    ///
    /// # Arguments
    ///
    /// * `partition_id` - The partition this backend manages
    /// * `partition_map` - Shared partition map for routing
    pub fn new(partition_id: u32, partition_map: Arc<RwLock<PartitionMap>>) -> Self {
        Self {
            partition_id,
            local_graph: G::default(),
            partition_map,
            id_allocator: PartitionIdAllocator::new(partition_id, 1),
            cross_partition_edges: HashMap::new(),
            #[cfg(feature = "distributed")]
            wal: None,
            #[cfg(feature = "distributed")]
            fencing_token: None,
        }
    }

    /// Attach a Write-Ahead Log to this backend.
    ///
    /// Opens (or creates) the WAL at the path derived from `config` and
    /// `storage_path`, then stores it on the backend so every subsequent
    /// mutation is logged before being applied in-memory.
    ///
    /// Returns `Err` if the WAL file cannot be opened.
    #[cfg(feature = "distributed")]
    pub fn with_wal(mut self, config: &WalConfig, storage_path: &str) -> CypherResult<Self>
    where
        G: GraphBackend + Default,
    {
        if config.enabled {
            self.wal = Some(WriteAheadLog::open(config, storage_path)?);
        }
        Ok(self)
    }

    /// Set the fencing token for split-brain prevention.
    ///
    /// Once set, the backend tracks the current epoch. Use `check_epoch()`
    /// to validate that a caller's epoch is not stale before mutations.
    #[cfg(feature = "distributed")]
    pub fn set_fencing_token(&mut self, token: FencingToken) {
        self.fencing_token = Some(token);
    }

    /// Get the current epoch from the fencing token (if set).
    #[cfg(feature = "distributed")]
    pub fn current_epoch(&self) -> Option<u64> {
        self.fencing_token.as_ref().map(|t| t.epoch())
    }

    /// Validate that the caller's epoch matches the current fencing token.
    ///
    /// Returns `Ok(())` if no fencing token is set (opt-in mechanism) or
    /// if the epoch is valid. Returns `Err(StaleEpoch)` if the epoch is stale.
    #[cfg(feature = "distributed")]
    pub fn check_epoch(&self, caller_epoch: u64) -> ExecutionResult<()> {
        if let Some(ref token) = self.fencing_token {
            token.validate(caller_epoch).map_err(|e| {
                ExecutionError::Internal(format!("{e}"))
            })
        } else {
            Ok(())
        }
    }

    /// Get the partition ID this backend manages.
    pub fn partition_id(&self) -> u32 {
        self.partition_id
    }

    /// Replay a `CreateNode` WAL entry with the exact stored ID.
    ///
    /// This bypasses the partition ID allocator and WAL logging, since the
    /// operation was already recorded in the WAL we are replaying from.
    pub fn replay_create_node(
        &mut self,
        id: u64,
        labels: Vec<String>,
        properties: IndexMap<String, PropertyValue>,
    ) {
        let mut node = GraphNode::new(id);
        node.properties = properties;
        for label in &labels {
            node.labels.insert(label.clone());
            if let Ok(mut map) = self.partition_map.write() {
                map.add_label(self.partition_id, label.clone());
            }
        }
        self.local_graph.add_existing_node(node);
        if let Ok(mut map) = self.partition_map.write() {
            map.assign_vertex(id, self.partition_id);
        }
    }

    /// Replay a `CreateRelationship` WAL entry with the exact stored edge ID.
    ///
    /// This bypasses the partition ID allocator and WAL logging.
    pub fn replay_create_relationship(
        &mut self,
        edge_id: u64,
        src: u64,
        dst: u64,
        rel_type: String,
        properties: IndexMap<String, PropertyValue>,
    ) -> ExecutionResult<u64> {
        let result = self.local_graph.create_relationship_with_id(
            edge_id, src, dst, rel_type.clone(), properties,
        );
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.assign_edge(edge_id, self.partition_id);
                map.add_relationship_type(self.partition_id, rel_type);
            }
        }
        result
    }

    /// Replay an `AddCrossPartitionEdge` WAL entry.
    pub fn replay_add_cross_partition_edge(
        &mut self,
        edge_id: u64,
        source_node_id: u64,
        target_node_id: u64,
        source_partition: u32,
        target_partition: u32,
        rel_type: String,
    ) {
        self.cross_partition_edges.insert(edge_id, CrossPartitionEdgeMetadata {
            edge_id,
            source_node_id,
            target_node_id,
            source_partition,
            target_partition,
            rel_type,
        });
    }

    /// Replay a `RemoveCrossPartitionEdge` WAL entry.
    pub fn replay_remove_cross_partition_edge(&mut self, edge_id: u64) {
        self.cross_partition_edges.remove(&edge_id);
    }

    /// Read-only access to cross-partition edge metadata.
    pub fn cross_partition_edges(&self) -> &HashMap<u64, CrossPartitionEdgeMetadata> {
        &self.cross_partition_edges
    }

    /// Check if this partition owns the given node.
    pub fn owns_node(&self, node_id: u64) -> bool {
        let partitioned_id = PartitionedId::from_raw(node_id);
        partitioned_id.partition_id() == self.partition_id
    }

    /// Check if this partition owns the given edge.
    pub fn owns_edge(&self, edge_id: u64) -> bool {
        let partitioned_id = PartitionedId::from_raw(edge_id);
        partitioned_id.partition_id() == self.partition_id
    }

    /// Get the partition ID for a node by querying the partition map.
    fn get_node_partition(&self, node_id: u64) -> Option<u32> {
        // First check if we can determine from the ID itself
        let partitioned_id = PartitionedId::from_raw(node_id);
        if partitioned_id.partition_id() > 0 {
            return Some(partitioned_id.partition_id());
        }

        // Otherwise check the partition map
        if let Ok(map) = self.partition_map.read() {
            map.get_vertex_partition(node_id)
        } else {
            None
        }
    }

    /// Allocate a new node ID for this partition.
    fn allocate_node_id(&mut self) -> u64 {
        self.id_allocator.allocate_node_id().as_raw()
    }

    /// Allocate a new edge ID for this partition.
    fn allocate_edge_id(&mut self) -> u64 {
        self.id_allocator.allocate_edge_id().as_raw()
    }
}

impl<G> GraphBackend for PartitionedBackend<G>
where
    G: GraphBackend + Default,
{
    type NodeRef<'a> = G::NodeRef<'a> where G: 'a;
    type EdgeRef<'a> = G::EdgeRef<'a> where G: 'a;
    type NodesIter<'a> = G::NodesIter<'a> where G: 'a;
    type EdgesIter<'a> = G::EdgesIter<'a> where G: 'a;

    // === Core Write Primitives (WAL replay / restore) ===

    fn add_existing_node(&mut self, node: GraphNode) {
        self.local_graph.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> {
        self.local_graph.create_relationship_with_id(edge_id, start_id, end_id, rel_type, properties)
    }

    // === Core Read Operations ===

    fn get_node(&self, id: u64) -> Option<Self::NodeRef<'_>> {
        if !self.owns_node(id) {
            return None;
        }
        self.local_graph.get_node(id)
    }

    fn get_node_mut(&mut self, id: u64) -> Option<&mut GraphNode> {
        if !self.owns_node(id) {
            return None;
        }
        self.local_graph.get_node_mut(id)
    }

    fn all_nodes(&self) -> Self::NodesIter<'_> {
        self.local_graph.all_nodes()
    }

    fn nodes_with_label(&self, label: &str) -> Vec<Self::NodeRef<'_>> {
        self.local_graph.nodes_with_label(label)
    }

    fn get_edge(&self, id: u64) -> Option<Self::EdgeRef<'_>> {
        if !self.owns_edge(id) {
            return None;
        }
        self.local_graph.get_edge(id)
    }

    fn get_edge_mut(&mut self, id: u64) -> Option<&mut GraphEdge> {
        if !self.owns_edge(id) {
            return None;
        }
        self.local_graph.get_edge_mut(id)
    }

    fn all_edges(&self) -> Self::EdgesIter<'_> {
        self.local_graph.all_edges()
    }

    fn edges_with_type(&self, rel_type: &str) -> Vec<Self::EdgeRef<'_>> {
        self.local_graph.edges_with_type(rel_type)
    }

    fn get_edge_endpoints(&self, edge_id: u64) -> Option<(u64, u64)> {
        if !self.owns_edge(edge_id) {
            return None;
        }
        self.local_graph.get_edge_endpoints(edge_id)
    }

    fn outgoing_edges(&self, node_id: u64) -> Vec<(u64, Self::EdgeRef<'_>)> {
        if !self.owns_node(node_id) {
            return vec![];
        }
        self.local_graph.outgoing_edges(node_id)
    }

    fn incoming_edges(&self, node_id: u64) -> Vec<(u64, Self::EdgeRef<'_>)> {
        if !self.owns_node(node_id) {
            return vec![];
        }
        self.local_graph.incoming_edges(node_id)
    }

    fn all_edges_for_node(&self, node_id: u64) -> Vec<(u64, u64, Self::EdgeRef<'_>)> {
        if !self.owns_node(node_id) {
            return vec![];
        }
        self.local_graph.all_edges_for_node(node_id)
    }

    fn count_node_relationships(&self, node_id: u64) -> usize {
        if !self.owns_node(node_id) {
            return 0;
        }
        self.local_graph.count_node_relationships(node_id)
    }

    // === Path Finding ===

    fn shortest_path(&self, start_id: u64, end_id: u64) -> Option<Path> {
        // For now, only support paths within the same partition
        if !self.owns_node(start_id) || !self.owns_node(end_id) {
            return None;
        }
        self.local_graph.shortest_path(start_id, end_id)
    }

    fn all_shortest_paths(&self, start_id: u64, end_id: u64) -> Vec<Path> {
        // For now, only support paths within the same partition
        if !self.owns_node(start_id) || !self.owns_node(end_id) {
            return vec![];
        }
        self.local_graph.all_shortest_paths(start_id, end_id)
    }

    // === Core Write Operations ===

    fn create_node(
        &mut self,
        labels: impl IntoIterator<Item = impl Into<String>>,
        properties: IndexMap<String, PropertyValue>,
    ) -> u64 {
        let labels_vec: Vec<String> = labels.into_iter().map(Into::into).collect();
        let props_clone = properties.clone();

        // Allocate a new ID for this partition
        let node_id = self.allocate_node_id();

        // Create the node in the local graph
        let mut node = GraphNode::new(node_id);
        node.properties = properties;

        for label in &labels_vec {
            node.labels.insert(label.clone());

            // Update partition map
            if let Ok(mut map) = self.partition_map.write() {
                map.add_label(self.partition_id, label.clone());
            }
        }

        // Add to local graph
        self.local_graph.add_existing_node(node);

        // Update partition map
        if let Ok(mut map) = self.partition_map.write() {
            map.assign_vertex(node_id, self.partition_id);
        }

        // WAL append (best-effort: log error but never fail the operation)
        #[cfg(feature = "distributed")]
        if let Some(ref mut wal) = self.wal {
            if let Err(e) = wal.append(WalOperation::CreateNode {
                id: node_id,
                labels: labels_vec,
                properties: props_clone,
            }) {
                tracing::error!(error = %e, node_id, "WAL write failed for CreateNode");
            }
        }

        node_id
    }

    fn create_relationship(
        &mut self,
        start_id: u64,
        end_id: u64,
        rel_type: impl Into<String>,
        properties: IndexMap<String, PropertyValue>,
    ) -> ExecutionResult<u64> {
        let rel_type = rel_type.into();
        let props_clone = properties.clone();

        // Check if we own the source node
        if !self.owns_node(start_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Source node {} not in partition {}",
                start_id, self.partition_id
            )));
        }

        // Verify source node exists locally
        if self.local_graph.get_node(start_id).is_none() {
            return Err(ExecutionError::NodeNotFound(start_id));
        }

        // Allocate edge ID
        let edge_id = self.allocate_edge_id();

        // Check if target is in another partition
        let target_partition = self.get_node_partition(end_id);
        let is_cross_partition = target_partition.map_or(false, |p| p != self.partition_id);

        if is_cross_partition {
            let target_part = target_partition.unwrap();

            // Track as cross-partition edge
            self.cross_partition_edges.insert(edge_id, CrossPartitionEdgeMetadata {
                edge_id,
                source_node_id: start_id,
                target_node_id: end_id,
                source_partition: self.partition_id,
                target_partition: target_part,
                rel_type: rel_type.clone(),
            });

            // WAL: log cross-partition edge metadata
            #[cfg(feature = "distributed")]
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::AddCrossPartitionEdge {
                    edge_id,
                    source_node_id: start_id,
                    target_node_id: end_id,
                    source_partition: self.partition_id,
                    target_partition: target_part,
                    rel_type: rel_type.clone(),
                }) {
                    tracing::error!(error = %e, edge_id, "WAL write failed for AddCrossPartitionEdge");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for AddCrossPartitionEdge: {e}")
                    ));
                }
            }

            // Update partition map
            if let Ok(mut map) = self.partition_map.write() {
                map.mark_cross_partition_edge(self.partition_id);
            }
        }

        // Create edge in local graph
        // Note: PropertyGraph will handle the case where end_id doesn't exist locally
        // by creating a "ghost node" reference
        let result = self.local_graph.create_relationship_with_id(
            edge_id,
            start_id,
            end_id,
            rel_type.clone(),
            properties,
        );

        // Update partition map
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.assign_edge(edge_id, self.partition_id);
                map.add_relationship_type(self.partition_id, rel_type.clone());
            }

            // WAL append — fail the mutation if WAL write fails to preserve durability.
            #[cfg(feature = "distributed")]
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::CreateRelationship {
                    id: edge_id,
                    src: start_id,
                    dst: end_id,
                    rel_type,
                    properties: props_clone,
                }) {
                    tracing::error!(error = %e, edge_id, "WAL write failed for CreateRelationship");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for CreateRelationship: {e}")
                    ));
                }
            }
        }

        result
    }

    fn delete_node(&mut self, id: u64) -> ExecutionResult<GraphNode> {
        if !self.owns_node(id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                id, self.partition_id
            )));
        }

        // WAL: log BEFORE deletion so we record the intent.
        // Fail the operation if WAL write fails to preserve durability.
        #[cfg(feature = "distributed")]
        if let Some(ref mut wal) = self.wal {
            if let Err(e) = wal.append(WalOperation::DeleteNode { id }) {
                tracing::error!(error = %e, node_id = id, "WAL write failed for DeleteNode");
                return Err(ExecutionError::Internal(
                    format!("WAL write failed for DeleteNode: {e}")
                ));
            }
        }

        let result = self.local_graph.delete_node(id);

        // Update partition map if successful
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.vertex_to_partition.remove(&id);
            }
        }

        result
    }

    fn delete_relationship(&mut self, id: u64) -> ExecutionResult<GraphEdge> {
        if !self.owns_edge(id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Edge {} not in partition {}",
                id, self.partition_id
            )));
        }

        // WAL: log BEFORE deletion.
        // Fail the operation if WAL write fails to preserve durability.
        #[cfg(feature = "distributed")]
        if let Some(ref mut wal) = self.wal {
            if let Err(e) = wal.append(WalOperation::DeleteRelationship { id }) {
                tracing::error!(error = %e, edge_id = id, "WAL write failed for DeleteRelationship");
                return Err(ExecutionError::Internal(
                    format!("WAL write failed for DeleteRelationship: {e}")
                ));
            }
        }

        // Remove from cross-partition tracking if present, logging removal to WAL.
        if self.cross_partition_edges.remove(&id).is_some() {
            #[cfg(feature = "distributed")]
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::RemoveCrossPartitionEdge { edge_id: id }) {
                    tracing::error!(error = %e, edge_id = id, "WAL write failed for RemoveCrossPartitionEdge");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for RemoveCrossPartitionEdge: {e}")
                    ));
                }
            }
        }

        let result = self.local_graph.delete_relationship(id);

        // Update partition map if successful
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.edge_to_partition.remove(&id);
            }
        }

        result
    }

    fn detach_delete_node(&mut self, id: u64) -> ExecutionResult<GraphNode> {
        if !self.owns_node(id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                id, self.partition_id
            )));
        }

        // Remove all edges connected to this node from cross-partition tracking
        let edges_to_remove: Vec<u64> = self.cross_partition_edges
            .iter()
            .filter(|(_, meta)| meta.source_node_id == id || meta.target_node_id == id)
            .map(|(edge_id, _)| *edge_id)
            .collect();

        for edge_id in edges_to_remove {
            self.cross_partition_edges.remove(&edge_id);
        }

        let result = self.local_graph.detach_delete_node(id);

        // Update partition map if successful
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.vertex_to_partition.remove(&id);
            }
        }

        result
    }

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

    fn set_node_property(
        &mut self,
        node_id: u64,
        key: impl Into<String>,
        value: PropertyValue,
    ) -> ExecutionResult<()> {
        if !self.owns_node(node_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                node_id, self.partition_id
            )));
        }
        let key_str = key.into();
        let value_clone = value.clone();
        let result = self.local_graph.set_node_property(node_id, key_str.clone(), value);

        #[cfg(feature = "distributed")]
        if result.is_ok() {
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::SetNodeProperty {
                    id: node_id,
                    key: key_str,
                    value: value_clone,
                }) {
                    tracing::error!(error = %e, node_id, "WAL write failed for SetNodeProperty");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for SetNodeProperty: {e}")
                    ));
                }
            }
        }

        result
    }

    fn set_edge_property(
        &mut self,
        edge_id: u64,
        key: impl Into<String>,
        value: PropertyValue,
    ) -> ExecutionResult<()> {
        if !self.owns_edge(edge_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Edge {} not in partition {}",
                edge_id, self.partition_id
            )));
        }
        let key_str = key.into();
        let value_clone = value.clone();
        let result = self.local_graph.set_edge_property(edge_id, key_str.clone(), value);

        #[cfg(feature = "distributed")]
        if result.is_ok() {
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::SetEdgeProperty {
                    id: edge_id,
                    key: key_str,
                    value: value_clone,
                }) {
                    tracing::error!(error = %e, edge_id, "WAL write failed for SetEdgeProperty");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for SetEdgeProperty: {e}")
                    ));
                }
            }
        }

        result
    }

    fn remove_node_property(
        &mut self,
        node_id: u64,
        key: &str,
    ) -> ExecutionResult<Option<PropertyValue>> {
        if !self.owns_node(node_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                node_id, self.partition_id
            )));
        }
        let result = self.local_graph.remove_node_property(node_id, key);

        #[cfg(feature = "distributed")]
        if result.is_ok() {
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::RemoveNodeProperty {
                    id: node_id,
                    key: key.to_string(),
                }) {
                    tracing::error!(error = %e, node_id, "WAL write failed for RemoveNodeProperty");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for RemoveNodeProperty: {e}")
                    ));
                }
            }
        }

        result
    }

    fn remove_edge_property(
        &mut self,
        edge_id: u64,
        key: &str,
    ) -> ExecutionResult<Option<PropertyValue>> {
        if !self.owns_edge(edge_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Edge {} not in partition {}",
                edge_id, self.partition_id
            )));
        }
        let result = self.local_graph.remove_edge_property(edge_id, key);

        #[cfg(feature = "distributed")]
        if result.is_ok() {
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::RemoveEdgeProperty {
                    id: edge_id,
                    key: key.to_string(),
                }) {
                    tracing::error!(error = %e, edge_id, "WAL write failed for RemoveEdgeProperty");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for RemoveEdgeProperty: {e}")
                    ));
                }
            }
        }

        result
    }

    fn add_label(&mut self, node_id: u64, label: impl Into<String>) -> ExecutionResult<()> {
        if !self.owns_node(node_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                node_id, self.partition_id
            )));
        }

        let label = label.into();
        let result = self.local_graph.add_label(node_id, label.clone());

        // Update partition map if successful
        if result.is_ok() {
            if let Ok(mut map) = self.partition_map.write() {
                map.add_label(self.partition_id, label.clone());
            }

            #[cfg(feature = "distributed")]
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::AddLabel {
                    id: node_id,
                    label: label.clone(),
                }) {
                    tracing::error!(error = %e, node_id, "WAL write failed for AddLabel");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for AddLabel: {e}")
                    ));
                }
            }
        }

        result
    }

    fn remove_label(&mut self, node_id: u64, label: &str) -> ExecutionResult<bool> {
        if !self.owns_node(node_id) {
            return Err(ExecutionError::PartitionMismatch(format!(
                "Node {} not in partition {}",
                node_id, self.partition_id
            )));
        }
        let result = self.local_graph.remove_label(node_id, label);

        #[cfg(feature = "distributed")]
        if result.is_ok() {
            if let Some(ref mut wal) = self.wal {
                if let Err(e) = wal.append(WalOperation::RemoveLabel {
                    id: node_id,
                    label: label.to_string(),
                }) {
                    tracing::error!(error = %e, node_id, "WAL write failed for RemoveLabel");
                    return Err(ExecutionError::Internal(
                        format!("WAL write failed for RemoveLabel: {e}")
                    ));
                }
            }
        }

        result
    }

    // === Metadata ===

    fn node_count(&self) -> usize {
        self.local_graph.node_count()
    }

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

/// Backwards-compatible type alias: `PartitionedGraphBackend` = `PartitionedBackend<NetworKitRustBackend>`.
///
/// New code should use `PartitionedBackend<G>` directly.
pub type PartitionedGraphBackend = PartitionedBackend<NetworKitRustBackend>;

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

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

        assert_eq!(backend.partition_id, 0);
        assert_eq!(backend.node_count(), 0);
        assert_eq!(backend.edge_count(), 0);
    }

    #[test]
    fn test_create_node_in_partition() {
        let partition_map = Arc::new(RwLock::new(PartitionMap::new(3)));
        let mut backend = PartitionedGraphBackend::new(1, partition_map.clone());

        let labels = vec!["Person"];
        let mut properties = IndexMap::new();
        properties.insert("name".to_string(), PropertyValue::String("Alice".to_string()));

        let node_id = backend.create_node(labels, properties);

        // Verify node ID has correct partition
        let partitioned_id = PartitionedId::from_raw(node_id);
        assert_eq!(partitioned_id.partition_id(), 1);

        // Verify node exists
        assert!(backend.get_node(node_id).is_some());
        assert_eq!(backend.node_count(), 1);

        // Verify partition map updated
        let map = partition_map.read().unwrap();
        assert_eq!(map.get_vertex_partition(node_id), Some(1));
    }

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

        // Create node in partition 1
        let node_id = backend.create_node(vec!["Test"], IndexMap::new());

        // Backend should own this node
        assert!(backend.owns_node(node_id));

        // Create a hypothetical node ID from partition 2
        let other_partition_id = PartitionedId::new(2, 1).as_raw();

        // Backend should not own node from other partition
        assert!(!backend.owns_node(other_partition_id));
    }

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

        // Create two nodes
        let node1 = backend.create_node(vec!["Person"], IndexMap::new());
        let node2 = backend.create_node(vec!["Person"], IndexMap::new());

        // Create relationship
        let edge_id = backend.create_relationship(
            node1,
            node2,
            "KNOWS",
            IndexMap::new(),
        ).unwrap();

        // Verify edge exists
        assert!(backend.get_edge(edge_id).is_some());
        assert_eq!(backend.edge_count(), 1);

        // Verify edge is not tracked as cross-partition
        assert!(!backend.cross_partition_edges.contains_key(&edge_id));
    }

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

        // Try to get node from another partition
        let other_partition_node = PartitionedId::new(1, 1).as_raw();
        assert!(backend.get_node(other_partition_node).is_none());

        // Try to delete node from another partition
        let result = backend.delete_node(other_partition_node);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), ExecutionError::PartitionMismatch(_)));
    }
}
