somatize_core/

graph.rs

//! Computational graph — DAG of filter nodes connected by edges.
//!
//! The graph is the user-facing representation of a pipeline topology.
//! It gets compiled into an `ExecutionPlan` by the compiler.

use crate::control::LoopCondition;
use crate::error::{Result, SomaError};
use crate::strategy::TrainingStrategy;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

/// Unique identifier for a node in a graph.
///
/// Currently a type alias. Will be promoted to a newtype in a future version
/// for stronger type safety. Deliberately deferred — see the
/// "NodeId stays a String" entry in docs design/decisions.
pub type NodeId = String;

/// Unique identifier for an edge in a graph.
pub type EdgeId = String;

/// What kind of computation a node represents.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum NodeKind {
    /// A single filter (the common case).
    Filter {
        /// Name the filter is registered under in the `NodeCatalog`.
        filter_name: String,
    },
    /// A nested sub-graph (compiled recursively).
    SubGraph {
        /// The inner graph, boxed to keep `NodeKind` a fixed size.
        graph: Box<Graph>,
    },
    /// A loop node. Its body is the sub-graph reached through its *control*
    /// edges; `until` names what decides to stop.
    Loop {
        /// Hard cap on iterations; `None` means the body runs until `until`
        /// signals stop.
        max_iterations: Option<usize>,
        /// Defaults to [`LoopCondition::BodyTerminal`], resolved by the
        /// compiler. Never inferred at runtime from execution order.
        #[serde(default)]
        until: LoopCondition,
    },
    /// A branch/conditional node. Arms are the labelled control edges
    /// leaving it; `arms` optionally declares the complete set of labels the
    /// condition may produce, so the compiler can catch a mislabelled edge
    /// before the run rather than at the moment the branch is taken.
    Branch {
        /// Declared labels. Empty means "infer from the edges" — the
        /// backwards-compatible default.
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        arms: Vec<String>,
    },
    /// An effectful node: calls models, tools, or other graphs, and decides
    /// what happens next. See [`crate::step::Step`].
    Step {
        /// Name the step is registered under in the `NodeCatalog`.
        step_name: String,
    },
}

/// A node in the computational graph.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node {
    /// Unique id within the graph; edges and trained states refer to it.
    pub id: NodeId,
    /// Human-readable name shown in diagrams; cosmetic, excluded from the
    /// architecture fingerprint.
    pub label: String,
    /// What kind of computation this node represents.
    pub kind: NodeKind,
    /// Execution target: "local" (reserved, always local), or a worker tag.
    /// None means: use default (remote if workers available, else local).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub target: Option<String>,
}

impl Node {
    /// Create a filter node (backward-compatible with old 3-arg constructor).
    pub fn new(
        id: impl Into<String>,
        label: impl Into<String>,
        filter_name: impl Into<String>,
    ) -> Self {
        Self {
            id: id.into(),
            label: label.into(),
            kind: NodeKind::Filter {
                filter_name: filter_name.into(),
            },
            target: None,
        }
    }

    /// Create a filter node with explicit id and filter_name.
    pub fn filter_with_id(id: impl Into<String>, filter_name: impl Into<String>) -> Self {
        let id = id.into();
        Self {
            label: id.clone(),
            id,
            kind: NodeKind::Filter {
                filter_name: filter_name.into(),
            },
            target: None,
        }
    }

    /// Create a filter node where id defaults to filter_name.
    pub fn filter(filter_name: impl Into<String>) -> Self {
        let name = filter_name.into();
        Self {
            id: name.clone(),
            label: name.clone(),
            kind: NodeKind::Filter { filter_name: name },
            target: None,
        }
    }

    /// Create a sub-graph node.
    pub fn subgraph(id: impl Into<String>, graph: Graph) -> Self {
        let id = id.into();
        Self {
            id: id.clone(),
            label: id,
            kind: NodeKind::SubGraph {
                graph: Box::new(graph),
            },
            target: None,
        }
    }

    /// Create a loop node whose stop condition is its body's terminal node.
    pub fn loop_node(id: impl Into<String>, max_iterations: Option<usize>) -> Self {
        Self::loop_until(id, max_iterations, LoopCondition::BodyTerminal)
    }

    /// Create a loop node with an explicit stop condition.
    pub fn loop_until(
        id: impl Into<String>,
        max_iterations: Option<usize>,
        until: LoopCondition,
    ) -> Self {
        let id = id.into();
        Self {
            id: id.clone(),
            label: id,
            kind: NodeKind::Loop {
                max_iterations,
                until,
            },
            target: None,
        }
    }

    /// Create an effectful step node.
    pub fn step(id: impl Into<String>, step_name: impl Into<String>) -> Self {
        let id = id.into();
        Self {
            label: id.clone(),
            id,
            kind: NodeKind::Step {
                step_name: step_name.into(),
            },
            target: None,
        }
    }

    /// Create a branch node whose arms are inferred from its control edges.
    pub fn branch(id: impl Into<String>) -> Self {
        Self::branch_over(id, Vec::<String>::new())
    }

    /// Create a branch node declaring the labels its condition may produce.
    ///
    /// The compiler then checks the edges against this list in both
    /// directions: a declared arm with no edge, or an edge labelling an arm
    /// that was never declared, is a compile error rather than a branch that
    /// silently never fires.
    pub fn branch_over(
        id: impl Into<String>,
        arms: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        let id = id.into();
        Self {
            id: id.clone(),
            label: id,
            kind: NodeKind::Branch {
                arms: arms.into_iter().map(Into::into).collect(),
            },
            target: None,
        }
    }

    /// Set the execution target for this node.
    pub fn with_target(mut self, target: impl Into<String>) -> Self {
        self.target = Some(target.into());
        self
    }

    /// Whether this node is forced local.
    pub fn is_local(&self) -> bool {
        self.target.as_deref() == Some("local")
    }

    /// Get the filter name if this is a Filter node.
    pub fn filter_name(&self) -> Option<&str> {
        match &self.kind {
            NodeKind::Filter { filter_name } => Some(filter_name),
            _ => None,
        }
    }
}

/// Type of connection between nodes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum EdgeKind {
    /// Normal data flow: output of source becomes input of target.
    Data,
    /// Control flow edge (for conditional/loop logic).
    Control,
}

/// A directed edge connecting two nodes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Edge {
    /// Unique id within the graph; cosmetic — excluded from the
    /// architecture fingerprint.
    pub id: EdgeId,
    /// The node this edge leaves.
    pub source: NodeId,
    /// The node this edge enters.
    pub target: NodeId,
    /// Whether the edge carries data or control.
    pub kind: EdgeKind,
    /// Optional label; on an edge leaving a `Branch` node it names the arm.
    pub label: Option<String>,
}

impl Edge {
    /// Create a data edge: `source`'s output becomes an input of `target`.
    ///
    /// This is what `Graph::connect` builds, and what input resolution
    /// follows — a node's inputs are the outputs of its data predecessors,
    /// not "whatever ran last".
    pub fn data(
        id: impl Into<String>,
        source: impl Into<String>,
        target: impl Into<String>,
    ) -> Self {
        Self {
            id: id.into(),
            source: source.into(),
            target: target.into(),
            kind: EdgeKind::Data,
            label: None,
        }
    }

    /// Create a control edge: `source` decides whether `target` runs, but
    /// hands it no data.
    ///
    /// Control edges are how the compiler claims loop bodies and branch
    /// arms (by dominance); a branch passes its *input* to the chosen arm,
    /// not the selector's output.
    pub fn control(
        id: impl Into<String>,
        source: impl Into<String>,
        target: impl Into<String>,
    ) -> Self {
        Self {
            id: id.into(),
            source: source.into(),
            target: target.into(),
            kind: EdgeKind::Control,
            label: None,
        }
    }

    /// Attach a label. On an edge leaving a `Branch` node the label names the
    /// arm; the branch condition's value is matched against it.
    pub fn with_label(mut self, label: impl Into<String>) -> Self {
        self.label = Some(label.into());
        self
    }
}

/// A directed graph of computational nodes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Graph {
    /// The nodes, in insertion order (execution order comes from the edges).
    pub nodes: Vec<Node>,
    /// The directed edges connecting them.
    pub edges: Vec<Edge>,
    /// Training strategy for distributed execution.
    /// Inherited by subgraphs unless overridden.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub training_strategy: Option<TrainingStrategy>,
}

impl Graph {
    /// Create an empty graph with no training strategy set.
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            edges: Vec::new(),
            training_strategy: None,
        }
    }

    /// Set the training strategy for this graph.
    pub fn with_strategy(mut self, strategy: TrainingStrategy) -> Self {
        self.training_strategy = Some(strategy);
        self
    }

    /// Set the training strategy (mutable).
    pub fn set_strategy(&mut self, strategy: TrainingStrategy) {
        self.training_strategy = Some(strategy);
    }

    /// Get the effective training strategy (defaults to Local).
    pub fn effective_strategy(&self) -> &TrainingStrategy {
        static LOCAL: TrainingStrategy = TrainingStrategy::Local;
        self.training_strategy.as_ref().unwrap_or(&LOCAL)
    }

    /// Add a node. Duplicate ids are not checked here; [`Self::validate`]
    /// rejects them at compile time.
    pub fn add_node(&mut self, node: Node) {
        self.nodes.push(node);
    }

    /// Add a filter node using the filter name as the node id.
    /// If a node with that name already exists, appends a suffix.
    pub fn add_filter(&mut self, filter_name: impl Into<String>) -> &str {
        let name = filter_name.into();
        let id = if self.nodes.iter().any(|n| n.id == name) {
            let mut i = 2;
            loop {
                let candidate = format!("{name}_{i}");
                if !self.nodes.iter().any(|n| n.id == candidate) {
                    break candidate;
                }
                i += 1;
            }
        } else {
            name.clone()
        };
        self.nodes.push(Node::filter_with_id(&id, &name));
        &self.nodes.last().unwrap().id
    }

    /// Add an edge. Endpoints are not checked here; [`Self::validate`]
    /// rejects edges to unknown nodes at compile time.
    pub fn add_edge(&mut self, edge: Edge) {
        self.edges.push(edge);
    }

    /// Connect two nodes with a data edge (auto-generates edge id).
    pub fn connect(&mut self, source: impl Into<String>, target: impl Into<String>) {
        let id = format!("e_{}", self.edges.len());
        self.edges.push(Edge::data(id, source, target));
    }

    /// Get a node by its ID.
    pub fn node(&self, id: &str) -> Option<&Node> {
        self.nodes.iter().find(|n| n.id == id)
    }

    /// Get all node IDs.
    pub fn node_ids(&self) -> Vec<&str> {
        self.nodes.iter().map(|n| n.id.as_str()).collect()
    }

    /// Get predecessors of a node (nodes with edges pointing to it).
    pub fn predecessors(&self, node_id: &str) -> Vec<&str> {
        self.edges
            .iter()
            .filter(|e| e.target == node_id)
            .map(|e| e.source.as_str())
            .collect()
    }

    /// Get successors of a node (nodes it points to).
    pub fn successors(&self, node_id: &str) -> Vec<&str> {
        self.edges
            .iter()
            .filter(|e| e.source == node_id)
            .map(|e| e.target.as_str())
            .collect()
    }

    /// Find root nodes (no incoming edges).
    pub fn roots(&self) -> Vec<&str> {
        let has_incoming: HashSet<&str> = self.edges.iter().map(|e| e.target.as_str()).collect();
        self.nodes
            .iter()
            .filter(|n| !has_incoming.contains(n.id.as_str()))
            .map(|n| n.id.as_str())
            .collect()
    }

    /// Find leaf nodes (no outgoing edges).
    pub fn leaves(&self) -> Vec<&str> {
        let has_outgoing: HashSet<&str> = self.edges.iter().map(|e| e.source.as_str()).collect();
        self.nodes
            .iter()
            .filter(|n| !has_outgoing.contains(n.id.as_str()))
            .map(|n| n.id.as_str())
            .collect()
    }

    /// Compute in-degree for each node.
    fn in_degrees(&self) -> HashMap<&str, usize> {
        let mut degrees: HashMap<&str, usize> =
            self.nodes.iter().map(|n| (n.id.as_str(), 0)).collect();
        for edge in &self.edges {
            *degrees.entry(edge.target.as_str()).or_insert(0) += 1;
        }
        degrees
    }

    /// Topological sort using Kahn's algorithm.
    /// Returns Err if the graph contains a cycle.
    pub fn topological_sort(&self) -> Result<Vec<&str>> {
        let mut in_deg = self.in_degrees();
        let mut queue: Vec<&str> = in_deg
            .iter()
            .filter(|(_, deg)| **deg == 0)
            .map(|(&id, _)| id)
            .collect();
        queue.sort(); // deterministic order

        let mut sorted = Vec::with_capacity(self.nodes.len());

        while let Some(node) = queue.pop() {
            sorted.push(node);
            let mut next = Vec::new();
            for succ in self.successors(node) {
                if let Some(deg) = in_deg.get_mut(succ) {
                    *deg -= 1;
                    if *deg == 0 {
                        next.push(succ);
                    }
                }
            }
            next.sort();
            // Insert at beginning so we process in deterministic order
            for n in next.into_iter().rev() {
                queue.push(n);
            }
        }

        if sorted.len() != self.nodes.len() {
            return Err(SomaError::CycleDetected);
        }

        Ok(sorted)
    }

    /// Validate the graph structure (recursively validates sub-graphs).
    pub fn validate(&self) -> Result<()> {
        // Check for duplicate node IDs
        let mut seen = HashSet::new();
        for node in &self.nodes {
            if !seen.insert(&node.id) {
                return Err(SomaError::Compilation(format!(
                    "duplicate node id: `{}`",
                    node.id
                )));
            }
        }

        // Check that all edge endpoints reference existing nodes
        let node_ids: HashSet<&str> = self.nodes.iter().map(|n| n.id.as_str()).collect();
        for edge in &self.edges {
            if !node_ids.contains(edge.source.as_str()) {
                return Err(SomaError::NodeNotFound(edge.source.clone()));
            }
            if !node_ids.contains(edge.target.as_str()) {
                return Err(SomaError::NodeNotFound(edge.target.clone()));
            }
        }

        // Check for cycles
        self.topological_sort()?;

        // Recursively validate sub-graphs
        for node in &self.nodes {
            if let NodeKind::SubGraph { graph } = &node.kind {
                graph.validate()?;
            }
        }

        Ok(())
    }

    /// Does this graph — or any sub-graph nested inside it — contain a step?
    ///
    /// A step calls models and tools, so a graph that contains one is not a
    /// deterministic function of its input. [`crate::effect::Effect::is_pure`]
    /// asks this before memoizing a graph effect by content.
    pub fn contains_steps(&self) -> bool {
        self.nodes.iter().any(|node| match &node.kind {
            NodeKind::Step { .. } => true,
            NodeKind::SubGraph { graph } => graph.contains_steps(),
            _ => false,
        })
    }
}

// ── Visualization ──

impl Graph {
    /// Render as a Mermaid diagram.
    ///
    /// ```text
    /// graph LR
    ///     scaler[scaler]
    ///     model[model]
    ///     scaler --> model
    /// ```
    pub fn to_mermaid(&self) -> String {
        self.to_mermaid_with(&crate::viz::GraphOverlay::default())
    }

    /// Render as a Mermaid diagram with per-node execution annotations.
    ///
    /// Each annotated node gets a second label line (duration, cache
    /// tier, health flags — see [`crate::viz::NodeOverlay::sublabel_text`])
    /// and a status `classDef` for coloring. An empty overlay produces
    /// exactly [`Graph::to_mermaid`]'s output.
    pub fn to_mermaid_with(&self, overlay: &crate::viz::GraphOverlay) -> String {
        use std::fmt::Write;
        let mut out = String::from("graph LR\n");
        for node in &self.nodes {
            let ov = overlay.nodes.get(&node.id);
            // A sublabel needs a quoted label to allow `<br/>`.
            let label_with = |base: &str| match ov.and_then(|o| o.sublabel_text()) {
                Some(sub) => format!("\"{base}<br/>{sub}\""),
                None => base.to_string(),
            };
            let shape = match &node.kind {
                NodeKind::Filter { .. } => {
                    format!("    {}[{}]", node.id, label_with(&node.label))
                }
                NodeKind::SubGraph { .. } => {
                    format!("    {}[[{}]]", node.id, label_with(&node.label))
                }
                NodeKind::Loop { max_iterations, .. } => {
                    let label = match max_iterations {
                        Some(n) => format!("{} (max {})", node.label, n),
                        None => node.label.clone(),
                    };
                    format!("    {}(({}))", node.id, label_with(&label))
                }
                NodeKind::Branch { .. } => {
                    format!("    {}{{{{{}}}}}", node.id, label_with(&node.label))
                }
                // Parallelogram: the I/O shape, which is what an effectful
                // node is — it reaches outside the graph.
                NodeKind::Step { .. } => {
                    format!("    {}[/{}/]", node.id, label_with(&node.label))
                }
            };
            let _ = writeln!(out, "{shape}");
        }
        for edge in &self.edges {
            let arrow = match edge.kind {
                EdgeKind::Data => "-->",
                EdgeKind::Control => "-.->",
            };
            if let Some(label) = &edge.label {
                let _ = writeln!(
                    out,
                    "    {} {}|{}| {}",
                    edge.source, arrow, label, edge.target
                );
            } else {
                let _ = writeln!(out, "    {} {} {}", edge.source, arrow, edge.target);
            }
        }
        let assignments: Vec<(&str, &'static str)> = self
            .nodes
            .iter()
            .filter_map(|n| {
                overlay
                    .nodes
                    .get(&n.id)
                    .and_then(|o| o.style_class())
                    .map(|class| (n.id.as_str(), class))
            })
            .collect();
        if !assignments.is_empty() {
            let mut used: Vec<&'static str> = assignments.iter().map(|(_, c)| *c).collect();
            used.sort_unstable();
            used.dedup();
            for class in used {
                let _ = writeln!(
                    out,
                    "    classDef {class} {}",
                    crate::viz::mermaid_class_style(class)
                );
            }
            for (id, class) in assignments {
                let _ = writeln!(out, "    class {id} {class}");
            }
        }
        out
    }

    /// Render as Graphviz DOT format.
    pub fn to_graphviz(&self) -> String {
        self.to_graphviz_with(&crate::viz::GraphOverlay::default())
    }

    /// Render as Graphviz DOT with per-node execution annotations:
    /// a second label line plus fill/border status colors. An empty
    /// overlay produces exactly [`Graph::to_graphviz`]'s output.
    pub fn to_graphviz_with(&self, overlay: &crate::viz::GraphOverlay) -> String {
        use std::fmt::Write;
        let mut out = String::from("digraph G {\n    rankdir=LR;\n");
        for node in &self.nodes {
            let shape = match &node.kind {
                NodeKind::Filter { .. } => "box",
                NodeKind::SubGraph { .. } => "doubleoctagon",
                NodeKind::Loop { .. } => "ellipse",
                NodeKind::Branch { .. } => "diamond",
                NodeKind::Step { .. } => "parallelogram",
            };
            let ov = overlay.nodes.get(&node.id);
            let label = match ov.and_then(|o| o.sublabel_text()) {
                Some(sub) => format!("{}\\n{}", node.label, sub),
                None => node.label.clone(),
            };
            let style = ov
                .and_then(|o| o.style_class())
                .map(crate::viz::dot_class_style)
                .unwrap_or_default();
            let _ = writeln!(
                out,
                "    \"{}\" [label=\"{}\" shape={}{}];",
                node.id, label, shape, style
            );
        }
        for edge in &self.edges {
            let style = match edge.kind {
                EdgeKind::Data => "",
                EdgeKind::Control => " [style=dashed]",
            };
            let label = edge
                .label
                .as_ref()
                .map(|l| format!(" [label=\"{l}\"]"))
                .unwrap_or_default();
            let attrs = if style.is_empty() && label.is_empty() {
                String::new()
            } else if label.is_empty() {
                style.to_string()
            } else {
                label
            };
            let _ = writeln!(
                out,
                "    \"{}\" -> \"{}\"{};",
                edge.source, edge.target, attrs
            );
        }
        out.push_str("}\n");
        out
    }

    /// Render as an ASCII text tree for terminal display.
    pub fn to_text(&self) -> String {
        use std::fmt::Write;
        let mut out = String::new();
        let sorted = self.topological_sort().unwrap_or_default();
        let total_nodes = self.nodes.len();
        let total_edges = self.edges.len();
        let _ = writeln!(out, "Graph ({total_nodes} nodes, {total_edges} edges)");

        for (i, node_id) in sorted.iter().enumerate() {
            let node = match self.node(node_id) {
                Some(n) => n,
                None => continue,
            };
            let is_last = i == sorted.len() - 1;
            let prefix = if is_last { "└── " } else { "├── " };
            let kind_tag = match &node.kind {
                NodeKind::Filter { filter_name } => {
                    if filter_name == &node.id {
                        String::new()
                    } else {
                        format!(" ({})", filter_name)
                    }
                }
                NodeKind::SubGraph { graph } => {
                    format!(" [subgraph: {} nodes]", graph.nodes.len())
                }
                NodeKind::Loop { max_iterations, .. } => match max_iterations {
                    Some(n) => format!(" [loop max={n}]"),
                    None => " [loop]".into(),
                },
                NodeKind::Branch { .. } => " [branch]".into(),
                NodeKind::Step { step_name } => format!(" [step: {step_name}]"),
            };
            let preds = self.predecessors(node_id);
            let pred_info = if preds.is_empty() {
                String::new()
            } else {
                format!(" ← {}", preds.join(", "))
            };
            let _ = writeln!(out, "{prefix}{}{kind_tag}{pred_info}", node.id);
        }
        out
    }
}

impl std::fmt::Display for Graph {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.to_text())
    }
}

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

/// Builder for constructing linear pipelines easily.
pub fn linear_pipeline(nodes: Vec<Node>) -> Graph {
    let mut graph = Graph::new();
    for (i, node) in nodes.iter().enumerate() {
        graph.add_node(node.clone());
        if i > 0 {
            graph.add_edge(Edge::data(format!("e_{}", i), &nodes[i - 1].id, &node.id));
        }
    }
    graph
}

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

    fn sample_linear_graph() -> Graph {
        linear_pipeline(vec![
            Node::new("a", "Scaler", "StandardScaler"),
            Node::new("b", "PCA", "PCA"),
            Node::new("c", "SVM", "SVM"),
        ])
    }

    #[test]
    fn linear_pipeline_structure() {
        let g = sample_linear_graph();
        assert_eq!(g.nodes.len(), 3);
        assert_eq!(g.edges.len(), 2);
    }

    #[test]
    fn roots_and_leaves() {
        let g = sample_linear_graph();
        assert_eq!(g.roots(), vec!["a"]);
        assert_eq!(g.leaves(), vec!["c"]);
    }

    #[test]
    fn predecessors_and_successors() {
        let g = sample_linear_graph();
        assert!(g.predecessors("a").is_empty());
        assert_eq!(g.predecessors("b"), vec!["a"]);
        assert_eq!(g.successors("a"), vec!["b"]);
        assert_eq!(g.successors("b"), vec!["c"]);
        assert!(g.successors("c").is_empty());
    }

    #[test]
    fn topological_sort_linear() {
        let g = sample_linear_graph();
        let sorted = g.topological_sort().unwrap();
        assert_eq!(sorted, vec!["a", "b", "c"]);
    }

    #[test]
    fn topological_sort_parallel() {
        let mut g = Graph::new();
        g.add_node(Node::new("root", "Root", "Input"));
        g.add_node(Node::new("b1", "Branch1", "F1"));
        g.add_node(Node::new("b2", "Branch2", "F2"));
        g.add_node(Node::new("merge", "Merge", "Merge"));
        g.add_edge(Edge::data("e1", "root", "b1"));
        g.add_edge(Edge::data("e2", "root", "b2"));
        g.add_edge(Edge::data("e3", "b1", "merge"));
        g.add_edge(Edge::data("e4", "b2", "merge"));

        let sorted = g.topological_sort().unwrap();
        // root must be first, merge must be last
        assert_eq!(sorted[0], "root");
        assert_eq!(sorted[3], "merge");
        // b1 and b2 can be in any order between root and merge
        let middle: HashSet<&str> = sorted[1..3].iter().copied().collect();
        assert!(middle.contains("b1"));
        assert!(middle.contains("b2"));
    }

    #[test]
    fn topological_sort_detects_cycle() {
        let mut g = Graph::new();
        g.add_node(Node::new("a", "A", "F"));
        g.add_node(Node::new("b", "B", "F"));
        g.add_edge(Edge::data("e1", "a", "b"));
        g.add_edge(Edge::data("e2", "b", "a")); // cycle!

        let result = g.topological_sort();
        assert!(matches!(result, Err(SomaError::CycleDetected)));
    }

    #[test]
    fn validate_accepts_valid_graph() {
        let g = sample_linear_graph();
        assert!(g.validate().is_ok());
    }

    #[test]
    fn validate_rejects_duplicate_ids() {
        let mut g = Graph::new();
        g.add_node(Node::new("a", "A", "F"));
        g.add_node(Node::new("a", "A2", "F"));
        assert!(matches!(g.validate(), Err(SomaError::Compilation(_))));
    }

    #[test]
    fn validate_rejects_missing_edge_target() {
        let mut g = Graph::new();
        g.add_node(Node::new("a", "A", "F"));
        g.add_edge(Edge::data("e1", "a", "nonexistent"));
        assert!(matches!(g.validate(), Err(SomaError::NodeNotFound(_))));
    }

    #[test]
    fn graph_serde_roundtrip() {
        let g = sample_linear_graph();
        let json = serde_json::to_string(&g).unwrap();
        let deserialized: Graph = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.nodes.len(), 3);
        assert_eq!(deserialized.edges.len(), 2);
    }

    #[test]
    fn empty_graph_is_valid() {
        let g = Graph::new();
        assert!(g.validate().is_ok());
        assert!(g.topological_sort().unwrap().is_empty());
    }

    #[test]
    fn single_node_graph() {
        let mut g = Graph::new();
        g.add_node(Node::new("solo", "Solo", "F"));
        assert_eq!(g.roots(), vec!["solo"]);
        assert_eq!(g.leaves(), vec!["solo"]);
        assert_eq!(g.topological_sort().unwrap(), vec!["solo"]);
    }

    // ── NodeKind tests ──

    #[test]
    fn node_filter_shorthand() {
        let n = Node::filter("StandardScaler");
        assert_eq!(n.id, "StandardScaler");
        assert_eq!(n.filter_name(), Some("StandardScaler"));
    }

    #[test]
    fn node_filter_with_id() {
        let n = Node::filter_with_id("my_scaler", "StandardScaler");
        assert_eq!(n.id, "my_scaler");
        assert_eq!(n.filter_name(), Some("StandardScaler"));
    }

    #[test]
    fn graph_add_filter_auto_names() {
        let mut g = Graph::new();
        g.add_filter("Scaler");
        g.add_filter("PCA");
        g.connect("Scaler", "PCA");

        assert!(g.validate().is_ok());
        assert_eq!(g.nodes.len(), 2);
        assert_eq!(g.nodes[0].id, "Scaler");
        assert_eq!(g.nodes[1].id, "PCA");
    }

    #[test]
    fn graph_add_filter_deduplicates() {
        let mut g = Graph::new();
        g.add_filter("Scaler");
        g.add_filter("Scaler"); // duplicate name → gets suffix

        assert_eq!(g.nodes.len(), 2);
        assert_eq!(g.nodes[0].id, "Scaler");
        assert_eq!(g.nodes[1].id, "Scaler_2");
    }

    #[test]
    fn subgraph_node() {
        let inner = linear_pipeline(vec![Node::new("a", "A", "F"), Node::new("b", "B", "F")]);

        let mut outer = Graph::new();
        outer.add_node(Node::new("input", "Input", "Input"));
        outer.add_node(Node::subgraph("pipeline", inner));
        outer.add_node(Node::new("output", "Output", "Output"));
        outer.add_edge(Edge::data("e1", "input", "pipeline"));
        outer.add_edge(Edge::data("e2", "pipeline", "output"));

        assert!(outer.validate().is_ok());
        assert_eq!(outer.nodes.len(), 3);

        // SubGraph node has no filter_name
        assert!(outer.node("pipeline").unwrap().filter_name().is_none());
    }

    #[test]
    fn loop_and_branch_nodes() {
        let mut g = Graph::new();
        g.add_node(Node::loop_node("train_loop", Some(100)));
        g.add_node(Node::branch("check_convergence"));
        g.add_edge(Edge::data("e1", "train_loop", "check_convergence"));

        assert!(g.validate().is_ok());
        assert!(matches!(
            g.node("train_loop").unwrap().kind,
            NodeKind::Loop {
                max_iterations: Some(100),
                ..
            }
        ));
        assert!(matches!(
            g.node("check_convergence").unwrap().kind,
            NodeKind::Branch { .. }
        ));
    }

    // ── Visualization tests ──

    #[test]
    fn to_mermaid_linear() {
        let g = sample_linear_graph();
        let m = g.to_mermaid();
        assert!(m.starts_with("graph LR"));
        assert!(m.contains("a[Scaler]"));
        assert!(m.contains("b[PCA]"));
        assert!(m.contains("c[SVM]"));
        assert!(m.contains("a --> b"));
        assert!(m.contains("b --> c"));
    }

    #[test]
    fn to_mermaid_branch_and_loop() {
        let mut g = Graph::new();
        g.add_node(Node::loop_node("train", Some(100)));
        g.add_node(Node::branch("check"));
        g.add_edge(Edge::data("e1", "train", "check"));

        let m = g.to_mermaid();
        assert!(m.contains("train((train (max 100)))"));
        assert!(m.contains("check{"));
        assert!(m.contains("train --> check"));
    }

    #[test]
    fn to_graphviz_output() {
        let g = sample_linear_graph();
        let dot = g.to_graphviz();
        assert!(dot.starts_with("digraph G {"));
        assert!(dot.contains("rankdir=LR"));
        assert!(dot.contains("\"a\" [label=\"Scaler\" shape=box]"));
        assert!(dot.contains("\"a\" -> \"b\""));
        assert!(dot.ends_with("}\n"));
    }

    #[test]
    fn overlay_empty_is_identical_to_plain_rendering() {
        use crate::viz::GraphOverlay;
        let g = sample_linear_graph();
        assert_eq!(g.to_mermaid(), g.to_mermaid_with(&GraphOverlay::default()));
        assert_eq!(
            g.to_graphviz(),
            g.to_graphviz_with(&GraphOverlay::default())
        );
        // No classDef/style leaks into the plain rendering.
        assert!(!g.to_mermaid().contains("classDef"));
        assert!(!g.to_graphviz().contains("fillcolor"));
    }

    #[test]
    fn to_mermaid_with_overlay_annotates_and_styles() {
        use crate::viz::{GraphOverlay, NodeOverlay, NodeStatus};
        let g = sample_linear_graph();
        let mut ov = GraphOverlay::default();
        ov.nodes.insert(
            "a".into(),
            NodeOverlay {
                status: Some(NodeStatus::Completed),
                duration_ms: Some(1_200),
                ..Default::default()
            },
        );
        ov.nodes.insert(
            "b".into(),
            NodeOverlay {
                status: Some(NodeStatus::Cached),
                duration_ms: Some(3),
                cache_tier: Some("memory".into()),
                ..Default::default()
            },
        );
        ov.nodes.insert(
            "c".into(),
            NodeOverlay {
                status: Some(NodeStatus::Completed),
                flags: vec!["LEAKAGE".into()],
                ..Default::default()
            },
        );

        let m = g.to_mermaid_with(&ov);
        assert!(m.contains("a[\"Scaler<br/>1.2s\"]"), "{m}");
        assert!(m.contains("b[\"PCA<br/>3ms · mem hit\"]"), "{m}");
        assert!(m.contains("c[\"SVM<br/>⚠ LEAKAGE\"]"), "{m}");
        assert!(m.contains("classDef soma_completed"));
        assert!(m.contains("classDef soma_cached"));
        assert!(m.contains("classDef soma_flagged"));
        assert!(m.contains("class a soma_completed"));
        assert!(m.contains("class b soma_cached"));
        assert!(m.contains("class c soma_flagged"), "flags win over status");
        // Edges unchanged.
        assert!(m.contains("a --> b"));
    }

    #[test]
    fn to_mermaid_with_overlay_ignores_unknown_nodes() {
        use crate::viz::{GraphOverlay, NodeOverlay, NodeStatus};
        let g = sample_linear_graph();
        let mut ov = GraphOverlay::default();
        ov.nodes.insert(
            "ghost".into(),
            NodeOverlay {
                status: Some(NodeStatus::Failed),
                ..Default::default()
            },
        );
        let m = g.to_mermaid_with(&ov);
        assert_eq!(m, g.to_mermaid(), "unknown node ids change nothing");
    }

    #[test]
    fn to_graphviz_with_overlay_annotates_and_styles() {
        use crate::viz::{GraphOverlay, NodeOverlay, NodeStatus};
        let g = sample_linear_graph();
        let mut ov = GraphOverlay::default();
        ov.nodes.insert(
            "a".into(),
            NodeOverlay {
                status: Some(NodeStatus::Failed),
                ..Default::default()
            },
        );
        ov.nodes.insert(
            "b".into(),
            NodeOverlay {
                flags: vec!["DEAD_CHANNELS".into()],
                ..Default::default()
            },
        );
        let dot = g.to_graphviz_with(&ov);
        assert!(
            dot.contains("\"a\" [label=\"Scaler\\nfailed\" shape=box"),
            "{dot}"
        );
        assert!(dot.contains("fillcolor=\"#ffebee\""), "failed fill: {dot}");
        assert!(dot.contains("penwidth=3"), "flagged border: {dot}");
        // Unannotated node keeps the plain attribute set.
        assert!(dot.contains("\"c\" [label=\"SVM\" shape=box];"));
    }

    #[test]
    fn to_text_output() {
        let g = sample_linear_graph();
        let text = g.to_text();
        assert!(text.contains("Graph (3 nodes, 2 edges)"));
        assert!(text.contains("a"));
        assert!(text.contains("b"));
        assert!(text.contains("c"));
        assert!(text.contains("← a"));
    }

    #[test]
    fn display_trait() {
        let g = sample_linear_graph();
        let s = format!("{g}");
        assert!(s.contains("Graph (3 nodes"));
    }

    #[test]
    fn node_kind_serde_roundtrip() {
        let inner = linear_pipeline(vec![Node::new("x", "X", "F")]);
        let nodes = vec![
            Node::filter("Scaler"),
            Node::subgraph("sub", inner),
            Node::loop_node("loop", Some(50)),
            Node::branch("cond"),
        ];

        for node in &nodes {
            let json = serde_json::to_string(node).unwrap();
            let parsed: Node = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed.id, node.id);
        }
    }
}