"""
Basic tests for YBAgent workflow engine.
"""

import pytest
from ybagent import Workflow, InputNode, OutputNode, TransformNode
from ybagent.nodes.control_nodes import ConditionalNode, MergeNode


def test_workflow_creation():
    """Test basic workflow creation."""
    workflow = Workflow(name="Test Workflow")
    assert workflow.name == "Test Workflow"
    assert len(workflow.nodes) == 0
    assert len(workflow.edges) == 0


def test_add_nodes():
    """Test adding nodes to workflow."""
    workflow = Workflow(name="Test")
    
    input_node = InputNode(id="input1")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(output_node)
    
    assert len(workflow.nodes) == 2
    assert "input1" in workflow.nodes
    assert "output1" in workflow.nodes


def test_connect_nodes():
    """Test connecting nodes."""
    workflow = Workflow(name="Test")
    
    input_node = InputNode(id="input1")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(output_node)
    workflow.connect("input1", "output1")
    
    assert len(workflow.edges) == 1
    assert workflow.edges[0].source_node_id == "input1"
    assert workflow.edges[0].target_node_id == "output1"


def test_simple_execution():
    """Test simple workflow execution."""
    workflow = Workflow(name="Simple Test")
    
    input_node = InputNode(id="input1")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(output_node)
    workflow.connect("input1", "output1")
    
    result = workflow.execute({"test": "data"})
    assert result is not None


def test_transform_node():
    """Test transform node."""
    workflow = Workflow(name="Transform Test")
    
    input_node = InputNode(id="input1")
    transform_node = TransformNode(id="transform1", transform_type="string")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(transform_node)
    workflow.add_node(output_node)
    
    workflow.connect("input1", "transform1")
    workflow.connect("transform1", "output1")
    
    result = workflow.execute({"input": 123})
    assert result is not None


def test_conditional_node():
    """Test conditional node."""
    workflow = Workflow(name="Conditional Test")
    
    input_node = InputNode(id="input1")
    conditional = ConditionalNode(id="cond1", condition="value > 5")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(conditional)
    workflow.add_node(output_node)
    
    workflow.connect("input1", "cond1")
    workflow.connect("cond1", "output1", source_output="true")
    
    result = workflow.execute({"value": 10})
    assert result is not None


def test_merge_node():
    """Test merge node."""
    workflow = Workflow(name="Merge Test")
    
    input1 = InputNode(id="input1")
    input2 = InputNode(id="input2")
    merge = MergeNode(id="merge1", merge_strategy="list")
    output = OutputNode(id="output1")
    
    workflow.add_node(input1)
    workflow.add_node(input2)
    workflow.add_node(merge)
    workflow.add_node(output)
    
    workflow.connect("input1", "merge1", target_input="input1")
    workflow.connect("input2", "merge1", target_input="input2")
    workflow.connect("merge1", "output1")
    
    result = workflow.execute({"data1": "test1", "data2": "test2"})
    assert result is not None


def test_workflow_serialization():
    """Test workflow serialization to JSON."""
    workflow = Workflow(name="Serialization Test")
    
    input_node = InputNode(id="input1")
    output_node = OutputNode(id="output1")
    
    workflow.add_node(input_node)
    workflow.add_node(output_node)
    workflow.connect("input1", "output1")
    
    json_str = workflow.to_json()
    assert json_str is not None
    assert "Serialization Test" in json_str
    assert "input1" in json_str
    assert "output1" in json_str


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
