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_orchestrator_full_flow():
    """Verify the full Orchestrator-Workers flow via direct node injection."""
    
    mock_analyzer = MagicMock()
    mock_planner = MagicMock()
    mock_delegate = MagicMock()
    mock_worker = MagicMock()
    mock_reflector = MagicMock()
    mock_synthesizer = MagicMock()
    mock_summarize_conv = MagicMock()
    
    # 1. Analyzer: Proceed to planning
    mock_analyzer.return_value = {"next_action": "plan"}
    
    # 2. Planner: Generate checklist
    mock_planner.return_value = {
        "checklist": [{"task": "T1", "worker": "data_analyst"}],
        "current_step": 0
    }
    
    # 3. Delegate: Route to data_analyst
    mock_delegate.side_effect = [
        {"next_action": "data_analyst"}, # First call
        {"next_action": "summarize"}    # Second call (after reflector)
    ]
    
    # 4. Worker: Success
    mock_worker.return_value = {
        "messages": [AIMessage(content="Worker result")],
        "vfs": {"res.txt": "data"}
    }
    
    # 5. Reflector: Advance
    mock_reflector.return_value = {
        "current_step": 1,
        "next_action": "delegate"
    }
    
    # 6. Synthesizer: Final answer
    mock_synthesizer.return_value = {
        "messages": [AIMessage(content="Final synthesized answer")],
        "next_action": "end"
    }
    
    # 7. Summarize Conv: End
    mock_summarize_conv.return_value = {"summary": "Done"}
    
    # Create workflow with injected mocks
    app = create_workflow(
        query_analyzer=mock_analyzer,
        planner=mock_planner,
        delegate=mock_delegate,
        data_analyst_worker=mock_worker,
        reflector=mock_reflector,
        synthesizer=mock_synthesizer,
        summarize_conversation=mock_summarize_conv
    )
    
    initial_state = AgentState(
        messages=[HumanMessage(content="Explain results")],
        question="Explain results",
        analysis={},
        next_action="",
        iterations=0,
        checklist=[],
        current_step=0,
        vfs={},
        plots=[],
        dfs={}
    )
    
    final_state = app.invoke(initial_state)
    
    assert mock_analyzer.called
    assert mock_planner.called
    assert mock_delegate.call_count == 2
    assert mock_worker.called
    assert mock_reflector.called
    assert mock_synthesizer.called
    assert "Final synthesized answer" in [m.content for m in final_state["messages"]]