import os
import sys
from unittest.mock import MagicMock, patch

# Add current directory to path for imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from main import run_flow_b

def test_run_flow_b_parsing():
    print("Testing run_flow_b parsing and extraction logic...")
    
    mock_response = MagicMock()
    mock_response.content = '{"s": "frustration", "r": "Let\'s simplify this. What do you get when you add the numbers first?"}'
    
    with patch("main.ChatGoogleGenerativeAI") as MockLLM:
        mock_llm_instance = MagicMock()
        mock_llm_instance.invoke.return_value = mock_response
        MockLLM.return_value = mock_llm_instance
        
        state, reply, context, est_in, est_out = run_flow_b(
            message="I am stuck on this step, it is too hard!",
            api_key="fake-api-key"
        )
        
        assert state == "frustration", f"Expected frustration, got {state}"
        assert "simplify" in reply, f"Expected reply to contain simplification, got {reply}"
        assert est_in > 0
        assert est_out > 0
        
        print("Mocked run_flow_b test passed!")

if __name__ == "__main__":
    test_run_flow_b_parsing()