"""
Tests for the DAG module.
"""

import unittest
import asyncio
from tygent import DAG, ToolNode, LLMNode

class TestDAG(unittest.TestCase):
    """Test the DAG implementation."""
    
    def test_dag_creation(self):
        """Test that a DAG can be created."""
        dag = DAG("test_dag")
        self.assertEqual(dag.name, "test_dag")
        self.assertEqual(len(dag.nodes), 0)
        self.assertEqual(len(dag.edges), 0)
    
    def test_add_node(self):
        """Test that nodes can be added to a DAG."""
        dag = DAG("test_dag")
        
        async def test_tool(inputs):
            return {"result": "test"}
        
        tool_node = ToolNode("test_tool", test_tool)
        dag.add_node(tool_node)
        
        self.assertEqual(len(dag.nodes), 1)
        self.assertIn("test_tool", dag.nodes)
        self.assertEqual(dag.nodes["test_tool"], tool_node)
    
    def test_add_edge(self):
        """Test that edges can be added between nodes."""
        dag = DAG("test_dag")
        
        async def tool1(inputs):
            return {"data": "from_tool1"}
        
        async def tool2(inputs):
            return {"data": f"processed_{inputs.get('data', '')}"}
        
        dag.add_node(ToolNode("tool1", tool1))
        dag.add_node(ToolNode("tool2", tool2))
        
        dag.add_edge("tool1", "tool2", {"data": "data"})
        
        self.assertEqual(len(dag.edges), 1)
        self.assertIn("tool1", dag.edges)
        self.assertEqual(dag.edges["tool1"], ["tool2"])
        self.assertEqual(dag.edge_mappings["tool1"]["tool2"], {"data": "data"})
    
    def test_topological_order(self):
        """Test that topological ordering works properly."""
        dag = DAG("test_dag")
        
        async def dummy_tool(inputs):
            return {}
        
        # Create a simple diamond-shaped DAG
        #    A
        #   / \
        #  B   C
        #   \ /
        #    D
        
        for node_id in ["A", "B", "C", "D"]:
            dag.add_node(ToolNode(node_id, dummy_tool))
        
        dag.add_edge("A", "B")
        dag.add_edge("A", "C")
        dag.add_edge("B", "D")
        dag.add_edge("C", "D")
        
        order = dag.get_topological_order()
        
        # Verify that A comes before B and C, and B and C come before D
        self.assertTrue(order.index("A") < order.index("B"))
        self.assertTrue(order.index("A") < order.index("C"))
        self.assertTrue(order.index("B") < order.index("D"))
        self.assertTrue(order.index("C") < order.index("D"))

class TestExecution(unittest.TestCase):
    """Test the execution of DAGs."""
    
    def test_execution(self):
        """Test that a simple DAG can be executed successfully."""
        from tygent import Scheduler
        
        async def simple_add(inputs):
            a = inputs.get("a", 0)
            b = inputs.get("b", 0)
            return {"sum": a + b}
        
        async def simple_multiply(inputs):
            sum_value = inputs.get("sum", 0)
            factor = inputs.get("factor", 1)
            return {"product": sum_value * factor}
        
        async def run_test():
            dag = DAG("math_dag")
            
            dag.add_node(ToolNode("add", simple_add))
            dag.add_node(ToolNode("multiply", simple_multiply))
            
            dag.add_edge("add", "multiply", {"sum": "sum"})
            
            scheduler = Scheduler(dag)
            
            # Execute with inputs a=2, b=3, factor=5, expecting product=25
            result = await scheduler.execute({"a": 2, "b": 3, "factor": 5})
            
            self.assertIn("results", result)
            self.assertIn("multiply", result["results"])
            self.assertEqual(result["results"]["multiply"]["product"], 25)
        
        asyncio.run(run_test())

if __name__ == "__main__":
    unittest.main()