from unittest.mock import MagicMock

from GramAddict.core.situational_awareness import SituationalAwarenessEngine, SituationType


def test_sae_escalation_reset_on_situation_change():
    """
    Test that the SAE resets its escalation counter if the situation type changes.
    This prevents the 'Nuclear Escalation' trap when transitioning from
    a system permission dialog to an in-app modal.
    """
    device_mock = MagicMock()
    # Mocking a sequence of dumps:
    # 1-3: System Permission Dialog
    # 4-6: In-App Modal
    # 7: Clear

    # Needs 8 calls because the first attempt gets initial_xml if provided,
    # but subsequent attempts call dump_hierarchy(). In execute_escape we also call dump_hierarchy.
    # Actually, let's just make dump_hierarchy yield from a generator, but also
    # the SAE perceive is what we care about.

    # We will patch `perceive` to directly return our mock situations
    sae = SituationalAwarenessEngine.get_instance(device_mock)

    situations = [
        # Attempt 0
        SituationType.OBSTACLE_SYSTEM,  # perceive
        SituationType.OBSTACLE_SYSTEM,  # post-perceive -> success=False (counter=1)
        # Attempt 1
        SituationType.OBSTACLE_SYSTEM,  # perceive
        SituationType.OBSTACLE_SYSTEM,  # post-perceive -> success=False (counter=2)
        # Attempt 2
        SituationType.OBSTACLE_SYSTEM,  # perceive
        SituationType.OBSTACLE_MODAL,  # post-perceive -> success=False (counter=3)
        # Attempt 3 (new situation perceived!) -> situation_attempts resets to 0
        SituationType.OBSTACLE_MODAL,  # perceive
        SituationType.OBSTACLE_MODAL,  # post-perceive -> success=False (counter=1)
        # Attempt 4
        SituationType.OBSTACLE_MODAL,  # perceive
        SituationType.OBSTACLE_MODAL,  # post-perceive -> success=False (counter=2)
        # Attempt 5
        SituationType.OBSTACLE_MODAL,  # perceive
        SituationType.NORMAL,  # post-perceive -> success=True!
    ]

    sae.perceive = MagicMock(side_effect=situations)
    sae._plan_escape_via_llm = MagicMock()
    sae._execute_escape = MagicMock()

    from GramAddict.core.situational_awareness import EscapeAction

    # Let the LLM "plan" something so it doesn't crash
    # Each plan needs a unique reason to not be caught by failed_this_session perfectly if it's the same coordinate?
    # Actually, the recall check does: failed_this_session.add(action_key)
    # If the LLM keeps returning the same coordinates, it might be an issue.
    # We can just return different EscapeActions on side_effect
    sae._plan_escape_via_llm.side_effect = [
        EscapeAction("tap_coordinates", 100, 100, "mock_1"),
        EscapeAction("tap_coordinates", 101, 100, "mock_2"),
        EscapeAction("tap_coordinates", 102, 100, "mock_3"),
        EscapeAction("tap_coordinates", 103, 100, "mock_4"),
        EscapeAction("tap_coordinates", 104, 100, "mock_5"),
        EscapeAction("tap_coordinates", 105, 100, "mock_6"),
    ]

    # Since execute_escape checks the device dump, we just mock device.dump_hierarchy to return garbage
    # The actual situation check relies on perceive
    device_mock.dump_hierarchy.return_value = "<mock/>"

    success = sae.ensure_clear_screen(max_attempts=10)

    assert success is True
    assert sae.perceive.call_count == 12

    # 6 LLM calls total
    assert sae._plan_escape_via_llm.call_count == 6

    # We should ensure that app_start (nuclear escalation) was NEVER called.
    # We can check the actions executed
    app_starts = [
        args[0][0].action_type for args in sae._execute_escape.call_args_list if args[0][0].action_type == "app_start"
    ]
    assert len(app_starts) == 0
