impl DependencyGraph {
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    /// Create a new instance.
    pub fn new() -> Self {
        Self {
            nodes: FxHashMap::default(),
            edges: Vec::new(),
        }
    }

    /// Adds a node to the dependency graph
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, NodeInfo, NodeType};
    /// use rustc_hash::FxHashMap;
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_node(NodeInfo {
    ///     id: "main::hello".to_string(),
    ///     label: "hello".to_string(),
    ///     node_type: NodeType::Function,
    ///     file_path: "src/main.rs".to_string(),
    ///     line_number: 10,
    ///     complexity: 1,
    ///     metadata: FxHashMap::default(),
    /// });
    ///
    /// assert_eq!(graph.nodes.len(), 1);
    /// ```
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn add_node(&mut self, node: NodeInfo) {
        self.nodes.insert(node.id.clone(), node);
    }

    /// Get the number of nodes in the graph
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of edges in the graph  
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }

    /// Adds an edge between two nodes in the dependency graph
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, Edge, EdgeType};
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_edge(Edge {
    ///     from: "main::hello".to_string(),
    ///     to: "utils::print".to_string(),
    ///     edge_type: EdgeType::Calls,
    ///     weight: 1,
    /// });
    ///
    /// assert_eq!(graph.edges.len(), 1);
    /// assert_eq!(graph.edges[0].edge_type, EdgeType::Calls);
    /// ```
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn add_edge(&mut self, edge: Edge) {
        self.edges.push(edge);
    }

    /// Creates a new graph containing only edges of the specified type
    ///
    /// # Examples
    ///
    /// ```rust
    /// use pmat::models::dag::{DependencyGraph, Edge, EdgeType};
    ///
    /// let mut graph = DependencyGraph::new();
    /// graph.add_edge(Edge {
    ///     from: "a".to_string(),
    ///     to: "b".to_string(),
    ///     edge_type: EdgeType::Calls,
    ///     weight: 1,
    /// });
    /// graph.add_edge(Edge {
    ///     from: "c".to_string(),
    ///     to: "d".to_string(),
    ///     edge_type: EdgeType::Imports,
    ///     weight: 1,
    /// });
    ///
    /// let calls_only = graph.filter_by_edge_type(EdgeType::Calls);
    /// assert_eq!(calls_only.edges.len(), 1);
    /// assert_eq!(calls_only.edges[0].edge_type, EdgeType::Calls);
    /// ```
    #[must_use]
    #[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
    pub fn filter_by_edge_type(&self, edge_type: EdgeType) -> Self {
        let filtered_edges: Vec<Edge> = self
            .edges
            .iter()
            .filter(|e| e.edge_type == edge_type)
            .cloned()
            .collect();

        // If filtering results in no edges but we originally had edges,
        // only include nodes that were connected by the filtered edge type
        if filtered_edges.is_empty() && !self.edges.is_empty() {
            // Return empty nodes since no nodes are connected by this edge type
            return Self {
                nodes: FxHashMap::default(),
                edges: filtered_edges,
            };
        }

        // If we have no edges at all, return all nodes
        if self.edges.is_empty() {
            return Self {
                nodes: self.nodes.clone(),
                edges: Vec::new(),
            };
        }

        // Otherwise, filter nodes to only those connected by the filtered edges
        let used_nodes: FxHashSet<String> = filtered_edges
            .iter()
            .flat_map(|e| vec![e.from.clone(), e.to.clone()])
            .collect();

        let filtered_nodes: FxHashMap<String, NodeInfo> = self
            .nodes
            .iter()
            .filter(|(id, _)| used_nodes.contains(*id))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        Self {
            nodes: filtered_nodes,
            edges: filtered_edges,
        }
    }
}

impl Default for DependencyGraph {
    fn default() -> Self {
        Self::new()
    }
}
