from unittest.mock import MagicMock, patch

import pytest

from GramAddict.core.situational_awareness import EscapeAction, SituationalAwarenessEngine, SituationType


@pytest.fixture
def mock_device():
    device = MagicMock()
    device.app_id = "com.instagram.android"
    return device


def test_sae_state_transition_success(mock_device):
    """
    Test that if an action changes the situation from one obstacle to ANOTHER obstacle,
    it is considered a partial success and NOT marked as a failure for the previous situation.
    Also verifies that LLM queries use a sufficient max_tokens limit to prevent truncation.
    """
    sae = SituationalAwarenessEngine(mock_device)

    # We will simulate 3 dumps:
    # 1. FOREIGN_APP
    # 2. OBSTACLE_MODAL (Foreign app killed, but now we have a modal)
    # 3. NORMAL (Modal dismissed)

    # We don't actually need real XML if we mock perceive and _compress_xml
    mock_device.dump_hierarchy.side_effect = ["<xml>1</xml>", "<xml>2</xml>", "<xml>3</xml>"]

    # Mock compression to avoid real work
    sae._compress_xml = MagicMock(side_effect=["comp1", "comp2", "comp3"])

    # Mock perception
    sae.perceive = MagicMock(
        side_effect=[
            SituationType.OBSTACLE_FOREIGN_APP,  # Initial
            SituationType.OBSTACLE_MODAL,  # After attempt 1
            SituationType.OBSTACLE_MODAL,  # Start of attempt 2
            SituationType.NORMAL,  # After attempt 2
        ]
    )

    # Mock LLM fallback planning
    llm_actions = [EscapeAction(action_type="click", x=100, y=100, reason="LLM Action to dismiss modal")]
    sae._plan_escape_via_llm = MagicMock(side_effect=llm_actions)

    # Mock memory to return nothing (force LLM/heuristic)
    sae.episodes.recall = MagicMock(return_value=None)
    sae.episodes.learn = MagicMock()

    # Mock execution
    sae._kill_foreign_apps = MagicMock()
    sae._execute_escape = MagicMock()

    # Let's use the REAL _plan_escape_via_llm but mock `query_llm`
    sae._plan_escape_via_llm = SituationalAwarenessEngine._plan_escape_via_llm.__get__(sae, SituationalAwarenessEngine)

    with patch("GramAddict.core.llm_provider.query_llm") as mock_query_llm:
        mock_query_llm.return_value = {"response": '{"action": "click", "x": 100, "y": 100, "reason": "test"}'}

        result = sae.ensure_clear_screen(max_attempts=5, initial_xml="<xml>0</xml>")

        assert result is True, "SAE should eventually clear the screen"

        # Check that query_llm was called with max_tokens >= 300
        assert mock_query_llm.called
        kwargs = mock_query_llm.call_args[1]
        assert kwargs.get("max_tokens", 0) >= 300, f"max_tokens is too low: {kwargs.get('max_tokens')}"

        # Check that the first action (killing foreign apps) was NOT marked as a failure,
        # because it successfully transitioned from FOREIGN_APP to OBSTACLE_MODAL.
        # Wait, the failure is tracked in `failed_this_session`. We can't easily inspect it directly
        # since it's a local variable. But we can check `sae.episodes.learn` calls!
        # The first learn call should be success=True because the state changed!

        learn_calls = sae.episodes.learn.call_args_list
        assert len(learn_calls) >= 2

        # First action (kill_foreign_apps)
        assert learn_calls[0][0][2] is True, "kill_foreign_apps should be marked as success because situation changed"

        # Second action (click from LLM)
        assert learn_calls[1][0][2] is True, "click should be marked as success because we reached NORMAL"
