import time

from loguru import logger
from swarms import Agent

from experimental.airflow_swarm import (
    AirflowDAGSwarm,
    NodeType,
    Conversation,
)

# Configure logger
logger.remove()
logger.add(lambda msg: print(msg, end=""), level="DEBUG")


def test_swarm_initialization():
    """Test basic swarm initialization and configuration."""
    try:
        swarm = AirflowDAGSwarm(
            dag_id="test_dag",
            name="Test DAG",
            initial_message="Test message",
        )
        assert swarm.dag_id == "test_dag", "DAG ID not set correctly"
        assert swarm.name == "Test DAG", "Name not set correctly"
        assert (
            len(swarm.nodes) == 0
        ), "Nodes should be empty on initialization"
        assert (
            len(swarm.edges) == 0
        ), "Edges should be empty on initialization"

        # Test initial message
        conv_json = swarm.get_conversation_history()
        assert (
            "Test message" in conv_json
        ), "Initial message not set correctly"
        print("✅ Swarm initialization test passed")
        return True
    except AssertionError as e:
        print(f"❌ Swarm initialization test failed: {str(e)}")
        return False


def test_node_addition():
    """Test adding different types of nodes to the swarm."""
    try:
        swarm = AirflowDAGSwarm(dag_id="test_dag")

        # Test adding an agent node
        agent = Agent(
            agent_name="Test-Agent",
            system_prompt="Test prompt",
            model_name="gpt-4o-mini",
            max_loops=1,
        )
        agent_id = swarm.add_node(
            "test_agent",
            agent,
            NodeType.AGENT,
            query="Test query",
            concurrent=True,
        )
        assert (
            agent_id == "test_agent"
        ), "Agent node ID not returned correctly"
        assert (
            "test_agent" in swarm.nodes
        ), "Agent node not added to nodes dict"

        # Test adding a callable node
        def test_callable(x: int, conversation: Conversation) -> str:
            return f"Test output {x}"

        callable_id = swarm.add_node(
            "test_callable",
            test_callable,
            NodeType.CALLABLE,
            args=[42],
            concurrent=False,
        )
        assert (
            callable_id == "test_callable"
        ), "Callable node ID not returned correctly"
        assert (
            "test_callable" in swarm.nodes
        ), "Callable node not added to nodes dict"

        print("✅ Node addition test passed")
        return True
    except AssertionError as e:
        print(f"❌ Node addition test failed: {str(e)}")
        return False
    except Exception as e:
        print(
            f"❌ Node addition test failed with unexpected error: {str(e)}"
        )
        return False


def test_edge_addition():
    """Test adding edges between nodes."""
    try:
        swarm = AirflowDAGSwarm(dag_id="test_dag")

        # Add two nodes
        def node1_fn(conversation: Conversation) -> str:
            return "Node 1 output"

        def node2_fn(conversation: Conversation) -> str:
            return "Node 2 output"

        swarm.add_node("node1", node1_fn, NodeType.CALLABLE)
        swarm.add_node("node2", node2_fn, NodeType.CALLABLE)

        # Add edge between them
        swarm.add_edge("node1", "node2")

        assert (
            "node2" in swarm.edges["node1"]
        ), "Edge not added correctly"
        assert (
            len(swarm.edges["node1"]) == 1
        ), "Incorrect number of edges"

        # Test adding edge with non-existent node
        try:
            swarm.add_edge("node1", "non_existent")
            assert (
                False
            ), "Should raise ValueError for non-existent node"
        except ValueError:
            pass

        print("✅ Edge addition test passed")
        return True
    except AssertionError as e:
        print(f"❌ Edge addition test failed: {str(e)}")
        return False


def test_execution_order():
    """Test that nodes are executed in the correct order based on dependencies."""
    try:
        swarm = AirflowDAGSwarm(dag_id="test_dag")
        execution_order = []

        def node1(conversation: Conversation) -> str:
            execution_order.append("node1")
            return "Node 1 output"

        def node2(conversation: Conversation) -> str:
            execution_order.append("node2")
            return "Node 2 output"

        def node3(conversation: Conversation) -> str:
            execution_order.append("node3")
            return "Node 3 output"

        # Add nodes
        swarm.add_node(
            "node1", node1, NodeType.CALLABLE, concurrent=False
        )
        swarm.add_node(
            "node2", node2, NodeType.CALLABLE, concurrent=False
        )
        swarm.add_node(
            "node3", node3, NodeType.CALLABLE, concurrent=False
        )

        # Add edges to create a chain: node1 -> node2 -> node3
        swarm.add_edge("node1", "node2")
        swarm.add_edge("node2", "node3")

        # Execute
        swarm.run()

        # Check execution order
        assert execution_order == [
            "node1",
            "node2",
            "node3",
        ], "Incorrect execution order"
        print("✅ Execution order test passed")
        return True
    except AssertionError as e:
        print(f"❌ Execution order test failed: {str(e)}")
        return False


def test_concurrent_execution():
    """Test concurrent execution of nodes."""
    try:
        swarm = AirflowDAGSwarm(dag_id="test_dag")

        def slow_node1(conversation: Conversation) -> str:
            time.sleep(0.5)
            return "Slow node 1 output"

        def slow_node2(conversation: Conversation) -> str:
            time.sleep(0.5)
            return "Slow node 2 output"

        # Add nodes with concurrent=True
        swarm.add_node(
            "slow1", slow_node1, NodeType.CALLABLE, concurrent=True
        )
        swarm.add_node(
            "slow2", slow_node2, NodeType.CALLABLE, concurrent=True
        )

        # Measure execution time
        start_time = time.time()
        swarm.run()
        execution_time = time.time() - start_time

        # Should take ~0.5s for concurrent execution, not ~1s
        assert (
            execution_time < 0.8
        ), "Concurrent execution took too long"
        print("✅ Concurrent execution test passed")
        return True
    except AssertionError as e:
        print(f"❌ Concurrent execution test failed: {str(e)}")
        return False


def test_conversation_handling():
    """Test conversation management within the swarm."""
    try:
        swarm = AirflowDAGSwarm(
            dag_id="test_dag", initial_message="Initial test message"
        )

        # Test adding user messages
        swarm.add_user_message("Test message 1")
        swarm.add_user_message("Test message 2")

        history = swarm.get_conversation_history()
        assert (
            "Initial test message" in history
        ), "Initial message not in history"
        assert (
            "Test message 1" in history
        ), "First message not in history"
        assert (
            "Test message 2" in history
        ), "Second message not in history"

        print("✅ Conversation handling test passed")
        return True
    except AssertionError as e:
        print(f"❌ Conversation handling test failed: {str(e)}")
        return False


def test_error_handling():
    """Test error handling in node execution."""
    try:
        swarm = AirflowDAGSwarm(dag_id="test_dag")

        def failing_node(conversation: Conversation) -> str:
            raise ValueError("Test error")

        swarm.add_node("failing", failing_node, NodeType.CALLABLE)

        # Execute should not raise an exception
        result = swarm.run()

        assert (
            "Error" in result
        ), "Error not captured in execution result"
        assert (
            "Test error" in result
        ), "Specific error message not captured"

        print("✅ Error handling test passed")
        return True
    except Exception as e:
        print(f"❌ Error handling test failed: {str(e)}")
        return False


def run_all_tests():
    """Run all test functions and report results."""
    tests = [
        test_swarm_initialization,
        test_node_addition,
        test_edge_addition,
        test_execution_order,
        test_concurrent_execution,
        test_conversation_handling,
        test_error_handling,
    ]

    results = []
    for test in tests:
        print(f"\nRunning {test.__name__}...")
        result = test()
        results.append(result)

    total = len(results)
    passed = sum(results)
    print("\n=== Test Results ===")
    print(f"Total tests: {total}")
    print(f"Passed: {passed}")
    print(f"Failed: {total - passed}")
    print("==================")


if __name__ == "__main__":
    run_all_tests()
