bounded_graph 0.3.0

use std::cmp;
use std::marker::PhantomData;

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

use petgraph::data::{Build, DataMap, DataMapMut};
use petgraph::visit::{
    Data, EdgeCount, EdgeRef, GraphBase, GraphProp, IntoEdges, IntoEdgesDirected, NodeCount,
};
use petgraph::{
    graph::{DefaultIx, EdgeIndex, IndexType, NodeIndex},
    stable_graph::StableGraph,
    Directed, Direction, EdgeType, Graph, IntoWeightedEdge, Undirected,
};

use crate::{BoundedGraphError, BoundedNode};

/// A graph wrapper that enforces edge count constraints on nodes.
///
/// `BoundedGraph` wraps a petgraph graph type (either [`Graph`] or [`StableGraph`])  
/// and ensures that nodes cannot exceed their defined edge limits when adding connections.
/// All edge addition operations check the bounds before modifying the graph.
///
/// # Type Parameters
///
/// * `N` - The node weight type, which must implement [`BoundedNode`]
/// * `E` - The edge weight type
/// * `G` - The underlying graph type (defaults to `Graph<N, E, Directed, DefaultIx>`)
///
/// # Examples
///
/// Using the convenience [`FixedEdgeCount`](crate::FixedEdgeCount) type with Graph:
///
/// ```
/// use bounded_graph::{BoundedGraph, FixedEdgeCount};
///
/// let mut graph = BoundedGraph::<FixedEdgeCount<2>, ()>::new();
/// let n1 = graph.add_node(FixedEdgeCount::empty());
/// let n2 = graph.add_node(FixedEdgeCount::empty());
///
/// // This will succeed
/// assert!(graph.add_edge(n1, n2, ()).is_ok());
/// ```
///
/// Using StableGraph:
///
/// ```
/// use bounded_graph::{BoundedStableDiGraph, FixedEdgeCount};
///
/// let mut graph = BoundedStableDiGraph::<FixedEdgeCount<2>, ()>::new();
/// let n1 = graph.add_node(FixedEdgeCount::empty());
/// let n2 = graph.add_node(FixedEdgeCount::empty());
///
/// // This will succeed
/// assert!(graph.add_edge(n1, n2, ()).is_ok());
/// ```
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde-1", derive(Serialize, Deserialize))]
pub struct BoundedGraph<N, E, G = Graph<N, E, Directed, DefaultIx>> {
    pub(crate) graph: G,
    pub(crate) _phantom: PhantomData<(N, E)>,
}

/// A directed bounded graph using [`Graph`] with default index type.
pub type BoundedDiGraph<N, E> = BoundedGraph<N, E, Graph<N, E, Directed, DefaultIx>>;

/// An undirected bounded graph using [`Graph`] with default index type.
pub type BoundedUnGraph<N, E> = BoundedGraph<N, E, Graph<N, E, Undirected, DefaultIx>>;

/// A directed bounded graph using [`StableGraph`] with default index type.
pub type BoundedStableDiGraph<N, E> = BoundedGraph<N, E, StableGraph<N, E, Directed, DefaultIx>>;

/// An undirected bounded graph using [`StableGraph`] with default index type.
pub type BoundedStableUnGraph<N, E> = BoundedGraph<N, E, StableGraph<N, E, Undirected, DefaultIx>>;

// Core implementation - works with any graph type that implements the necessary traits
impl<N, E, Ty, Ix, G> BoundedGraph<N, E, G>
where
    Ty: EdgeType,
    Ix: IndexType,
    N: BoundedNode<Ix>,
    E: Clone,
    G: GraphBase<NodeId = NodeIndex<Ix>, EdgeId = EdgeIndex<Ix>>
        + Data<NodeWeight = N, EdgeWeight = E>
        + GraphProp<EdgeType = Ty>
        + DataMap
        + DataMapMut
        + Build
        + NodeCount
        + EdgeCount,
{
    /// Creates a new empty bounded graph.
    pub fn new() -> Self
    where
        G: Default,
    {
        Self {
            graph: G::default(),
            _phantom: PhantomData,
        }
    }

    /// Checks if an edge can be added between two nodes without exceeding bounds.
    ///
    /// This method verifies that both the source node has space for an outgoing edge
    /// and the target node has space for an incoming edge, according to their
    /// [`BoundedNode`] implementations.
    ///
    /// # Arguments
    ///
    /// * `node` - The source node index
    /// * `other` - The target node index
    ///
    /// # Returns
    ///
    /// `true` if the edge can be added, `false` if either node is at capacity or doesn't exist.
    #[must_use = "checking edge capacity is only useful if the result is used"]
    pub fn can_add_edge(&self, node: NodeIndex<Ix>, other: NodeIndex<Ix>) -> bool
    where
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdgesDirected,
    {
        // Get node weights first - early return if either doesn't exist
        let (o_weight, i_weight) =
            match (self.graph.node_weight(node), self.graph.node_weight(other)) {
                (Some(o), Some(i)) => (o, i),
                _ => return false,
            };

        // Count edges in each direction
        let outgoing_count = (&self.graph)
            .edges_directed(node, Direction::Outgoing)
            .count();
        let incoming_count = (&self.graph)
            .edges_directed(other, Direction::Incoming)
            .count();

        o_weight.can_add_edge(Direction::Outgoing, outgoing_count, i_weight)
            && i_weight.can_add_edge(Direction::Incoming, incoming_count, o_weight)
    }

    /// Adds a new node with the given weight to the graph.
    ///
    /// # Returns
    ///
    /// The index of the newly added node.
    #[must_use = "the node index is needed to reference the added node"]
    pub fn add_node(&mut self, node: N) -> NodeIndex<Ix> {
        self.graph.add_node(node)
    }

    /// Adds a new edge from source to target with the given weight.
    ///
    /// This method checks that both nodes have capacity for the edge according to
    /// their [`BoundedNode`] implementations before adding it.
    ///
    /// # Arguments
    ///
    /// * `source` - The source node index
    /// * `target` - The target node index
    /// * `edge` - The edge weight
    ///
    /// # Returns
    ///
    /// * `Ok(EdgeIndex)` - The index of the newly added edge
    /// * `Err(BoundedGraphError)` - If the edge cannot be added due to capacity limits
    ///
    /// # Errors
    ///
    /// Returns an error if either node is at capacity or if either node doesn't exist.
    #[must_use = "edge addition may fail due to capacity constraints"]
    pub fn add_edge(
        &mut self,
        source: NodeIndex<Ix>,
        target: NodeIndex<Ix>,
        edge: E,
    ) -> Result<EdgeIndex<Ix>, BoundedGraphError<Ix>>
    where
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdgesDirected,
    {
        // Get node weights first - return NodeNotFound if either doesn't exist
        let (source_weight, target_weight) = match (
            self.graph.node_weight(source),
            self.graph.node_weight(target),
        ) {
            (Some(s), Some(t)) => (s, t),
            (None, None) => return Err(BoundedGraphError::not_found(None, None)),
            (None, Some(_)) => return Err(BoundedGraphError::not_found(None, Some(target))),
            (Some(_), None) => return Err(BoundedGraphError::not_found(Some(source), None)),
        };

        // Count edges in each direction (not neighbors - we support parallel edges)
        let outgoing_count = (&self.graph)
            .edges_directed(source, Direction::Outgoing)
            .count();
        let incoming_count = (&self.graph)
            .edges_directed(target, Direction::Incoming)
            .count();

        // Check capacity for both nodes
        let source_has_space =
            source_weight.can_add_edge(Direction::Outgoing, outgoing_count, target_weight);
        let target_has_space =
            target_weight.can_add_edge(Direction::Incoming, incoming_count, source_weight);

        // Return appropriate error based on which node(s) rejected the edge
        match (source_has_space, target_has_space) {
            (true, true) => {
                // Both nodes have space - add the edge via Build trait
                // Returns Option<EdgeIndex> - None if nodes don't exist (but we checked they do)
                <G as petgraph::data::Build>::add_edge(&mut self.graph, source, target, edge)
                    .ok_or_else(|| BoundedGraphError::not_found(Some(source), Some(target)))
            }
            (false, true) => Err(BoundedGraphError::source_rejected(source, target)),
            (true, false) => Err(BoundedGraphError::target_rejected(source, target)),
            (false, false) => Err(BoundedGraphError::both_rejected(source, target)),
        }
    }

    /// Updates the edge from `a` to `b` with the given weight, or adds it if it doesn't exist.
    ///
    /// If an edge already exists, its weight is updated. Otherwise, this method attempts
    /// to add a new edge, respecting the bounded constraints.
    ///
    /// # Returns
    ///
    /// * `Ok(EdgeIndex)` - The index of the updated or newly added edge
    /// * `Err(BoundedGraphError)` - If adding a new edge would exceed capacity limits
    #[must_use = "edge update may fail due to capacity constraints"]
    pub fn update_edge(
        &mut self,
        a: NodeIndex<Ix>,
        b: NodeIndex<Ix>,
        weight: E,
    ) -> Result<EdgeIndex<Ix>, BoundedGraphError<Ix>>
    where
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdges + IntoEdgesDirected,
    {
        // Check if edge already exists using petgraph's edge finding
        let edge_exists = (&self.graph).edges(a).any(|e| e.target() == b);

        if edge_exists {
            // Edge exists, update it directly (won't fail capacity check)
            Ok(self.graph.update_edge(a, b, weight))
        } else {
            // Edge doesn't exist, use add_edge which checks capacity
            self.add_edge(a, b, weight)
        }
    }

    /// Extends the graph by adding edges from an iterator.
    ///
    /// Creates default nodes as needed to accommodate the edges. Silently skips
    /// edges that would violate the bounded constraints.
    ///
    /// # Note
    ///
    /// This method will continue processing edges even if some cannot be added
    /// due to capacity constraints.
    pub fn extend_with_edges<I>(&mut self, iterable: I)
    where
        I: IntoIterator,
        I::Item: IntoWeightedEdge<E>,
        <I::Item as IntoWeightedEdge<E>>::NodeId: Into<NodeIndex<Ix>>,
        N: Default,
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdgesDirected,
    {
        let iter = iterable.into_iter();
        let (_low, _) = iter.size_hint();

        // Note: reserve_edges is a concrete method on Graph/StableGraph, not a trait method
        // We skip the reservation for generic code - it's an optimization, not correctness

        for elt in iter {
            let (source, target, weight) = elt.into_weighted_edge();
            let (source, target) = (source.into(), target.into());
            let nx = cmp::max(source, target);
            while nx.index() >= self.graph.node_count() {
                let _ = self.add_node(N::default());
            }
            let _ = self.add_edge(source, target, weight);
        }
    }

    /// Attempts to extend the graph with an iterator of edges.
    ///
    /// Similar to [`extend_with_edges`](Self::extend_with_edges), but returns the
    /// indices of successfully added edges or an error with partial results. Creates
    /// default nodes as needed.
    ///
    /// # Returns
    ///
    /// - `Ok(Vec<EdgeIndex<Ix>>)` - All edges were successfully added
    /// - `Err((Vec<EdgeIndex<Ix>>, BoundedGraphError))` - Partial success; returns the
    ///   edges that were added before the error occurred, along with the error
    #[must_use = "returns edge indices and potential errors that should be checked"]
    pub fn try_extend_with_edges<I>(
        &mut self,
        iterable: I,
    ) -> Result<Vec<EdgeIndex<Ix>>, (Vec<EdgeIndex<Ix>>, BoundedGraphError<Ix>)>
    where
        I: IntoIterator,
        I::Item: IntoWeightedEdge<E>,
        <I::Item as IntoWeightedEdge<E>>::NodeId: Into<NodeIndex<Ix>>,
        N: Default,
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdgesDirected,
    {
        let iter = iterable.into_iter();
        let (low, _) = iter.size_hint();

        let mut added_edges = Vec::with_capacity(low);

        for elt in iter {
            let (source, target, weight) = elt.into_weighted_edge();
            let (source, target) = (source.into(), target.into());
            let nx = cmp::max(source, target);
            while nx.index() >= self.graph.node_count() {
                let _ = self.add_node(N::default());
            }
            match self.add_edge(source, target, weight) {
                Ok(edge_idx) => added_edges.push(edge_idx),
                Err(err) => return Err((added_edges, err)),
            }
        }
        Ok(added_edges)
    }

    /// Creates a new bounded graph from an iterator of edges.
    ///
    /// Creates default nodes as needed. Silently skips edges that would violate
    /// the bounded constraints.
    #[must_use = "creates a new graph that should be used"]
    pub fn from_edges<I>(iterable: I) -> Self
    where
        I: IntoIterator,
        I::Item: IntoWeightedEdge<E>,
        <I::Item as IntoWeightedEdge<E>>::NodeId: Into<NodeIndex<Ix>>,
        N: Default,
        G: Default,
        for<'a> &'a G: GraphBase<NodeId = NodeIndex<Ix>> + IntoEdgesDirected,
    {
        let mut g = Self::new();
        g.extend_with_edges(iterable);
        g
    }

    /// Returns the number of nodes in the graph.
    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    /// Returns the number of edges in the graph.
    pub fn edge_count(&self) -> usize {
        self.graph.edge_count()
    }

    /// Returns a reference to the underlying graph.
    #[must_use]
    pub fn inner(&self) -> &G {
        &self.graph
    }

    /// Consumes the wrapper and returns the underlying graph.
    #[must_use = "consumes self and returns the inner graph"]
    pub fn into_inner(self) -> G {
        self.graph
    }
}

// Default implementation
impl<N, E, G> Default for BoundedGraph<N, E, G>
where
    G: Default,
{
    fn default() -> Self {
        Self {
            graph: G::default(),
            _phantom: PhantomData,
        }
    }
}

// Specialized implementation for Graph to provide with_capacity
impl<N, E, Ty, Ix> BoundedGraph<N, E, Graph<N, E, Ty, Ix>>
where
    Ty: EdgeType,
    Ix: IndexType,
    N: BoundedNode<Ix>,
    E: Clone,
{
    /// Creates a new bounded graph with estimated capacity for the given number of nodes and edges.
    ///
    /// This pre-allocates space to avoid reallocations during graph construction.
    #[must_use = "creates a new graph that should be used"]
    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self {
            graph: Graph::with_capacity(nodes, edges),
            _phantom: PhantomData,
        }
    }
}

// Specialized implementation for StableGraph to provide with_capacity
impl<N, E, Ty, Ix> BoundedGraph<N, E, StableGraph<N, E, Ty, Ix>>
where
    Ty: EdgeType,
    Ix: IndexType,
    N: BoundedNode<Ix>,
    E: Clone,
{
    /// Creates a new bounded graph with estimated capacity for the given number of nodes and edges.
    ///
    /// This pre-allocates space to avoid reallocations during graph construction.
    #[must_use = "creates a new graph that should be used"]
    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
        Self {
            graph: StableGraph::with_capacity(nodes, edges),
            _phantom: PhantomData,
        }
    }
}

// TryFrom implementation for Graph
impl<N, E, Ty, Ix> TryFrom<Graph<N, E, Ty, Ix>> for BoundedGraph<N, E, Graph<N, E, Ty, Ix>>
where
    Ty: EdgeType,
    Ix: IndexType,
    N: BoundedNode<Ix>,
    E: Clone,
{
    type Error = BoundedGraphError<Ix>;

    /// Attempts to create a bounded graph from an existing petgraph Graph.
    ///
    /// This validates that all existing edges in the graph respect the node capacity constraints.
    /// If any edge would violate the bounds, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use petgraph::Graph;
    /// use bounded_graph::{BoundedGraph, FixedEdgeCount};
    /// use std::convert::TryFrom;
    ///
    /// let mut pg = Graph::new();
    /// let n1 = pg.add_node(FixedEdgeCount::<2, 2>::empty());
    /// let n2 = pg.add_node(FixedEdgeCount::<2, 2>::empty());
    /// pg.add_edge(n1, n2, ());
    ///
    /// // This succeeds because the edge respects bounds
    /// let bounded = BoundedGraph::try_from(pg).unwrap();
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if any existing edge violates the capacity constraints of its nodes.
    fn try_from(graph: Graph<N, E, Ty, Ix>) -> Result<Self, Self::Error> {
        // Validate all edges respect bounds
        for node in graph.node_indices() {
            let node_weight = graph
                .node_weight(node)
                .ok_or_else(|| BoundedGraphError::not_found(Some(node), None))?;

            // Count outgoing edges
            let outgoing_count = graph.edges_directed(node, Direction::Outgoing).count();

            // Count incoming edges
            let incoming_count = graph.edges_directed(node, Direction::Incoming).count();

            // Check outgoing capacity
            for edge in graph.edges_directed(node, Direction::Outgoing) {
                let target = edge.target();
                let target_weight = graph
                    .node_weight(target)
                    .ok_or_else(|| BoundedGraphError::not_found(None, Some(target)))?;

                if !node_weight.can_add_edge(Direction::Outgoing, outgoing_count - 1, target_weight)
                {
                    return Err(BoundedGraphError::source_rejected(node, target));
                }
            }

            // Check incoming capacity
            for edge in graph.edges_directed(node, Direction::Incoming) {
                let source = edge.source();
                let source_weight = graph
                    .node_weight(source)
                    .ok_or_else(|| BoundedGraphError::not_found(Some(source), None))?;

                if !node_weight.can_add_edge(Direction::Incoming, incoming_count - 1, source_weight)
                {
                    return Err(BoundedGraphError::target_rejected(source, node));
                }
            }
        }

        // All edges are valid, wrap the graph
        Ok(BoundedGraph {
            graph,
            _phantom: PhantomData,
        })
    }
}

// TryFrom implementation for StableGraph
impl<N, E, Ty, Ix> TryFrom<StableGraph<N, E, Ty, Ix>>
    for BoundedGraph<N, E, StableGraph<N, E, Ty, Ix>>
where
    Ty: EdgeType,
    Ix: IndexType,
    N: BoundedNode<Ix>,
    E: Clone,
{
    type Error = BoundedGraphError<Ix>;

    /// Attempts to create a bounded graph from an existing petgraph StableGraph.
    ///
    /// This validates that all existing edges in the graph respect the node capacity constraints.
    /// If any edge would violate the bounds, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use petgraph::stable_graph::StableGraph;
    /// use bounded_graph::{BoundedGraph, FixedEdgeCount};
    /// use std::convert::TryFrom;
    ///
    /// let mut pg = StableGraph::new();
    /// let n1 = pg.add_node(FixedEdgeCount::<2, 2>::empty());
    /// let n2 = pg.add_node(FixedEdgeCount::<2, 2>::empty());
    /// pg.add_edge(n1, n2, ());
    ///
    /// // This succeeds because the edge respects bounds
    /// let bounded = BoundedGraph::try_from(pg).unwrap();
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if any existing edge violates the capacity constraints of its nodes.
    fn try_from(graph: StableGraph<N, E, Ty, Ix>) -> Result<Self, Self::Error> {
        // Validate all edges respect bounds
        for node in graph.node_indices() {
            let node_weight = graph
                .node_weight(node)
                .ok_or_else(|| BoundedGraphError::not_found(Some(node), None))?;

            // Count outgoing edges
            let outgoing_count = graph.edges_directed(node, Direction::Outgoing).count();

            // Count incoming edges
            let incoming_count = graph.edges_directed(node, Direction::Incoming).count();

            // Check outgoing capacity
            for edge in graph.edges_directed(node, Direction::Outgoing) {
                let target = edge.target();
                let target_weight = graph
                    .node_weight(target)
                    .ok_or_else(|| BoundedGraphError::not_found(None, Some(target)))?;

                if !node_weight.can_add_edge(Direction::Outgoing, outgoing_count - 1, target_weight)
                {
                    return Err(BoundedGraphError::source_rejected(node, target));
                }
            }

            // Check incoming capacity
            for edge in graph.edges_directed(node, Direction::Incoming) {
                let source = edge.source();
                let source_weight = graph
                    .node_weight(source)
                    .ok_or_else(|| BoundedGraphError::not_found(Some(source), None))?;

                if !node_weight.can_add_edge(Direction::Incoming, incoming_count - 1, source_weight)
                {
                    return Err(BoundedGraphError::target_rejected(source, node));
                }
            }
        }

        // All edges are valid, wrap the graph
        Ok(BoundedGraph {
            graph,
            _phantom: PhantomData,
        })
    }
}