import pytest
from unittest.mock import MagicMock
from ea_chatbot.graph.workflow import create_workflow
from ea_chatbot.graph.state import AgentState
from langchain_core.messages import AIMessage, HumanMessage

def test_deepagents_multi_worker_sequential_flow():
    """Verify that the Orchestrator can handle a sequence of different workers."""
    
    mock_analyzer = MagicMock()
    mock_planner = MagicMock()
    mock_delegate = MagicMock()
    mock_analyst = MagicMock()
    mock_researcher = MagicMock()
    mock_reflector = MagicMock()
    mock_synthesizer = MagicMock()
    
    # 1. Analyzer: Plan
    mock_analyzer.return_value = {"next_action": "plan"}
    
    # 2. Planner: Two tasks
    mock_planner.return_value = {
        "checklist": [
            {"task": "Get Numbers", "worker": "data_analyst"},
            {"task": "Get Facts", "worker": "researcher"}
        ],
        "current_step": 0
    }
    
    # 3. Delegate & Reflector: Loop through tasks
    mock_delegate.side_effect = [
        {"next_action": "data_analyst"}, # Step 0
        {"next_action": "researcher"},   # Step 1
        {"next_action": "summarize"}     # Step 2 (done)
    ]
    
    mock_analyst.return_value = {"messages": [AIMessage(content="Analyst summary")], "vfs": {"data.csv": "..."}}
    mock_researcher.return_value = {"messages": [AIMessage(content="Researcher summary")]}
    
    mock_reflector.side_effect = [
        {"current_step": 1, "next_action": "delegate"}, # Done with Analyst
        {"current_step": 2, "next_action": "delegate"}  # Done with Researcher
    ]
    
    mock_synthesizer.return_value = {
        "messages": [AIMessage(content="Final multi-agent response")],
        "next_action": "end"
    }
    
    app = create_workflow(
        query_analyzer=mock_analyzer,
        planner=mock_planner,
        delegate=mock_delegate,
        data_analyst_worker=mock_analyst,
        researcher_worker=mock_researcher,
        reflector=mock_reflector,
        synthesizer=mock_synthesizer
    )
    
    initial_state = AgentState(
        messages=[HumanMessage(content="Numbers and facts please")],
        question="Numbers and facts please",
        analysis={},
        next_action="",
        iterations=0,
        checklist=[],
        current_step=0,
        vfs={},
        plots=[],
        dfs={}
    )
    
    final_state = app.invoke(initial_state, config={"recursion_limit": 30})
    
    assert mock_analyst.called
    assert mock_researcher.called
    assert mock_reflector.call_count == 2
    assert "Final multi-agent response" in [m.content for m in final_state["messages"]]
    assert final_state["current_step"] == 2