reasonkit-core 0.1.8

//! Graph-Based Execution Engine
//!
//! This module provides DAG-based workflow execution for ThinkTool
//! orchestration using petgraph and daggy.
//!
//! # Features
//! - DAG construction with cycle detection
//! - Topological execution ordering
//! - Parallel execution of independent nodes
//! - Dependency tracking
//!
//! Enable with: `cargo build --features graph-execution`

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// Re-exports
pub use daggy;
pub use petgraph;

use daggy::petgraph::visit::Topo;
use daggy::{Dag, NodeIndex, Walker};

/// A node in the execution graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionNode {
    /// Unique node identifier
    pub id: String,
    /// Node type (e.g., "thinktool", "transform", "merge")
    pub node_type: String,
    /// Configuration for this node
    pub config: serde_json::Value,
    /// Current state
    pub state: NodeState,
}

/// Node execution state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeState {
    /// Not yet started
    Pending,
    /// Currently executing
    Running,
    /// Completed successfully
    Completed,
    /// Failed with error
    Failed,
    /// Skipped (dependency failed)
    Skipped,
}

/// Edge weight representing data flow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataFlow {
    /// Source output name
    pub source_output: String,
    /// Target input name
    pub target_input: String,
}

/// Execution graph for workflow orchestration
pub struct ExecutionGraph {
    dag: Dag<ExecutionNode, DataFlow>,
    node_map: HashMap<String, NodeIndex>,
}

impl ExecutionGraph {
    /// Create a new empty execution graph
    pub fn new() -> Self {
        Self {
            dag: Dag::new(),
            node_map: HashMap::new(),
        }
    }

    /// Add a node to the graph
    pub fn add_node(&mut self, node: ExecutionNode) -> Result<NodeIndex> {
        let id = node.id.clone();
        if self.node_map.contains_key(&id) {
            anyhow::bail!("Node with ID '{}' already exists", id);
        }

        let idx = self.dag.add_node(node);
        self.node_map.insert(id, idx);
        Ok(idx)
    }

    /// Add an edge (dependency) between nodes
    pub fn add_edge(&mut self, from_id: &str, to_id: &str, data_flow: DataFlow) -> Result<()> {
        let from_idx = self
            .node_map
            .get(from_id)
            .ok_or_else(|| anyhow::anyhow!("Source node '{}' not found", from_id))?;
        let to_idx = self
            .node_map
            .get(to_id)
            .ok_or_else(|| anyhow::anyhow!("Target node '{}' not found", to_id))?;

        self.dag
            .add_edge(*from_idx, *to_idx, data_flow)
            .map_err(|_| anyhow::anyhow!("Adding edge would create a cycle"))?;

        Ok(())
    }

    /// Get nodes in topological order for execution
    pub fn execution_order(&self) -> Vec<&ExecutionNode> {
        let mut order = Vec::new();
        let mut topo = Topo::new(self.dag.graph());

        while let Some(idx) = topo.next(self.dag.graph()) {
            if let Some(node) = self.dag.node_weight(idx) {
                order.push(node);
            }
        }

        order
    }

    /// Get nodes ready for execution (all dependencies completed)
    pub fn ready_nodes(&self) -> Vec<&ExecutionNode> {
        let mut ready = Vec::new();

        for (idx, node) in self
            .dag
            .graph()
            .node_indices()
            .zip(self.dag.raw_nodes().iter())
        {
            if node.weight.state != NodeState::Pending {
                continue;
            }

            // Check if all parents are completed
            let all_parents_done = self
                .dag
                .parents(idx)
                .iter(&self.dag)
                .all(|(_, parent_idx)| {
                    self.dag
                        .node_weight(parent_idx)
                        .map(|p| p.state == NodeState::Completed)
                        .unwrap_or(false)
                });

            if all_parents_done {
                ready.push(&node.weight);
            }
        }

        ready
    }

    /// Update node state
    pub fn set_state(&mut self, node_id: &str, state: NodeState) -> Result<()> {
        let idx = self
            .node_map
            .get(node_id)
            .ok_or_else(|| anyhow::anyhow!("Node '{}' not found", node_id))?;

        if let Some(node) = self.dag.node_weight_mut(*idx) {
            node.state = state;
        }

        Ok(())
    }

    /// Get node by ID
    pub fn get_node(&self, node_id: &str) -> Option<&ExecutionNode> {
        self.node_map
            .get(node_id)
            .and_then(|idx| self.dag.node_weight(*idx))
    }

    /// Get mutable node by ID
    pub fn get_node_mut(&mut self, node_id: &str) -> Option<&mut ExecutionNode> {
        self.node_map
            .get(node_id)
            .copied()
            .and_then(move |idx| self.dag.node_weight_mut(idx))
    }

    /// Get all direct dependencies of a node
    pub fn dependencies(&self, node_id: &str) -> Vec<&ExecutionNode> {
        let Some(idx) = self.node_map.get(node_id) else {
            return Vec::new();
        };

        self.dag
            .parents(*idx)
            .iter(&self.dag)
            .filter_map(|(_, parent_idx)| self.dag.node_weight(parent_idx))
            .collect()
    }

    /// Get all nodes that depend on this node
    pub fn dependents(&self, node_id: &str) -> Vec<&ExecutionNode> {
        let Some(idx) = self.node_map.get(node_id) else {
            return Vec::new();
        };

        self.dag
            .children(*idx)
            .iter(&self.dag)
            .filter_map(|(_, child_idx)| self.dag.node_weight(child_idx))
            .collect()
    }

    /// Get total node count
    pub fn node_count(&self) -> usize {
        self.dag.node_count()
    }

    /// Get edge count
    pub fn edge_count(&self) -> usize {
        self.dag.edge_count()
    }

    /// Check if the graph is empty
    pub fn is_empty(&self) -> bool {
        self.dag.node_count() == 0
    }

    /// Check if all nodes are completed
    pub fn is_complete(&self) -> bool {
        self.dag.raw_nodes().iter().all(|n| {
            matches!(
                n.weight.state,
                NodeState::Completed | NodeState::Skipped | NodeState::Failed
            )
        })
    }
}

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

/// Builder for creating execution graphs
pub struct GraphBuilder {
    graph: ExecutionGraph,
}

impl GraphBuilder {
    /// Create a new graph builder
    pub fn new() -> Self {
        Self {
            graph: ExecutionGraph::new(),
        }
    }

    /// Add a ThinkTool node
    pub fn thinktool(mut self, id: impl Into<String>, tool_name: impl Into<String>) -> Self {
        let _ = self.graph.add_node(ExecutionNode {
            id: id.into(),
            node_type: "thinktool".to_string(),
            config: serde_json::json!({ "tool": tool_name.into() }),
            state: NodeState::Pending,
        });
        self
    }

    /// Add a transform node
    pub fn transform(mut self, id: impl Into<String>, transform_type: impl Into<String>) -> Self {
        let _ = self.graph.add_node(ExecutionNode {
            id: id.into(),
            node_type: "transform".to_string(),
            config: serde_json::json!({ "type": transform_type.into() }),
            state: NodeState::Pending,
        });
        self
    }

    /// Add an edge with default data flow
    pub fn edge(mut self, from: &str, to: &str) -> Self {
        let _ = self.graph.add_edge(
            from,
            to,
            DataFlow {
                source_output: "output".to_string(),
                target_input: "input".to_string(),
            },
        );
        self
    }

    /// Add an edge with named ports
    pub fn edge_named(mut self, from: &str, output: &str, to: &str, input: &str) -> Self {
        let _ = self.graph.add_edge(
            from,
            to,
            DataFlow {
                source_output: output.to_string(),
                target_input: input.to_string(),
            },
        );
        self
    }

    /// Build the execution graph
    pub fn build(self) -> ExecutionGraph {
        self.graph
    }
}

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

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

    #[test]
    fn test_graph_construction() {
        let graph = GraphBuilder::new()
            .thinktool("gigathink", "GigaThink")
            .thinktool("laserlogic", "LaserLogic")
            .transform("merge", "combine")
            .edge("gigathink", "merge")
            .edge("laserlogic", "merge")
            .build();

        assert_eq!(graph.node_count(), 3);
        assert_eq!(graph.edge_count(), 2);
    }

    #[test]
    fn test_execution_order() {
        let graph = GraphBuilder::new()
            .thinktool("a", "ToolA")
            .thinktool("b", "ToolB")
            .thinktool("c", "ToolC")
            .edge("a", "c")
            .edge("b", "c")
            .build();

        let order = graph.execution_order();
        assert_eq!(order.len(), 3);

        // 'c' should come after 'a' and 'b'
        let c_pos = order.iter().position(|n| n.id == "c").unwrap();
        let a_pos = order.iter().position(|n| n.id == "a").unwrap();
        let b_pos = order.iter().position(|n| n.id == "b").unwrap();

        assert!(c_pos > a_pos);
        assert!(c_pos > b_pos);
    }

    #[test]
    fn test_cycle_detection() {
        let mut graph = ExecutionGraph::new();
        graph
            .add_node(ExecutionNode {
                id: "a".to_string(),
                node_type: "test".to_string(),
                config: serde_json::json!({}),
                state: NodeState::Pending,
            })
            .unwrap();
        graph
            .add_node(ExecutionNode {
                id: "b".to_string(),
                node_type: "test".to_string(),
                config: serde_json::json!({}),
                state: NodeState::Pending,
            })
            .unwrap();

        // This should succeed
        graph
            .add_edge(
                "a",
                "b",
                DataFlow {
                    source_output: "out".to_string(),
                    target_input: "in".to_string(),
                },
            )
            .unwrap();

        // This would create a cycle - should fail
        let result = graph.add_edge(
            "b",
            "a",
            DataFlow {
                source_output: "out".to_string(),
                target_input: "in".to_string(),
            },
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_ready_nodes() {
        let mut graph = GraphBuilder::new()
            .thinktool("a", "ToolA")
            .thinktool("b", "ToolB")
            .thinktool("c", "ToolC")
            .edge("a", "c")
            .edge("b", "c")
            .build();

        // Initially, 'a' and 'b' should be ready
        let ready = graph.ready_nodes();
        assert_eq!(ready.len(), 2);

        // Mark 'a' as completed
        graph.set_state("a", NodeState::Completed).unwrap();

        // 'b' should still be ready, 'c' not yet
        let ready = graph.ready_nodes();
        assert_eq!(ready.len(), 1);
        assert_eq!(ready[0].id, "b");

        // Mark 'b' as completed
        graph.set_state("b", NodeState::Completed).unwrap();

        // Now 'c' should be ready
        let ready = graph.ready_nodes();
        assert_eq!(ready.len(), 1);
        assert_eq!(ready[0].id, "c");
    }
}