import pytest
from unittest.mock import MagicMock, patch
from ea_chatbot.graph.nodes.planner import planner_node

@pytest.fixture
def mock_state():
    return {
        "messages": [],
        "question": "Show me results for New Jersey",
        "analysis": {
            # "requires_dataset" removed as it's no longer used
            "expert": "Data Analyst",
            "data": "NJ data",
            "unknown": "results",
            "condition": "state=NJ"
        },
        "next_action": "plan",
        "plan": None
    }

@patch("ea_chatbot.graph.nodes.planner.get_llm_model")
@patch("ea_chatbot.utils.database_inspection.get_data_summary")
def test_planner_node(mock_get_summary, mock_get_llm, mock_state):
    """Test planner node with unified prompt."""
    mock_get_summary.return_value = "Column: Name, Type: text"
    
    mock_llm = MagicMock()
    mock_get_llm.return_value = mock_llm
    
    from ea_chatbot.schemas import TaskPlanResponse, TaskPlanContext
    mock_plan = TaskPlanResponse(
        goal="Get NJ results",
        reflection="The user wants NJ results",
        context=TaskPlanContext(initial_context="NJ data", assumptions=[], constraints=[]),
        steps=["Step 1: Load data", "Step 2: Filter by NJ"]
    )
    mock_llm.with_structured_output.return_value.invoke.return_value = mock_plan
    
    result = planner_node(mock_state)
    
    assert "plan" in result
    assert "Step 1: Load data" in result["plan"]
    assert "Step 2: Filter by NJ" in result["plan"]
    
    # Verify helper was called
    mock_get_summary.assert_called_once()