import asyncio
import concurrent
import threading
import unittest.mock
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch

import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

import strands
import strands.telemetry
from strands import Agent
from strands.event_loop._retry import ModelRetryStrategy
from strands.experimental.checkpoint import Checkpoint
from strands.hooks import (
    AfterModelCallEvent,
    BeforeModelCallEvent,
    BeforeToolCallEvent,
    HookRegistry,
    MessageAddedEvent,
)
from strands.interrupt import Interrupt, _InterruptState
from strands.telemetry.metrics import EventLoopMetrics
from strands.telemetry.tracer import Tracer
from strands.tools.executors import SequentialToolExecutor
from strands.tools.registry import ToolRegistry
from strands.types._events import EventLoopStopEvent
from strands.types.exceptions import (
    ContextWindowOverflowException,
    EventLoopException,
    MaxTokensReachedException,
    ModelThrottledException,
)
from tests.fixtures.mock_hook_provider import MockHookProvider
from tests.fixtures.mocked_model_provider import MockedModelProvider


@pytest.fixture
def mock_sleep():
    with patch.object(strands.event_loop._retry.asyncio, "sleep", new_callable=AsyncMock) as mock:
        yield mock


@pytest.fixture
def model():
    return unittest.mock.Mock()


@pytest.fixture
def system_prompt():
    return "p1"


@pytest.fixture
def messages():
    return [{"role": "user", "content": [{"text": "Hello"}]}]


@pytest.fixture
def tool_registry():
    return ToolRegistry()


@pytest.fixture
def thread_pool():
    return concurrent.futures.ThreadPoolExecutor(max_workers=1)


@pytest.fixture
def tool(tool_registry):
    @strands.tool
    def tool_for_testing(random_string: str):
        return random_string

    tool_registry.register_tool(tool_for_testing)

    return tool_for_testing


@pytest.fixture
def tool_times_2(tool_registry):
    @strands.tools.tool
    def multiply_by_2(x: int) -> int:
        return x * 2

    tool_registry.register_tool(multiply_by_2)

    return multiply_by_2


@pytest.fixture
def tool_times_5(tool_registry):
    @strands.tools.tool
    def multiply_by_5(x: int) -> int:
        return x * 5

    tool_registry.register_tool(multiply_by_5)

    return multiply_by_5


@pytest.fixture
def tool_stream(tool):
    return [
        {
            "contentBlockStart": {
                "start": {
                    "toolUse": {
                        "toolUseId": "t1",
                        "name": tool.tool_spec["name"],
                    },
                },
            },
        },
        {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"random_string": "abcdEfghI123"}'}}}},
        {"contentBlockStop": {}},
        {"messageStop": {"stopReason": "tool_use"}},
    ]


@pytest.fixture
def hook_registry():
    registry = HookRegistry()
    # Register default retry strategy
    retry_strategy = ModelRetryStrategy()
    retry_strategy.register_hooks(registry)
    return registry


@pytest.fixture
def hook_provider(hook_registry):
    provider = MockHookProvider(event_types="all")
    hook_registry.add_hook(provider)
    return provider


@pytest.fixture
def tool_executor():
    return SequentialToolExecutor()


@pytest.fixture
def agent(model, system_prompt, messages, tool_registry, thread_pool, hook_registry, tool_executor):
    mock = unittest.mock.Mock(name="agent")
    mock.__class__ = Agent
    mock.config.cache_points = []
    mock.model = model
    mock.system_prompt = system_prompt
    mock.messages = messages
    mock.tool_registry = tool_registry
    mock.thread_pool = thread_pool
    mock.event_loop_metrics = EventLoopMetrics()
    mock.event_loop_metrics.reset_usage_metrics()
    mock.hooks = hook_registry
    mock.tool_executor = tool_executor
    mock._interrupt_state = _InterruptState()
    mock._cancel_signal = threading.Event()
    mock._model_state = {}
    mock._checkpointing = False
    mock._checkpoint = None
    mock._checkpoint_cycle_index = 0
    mock._checkpoint_resume_position = None
    mock.trace_attributes = {}
    mock.retry_strategy = ModelRetryStrategy()

    return mock


@pytest.fixture
def mock_tracer():
    tracer = MagicMock()
    tracer.start_event_loop_cycle_span.return_value = MagicMock()
    tracer.start_model_invoke_span.return_value = MagicMock()
    return tracer


@pytest.mark.asyncio
async def test_event_loop_cycle_text_response(
    agent,
    model,
    agenerator,
    alist,
):
    model.stream.return_value = agenerator(
        [
            {"contentBlockDelta": {"delta": {"text": "test text"}}},
            {"contentBlockStop": {}},
        ]
    )

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, tru_message, _, tru_request_state, _, _, _ = events[-1]["stop"]

    exp_stop_reason = "end_turn"
    exp_message = {"role": "assistant", "content": [{"text": "test text"}], "metadata": ANY}
    exp_request_state = {}

    assert tru_stop_reason == exp_stop_reason and tru_message == exp_message and tru_request_state == exp_request_state


@pytest.mark.asyncio
async def test_event_loop_cycle_text_response_throttling(
    mock_sleep,
    agent,
    model,
    agenerator,
    alist,
):
    model.stream.side_effect = [
        ModelThrottledException("ThrottlingException | ConverseStream"),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, tru_message, _, tru_request_state, _, _, _ = events[-1]["stop"]

    exp_stop_reason = "end_turn"
    exp_message = {"role": "assistant", "content": [{"text": "test text"}], "metadata": ANY}
    exp_request_state = {}

    assert tru_stop_reason == exp_stop_reason and tru_message == exp_message and tru_request_state == exp_request_state
    # Verify that sleep was called once with the initial delay
    mock_sleep.assert_called_once()


@pytest.mark.asyncio
async def test_event_loop_cycle_exponential_backoff(
    mock_sleep,
    agent,
    model,
    agenerator,
    alist,
):
    """Test that the exponential backoff works correctly with multiple retries."""
    # Set up the model to raise throttling exceptions multiple times before succeeding
    model.stream.side_effect = [
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, tru_message, _, tru_request_state, _, _, _ = events[-1]["stop"]

    # Verify the final response
    assert tru_stop_reason == "end_turn"
    assert tru_message == {"role": "assistant", "content": [{"text": "test text"}], "metadata": ANY}
    assert tru_request_state == {}

    # Verify that sleep was called with increasing delays
    # Initial delay is 4, then 8, then 16
    assert mock_sleep.call_count == 3
    assert mock_sleep.call_args_list == [call(4), call(8), call(16)]


@pytest.mark.asyncio
async def test_event_loop_cycle_text_response_throttling_exceeded(
    mock_sleep,
    agent,
    model,
    alist,
):
    model.stream.side_effect = [
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
        ModelThrottledException("ThrottlingException | ConverseStream"),
    ]

    with pytest.raises(ModelThrottledException):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

    mock_sleep.assert_has_calls(
        [
            call(4),
            call(8),
            call(16),
            call(32),
            call(64),
        ]
    )


@pytest.mark.asyncio
async def test_event_loop_cycle_text_response_error(
    agent,
    model,
    alist,
):
    model.stream.side_effect = RuntimeError("Unhandled error")

    with pytest.raises(RuntimeError):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)


@patch("strands.event_loop.event_loop.recover_message_on_max_tokens_reached")
@pytest.mark.asyncio
async def test_event_loop_cycle_tool_result(
    mock_recover_message,
    agent,
    model,
    system_prompt,
    messages,
    tool_stream,
    tool_registry,
    agenerator,
    alist,
):
    model.stream.side_effect = [
        agenerator(tool_stream),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, tru_message, _, tru_request_state, _, _, _ = events[-1]["stop"]

    exp_stop_reason = "end_turn"
    exp_message = {"role": "assistant", "content": [{"text": "test text"}], "metadata": ANY}
    exp_request_state = {}

    assert tru_stop_reason == exp_stop_reason and tru_message == exp_message and tru_request_state == exp_request_state

    # Verify that recover_message_on_max_tokens_reached was NOT called for tool_use stop reason
    mock_recover_message.assert_not_called()

    model.stream.assert_called_with(
        [
            {"role": "user", "content": [{"text": "Hello"}]},
            {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": "t1",
                            "name": "tool_for_testing",
                            "input": {"random_string": "abcdEfghI123"},
                        }
                    }
                ],
            },
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": "t1",
                            "status": "success",
                            "content": [{"text": "abcdEfghI123"}],
                        },
                    },
                ],
            },
        ],
        tool_registry.get_all_tool_specs(),
        "p1",
        tool_choice=None,
        system_prompt_content=unittest.mock.ANY,
        invocation_state=unittest.mock.ANY,
        model_state=unittest.mock.ANY,
    )


@pytest.mark.asyncio
async def test_event_loop_cycle_tool_result_error(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    model.stream.side_effect = [agenerator(tool_stream)]

    with pytest.raises(EventLoopException):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)


@pytest.mark.asyncio
async def test_event_loop_cycle_tool_result_no_tool_handler(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    model.stream.side_effect = [agenerator(tool_stream)]
    # Set tool_handler to None for this test
    agent.tool_handler = None

    with pytest.raises(EventLoopException):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)


@pytest.mark.asyncio
async def test_event_loop_cycle_stop(
    agent,
    model,
    tool,
    agenerator,
    alist,
):
    model.stream.side_effect = [
        agenerator(
            [
                {
                    "contentBlockStart": {
                        "start": {
                            "toolUse": {
                                "toolUseId": "t1",
                                "name": tool.tool_spec["name"],
                            },
                        },
                    },
                },
                {"contentBlockStop": {}},
                {"messageStop": {"stopReason": "tool_use"}},
            ]
        ),
    ]

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={"request_state": {"stop_event_loop": True}},
    )
    events = await alist(stream)
    tru_stop_reason, tru_message, _, tru_request_state, _, _, _ = events[-1]["stop"]

    exp_stop_reason = "tool_use"
    exp_message = {
        "role": "assistant",
        "content": [
            {
                "toolUse": {
                    "input": {},
                    "name": "tool_for_testing",
                    "toolUseId": "t1",
                }
            }
        ],
        "metadata": ANY,
    }
    exp_request_state = {"stop_event_loop": True}

    assert tru_stop_reason == exp_stop_reason and tru_message == exp_message and tru_request_state == exp_request_state


@pytest.mark.asyncio
async def test_cycle_exception(
    agent,
    model,
    tool_stream,
    agenerator,
):
    model.stream.side_effect = [
        agenerator(tool_stream),
        agenerator(tool_stream),
        agenerator(tool_stream),
        ValueError("Invalid error presented"),
    ]

    tru_stop_event = None
    exp_stop_event = {"force_stop": True, "force_stop_reason": "Invalid error presented"}

    with pytest.raises(EventLoopException):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        async for event in stream:
            tru_stop_event = event

    assert tru_stop_event == exp_stop_event


@patch("strands.event_loop.event_loop.get_tracer")
@pytest.mark.asyncio
async def test_event_loop_cycle_creates_spans(
    mock_get_tracer,
    agent,
    model,
    mock_tracer,
    agenerator,
    alist,
):
    # Setup
    mock_get_tracer.return_value = mock_tracer
    cycle_span = MagicMock()
    mock_tracer.start_event_loop_cycle_span.return_value = cycle_span
    model_span = MagicMock()
    mock_tracer.start_model_invoke_span.return_value = model_span

    model.stream.return_value = agenerator(
        [
            {"contentBlockDelta": {"delta": {"text": "test text"}}},
            {"contentBlockStop": {}},
        ]
    )

    # Call event_loop_cycle
    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    await alist(stream)

    # Verify tracer methods were called correctly
    mock_get_tracer.assert_called_once()
    mock_tracer.start_event_loop_cycle_span.assert_called_once()
    mock_tracer.start_model_invoke_span.assert_called_once()
    call_kwargs = mock_tracer.start_model_invoke_span.call_args[1]
    assert call_kwargs["system_prompt"] == agent.system_prompt
    assert call_kwargs["system_prompt_content"] == agent._system_prompt_content
    mock_tracer.end_model_invoke_span.assert_called_once()
    mock_tracer.end_event_loop_cycle_span.assert_called_once()


@patch("strands.event_loop.event_loop.get_tracer")
@pytest.mark.asyncio
async def test_event_loop_tracing_with_model_error(
    mock_get_tracer,
    agent,
    model,
    mock_tracer,
    alist,
):
    # Setup
    mock_get_tracer.return_value = mock_tracer
    cycle_span = MagicMock()
    mock_tracer.start_event_loop_cycle_span.return_value = cycle_span
    model_span = MagicMock()
    mock_tracer.start_model_invoke_span.return_value = model_span

    # Set up model to raise an exception
    model.stream.side_effect = ContextWindowOverflowException("Input too long")

    # Call event_loop_cycle, expecting it to handle the exception
    with pytest.raises(ContextWindowOverflowException):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

    assert mock_tracer.end_span_with_error.call_count == 2
    mock_tracer.end_span_with_error.assert_has_calls(
        [
            call(model_span, "Input too long", model.stream.side_effect),
            call(cycle_span, "Input too long", model.stream.side_effect),
        ]
    )


@pytest.mark.asyncio
async def test_event_loop_cycle_max_tokens_exception(
    agent,
    model,
    agenerator,
    alist,
):
    """Test that max_tokens stop reason calls _recover_message_on_max_tokens_reached then MaxTokensReachedException."""

    model.stream.side_effect = [
        agenerator(
            [
                {
                    "contentBlockStart": {
                        "start": {
                            "toolUse": {
                                "toolUseId": "t1",
                                "name": "asdf",
                                "input": {},  # empty
                            },
                        },
                    },
                },
                {"contentBlockStop": {}},
                {"messageStop": {"stopReason": "max_tokens"}},
            ]
        ),
    ]

    # Call event_loop_cycle, expecting it to raise MaxTokensReachedException
    expected_message = (
        "Agent has reached an unrecoverable state due to max_tokens limit. "
        "For more information see: "
        "https://strandsagents.com/latest/user-guide/concepts/agents/agent-loop/#maxtokensreachedexception"
    )
    with pytest.raises(MaxTokensReachedException, match=expected_message):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

    # Verify the exception message contains the expected content
    assert len(agent.messages) == 2
    assert "tool use was incomplete due" in agent.messages[1]["content"][0]["text"]


@patch("strands.event_loop.event_loop.get_tracer")
@pytest.mark.asyncio
async def test_event_loop_tracing_with_tool_execution(
    mock_get_tracer,
    agent,
    model,
    tool_stream,
    mock_tracer,
    agenerator,
    alist,
):
    # Setup
    mock_get_tracer.return_value = mock_tracer
    cycle_span = MagicMock()
    mock_tracer.start_event_loop_cycle_span.return_value = cycle_span
    model_span = MagicMock()
    mock_tracer.start_model_invoke_span.return_value = model_span

    # Set up model to return tool use and then text response
    model.stream.side_effect = [
        agenerator(tool_stream),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    # Call event_loop_cycle which should execute a tool
    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    await alist(stream)

    # Verify the parent_span parameter is passed to run_tools
    # At a minimum, verify both model spans were created (one for each model invocation)
    assert mock_tracer.start_model_invoke_span.call_count == 2
    assert mock_tracer.end_model_invoke_span.call_count == 2


@pytest.mark.asyncio
async def test_event_loop_cycle_closes_cycle_span_before_recursive_cycle(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    exporter = InMemorySpanExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))

    tracer = Tracer()
    tracer.tracer_provider = provider
    tracer.tracer = provider.get_tracer(tracer.service_name)

    async def delayed_text_stream():
        yield {"contentBlockDelta": {"delta": {"text": "test text"}}}
        await asyncio.sleep(0.05)
        yield {"contentBlockStop": {}}

    agent.trace_span = None
    agent._system_prompt_content = None
    model.config = {"model_id": "test-model"}
    model.stream.side_effect = [
        agenerator(tool_stream),
        delayed_text_stream(),
    ]

    with patch("strands.event_loop.event_loop.get_tracer", return_value=tracer):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

    provider.force_flush()
    cycle_spans = sorted(
        [span for span in exporter.get_finished_spans() if span.name == "execute_event_loop_cycle"],
        key=lambda span: span.start_time,
    )

    assert len(cycle_spans) == 2
    assert cycle_spans[0].end_time <= cycle_spans[1].start_time
    assert cycle_spans[0].end_time < cycle_spans[1].end_time


@patch("strands.event_loop.event_loop.get_tracer")
@pytest.mark.asyncio
async def test_event_loop_tracing_with_throttling_exception(
    mock_get_tracer,
    agent,
    model,
    mock_tracer,
    agenerator,
    alist,
):
    # Setup
    mock_get_tracer.return_value = mock_tracer
    cycle_span = MagicMock()
    mock_tracer.start_event_loop_cycle_span.return_value = cycle_span
    model_span = MagicMock()
    mock_tracer.start_model_invoke_span.return_value = model_span

    # Set up model to raise a throttling exception and then succeed
    model.stream.side_effect = [
        ModelThrottledException("Throttling Error"),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    # Mock the time.sleep function to speed up the test
    with patch.object(asyncio, "sleep", new_callable=unittest.mock.AsyncMock):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

    assert mock_tracer.end_span_with_error.call_count == 1
    # Verify span was created for the successful retry
    assert mock_tracer.start_model_invoke_span.call_count == 2
    assert mock_tracer.end_model_invoke_span.call_count == 1


@patch("strands.event_loop.event_loop.get_tracer")
@pytest.mark.asyncio
async def test_event_loop_cycle_with_parent_span(
    mock_get_tracer,
    agent,
    model,
    messages,
    mock_tracer,
    agenerator,
    alist,
):
    # Setup
    mock_get_tracer.return_value = mock_tracer
    parent_span = MagicMock()
    cycle_span = MagicMock()
    mock_tracer.start_event_loop_cycle_span.return_value = cycle_span

    model.stream.return_value = agenerator(
        [
            {"contentBlockDelta": {"delta": {"text": "test text"}}},
            {"contentBlockStop": {}},
        ]
    )

    # Set the parent span for this test
    agent.trace_span = parent_span

    # Call event_loop_cycle with a parent span
    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    await alist(stream)

    # Verify parent_span was used when creating cycle span
    mock_tracer.start_event_loop_cycle_span.assert_called_once_with(
        invocation_state=unittest.mock.ANY,
        parent_span=parent_span,
        messages=messages,
        custom_trace_attributes=unittest.mock.ANY,
    )


@pytest.mark.asyncio
async def test_request_state_initialization(alist):
    # Create a mock agent
    mock_agent = MagicMock()
    # not setting this to False results in endless recursion
    mock_agent._interrupt_state.activated = False
    mock_agent._cancel_signal = threading.Event()
    mock_agent.event_loop_metrics.start_cycle.return_value = (0, MagicMock())
    mock_agent.hooks.invoke_callbacks_async = AsyncMock()

    # Call without providing request_state
    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=mock_agent,
        invocation_state={},
    )
    events = await alist(stream)
    _, _, _, tru_request_state, _, _, _ = events[-1]["stop"]

    # Verify request_state was initialized to empty dict
    assert tru_request_state == {}

    # Call with pre-existing request_state
    initial_request_state = {"key": "value"}
    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=mock_agent,
        invocation_state={"request_state": initial_request_state},
    )
    events = await alist(stream)
    _, _, _, tru_request_state, _, _, _ = events[-1]["stop"]

    # Verify existing request_state was preserved
    assert tru_request_state == initial_request_state


@pytest.mark.asyncio
async def test_prepare_next_cycle_in_tool_execution(agent, model, tool_stream, agenerator, alist):
    """Test that cycle ID and metrics are properly updated during tool execution."""
    model.stream.side_effect = [
        agenerator(tool_stream),
        agenerator(
            [
                {"contentBlockStop": {}},
            ]
        ),
    ]

    # Create a mock for recurse_event_loop to capture the invocation_state passed to it
    with unittest.mock.patch.object(strands.event_loop.event_loop, "recurse_event_loop") as mock_recurse:
        # Set up mock to return a valid response
        mock_recurse.return_value = agenerator(
            [
                (
                    "end_turn",
                    {"role": "assistant", "content": [{"text": "test text"}]},
                    strands.telemetry.metrics.EventLoopMetrics(),
                    {},
                ),
            ]
        )

        # Call event_loop_cycle which should execute a tool and then call recurse_event_loop
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        await alist(stream)

        assert mock_recurse.called

        # Verify required properties are present
        recursive_args = mock_recurse.call_args[1]
        assert "event_loop_parent_cycle_id" in recursive_args["invocation_state"]
        assert (
            recursive_args["invocation_state"]["event_loop_parent_cycle_id"]
            == recursive_args["invocation_state"]["event_loop_cycle_id"]
        )


@pytest.mark.asyncio
async def test_event_loop_cycle_exception_model_hooks(mock_sleep, agent, model, agenerator, alist, hook_provider):
    """Test that model hooks are correctly emitted even when throttled."""
    # Set up the model to raise throttling exceptions multiple times before succeeding
    exception = ModelThrottledException("ThrottlingException | ConverseStream")
    model.stream.side_effect = [
        exception,
        exception,
        exception,
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    await alist(stream)

    count, events = hook_provider.get_events()

    assert count == 9

    # 1st call - throttled
    assert next(events) == BeforeModelCallEvent(agent=agent, invocation_state=ANY)
    expected_after = AfterModelCallEvent(agent=agent, invocation_state=ANY, stop_response=None, exception=exception)
    expected_after.retry = True
    assert next(events) == expected_after

    # 2nd call - throttled
    assert next(events) == BeforeModelCallEvent(agent=agent, invocation_state=ANY)
    expected_after = AfterModelCallEvent(agent=agent, invocation_state=ANY, stop_response=None, exception=exception)
    expected_after.retry = True
    assert next(events) == expected_after

    # 3rd call - throttled
    assert next(events) == BeforeModelCallEvent(agent=agent, invocation_state=ANY)
    expected_after = AfterModelCallEvent(agent=agent, invocation_state=ANY, stop_response=None, exception=exception)
    expected_after.retry = True
    assert next(events) == expected_after

    # 4th call - successful
    assert next(events) == BeforeModelCallEvent(agent=agent, invocation_state=ANY)
    assert next(events) == AfterModelCallEvent(
        agent=agent,
        invocation_state=ANY,
        stop_response=AfterModelCallEvent.ModelStopResponse(
            message={"content": [{"text": "test text"}], "role": "assistant", "metadata": ANY}, stop_reason="end_turn"
        ),
        exception=None,
    )

    # Final message
    assert next(events) == MessageAddedEvent(
        agent=agent, message={"content": [{"text": "test text"}], "role": "assistant", "metadata": ANY}
    )


@pytest.mark.asyncio
async def test_event_loop_cycle_interrupt(agent, model, tool_stream, agenerator, alist):
    def interrupt_callback(event):
        event.interrupt("test_name", "test reason")

    agent.hooks.add_callback(BeforeToolCallEvent, interrupt_callback)

    model.stream.side_effect = [agenerator(tool_stream)]

    stream = strands.event_loop.event_loop.event_loop_cycle(agent, invocation_state={})
    events = await alist(stream)

    tru_stop_reason, _, _, _, tru_interrupts, _, _ = events[-1]["stop"]
    exp_stop_reason = "interrupt"
    exp_interrupts = [
        Interrupt(
            id="v1:before_tool_call:t1:78714d6c-613c-5cf4-bf25-7037569941f9",
            name="test_name",
            reason="test reason",
        ),
    ]

    assert tru_stop_reason == exp_stop_reason and tru_interrupts == exp_interrupts

    tru_state = agent._interrupt_state.to_dict()
    exp_state = {
        "activated": True,
        "context": {
            "tool_results": [],
            "tool_use_message": {
                "content": [
                    {
                        "toolUse": {
                            "input": {"random_string": "abcdEfghI123"},
                            "name": "tool_for_testing",
                            "toolUseId": "t1",
                        },
                    },
                ],
                "role": "assistant",
                "metadata": ANY,
            },
        },
        "interrupts": {
            "v1:before_tool_call:t1:78714d6c-613c-5cf4-bf25-7037569941f9": {
                "id": "v1:before_tool_call:t1:78714d6c-613c-5cf4-bf25-7037569941f9",
                "name": "test_name",
                "reason": "test reason",
                "response": None,
            },
        },
    }
    assert tru_state == exp_state


@pytest.mark.asyncio
async def test_event_loop_cycle_interrupt_resume(agent, model, tool, tool_times_2, agenerator, alist):
    interrupt = Interrupt(
        id="v1:before_tool_call:t1:78714d6c-613c-5cf4-bf25-7037569941f9",
        name="test_name",
        reason="test reason",
        response="test response",
    )

    tool_use_message = {
        "role": "assistant",
        "content": [
            {
                "toolUse": {
                    "toolUseId": "t1",
                    "name": "tool_for_testing",
                    "input": {"random_string": "test input"},
                }
            },
            {
                "toolUse": {
                    "toolUseId": "t2",
                    "name": "tool_times_2",
                    "input": {},
                }
            },
        ],
    }
    tool_results = [
        {
            "toolUseId": "t2",
            "status": "success",
            "content": [{"text": "t2 result"}],
        },
    ]

    agent._interrupt_state.context = {"tool_use_message": tool_use_message, "tool_results": tool_results}
    agent._interrupt_state.interrupts[interrupt.id] = interrupt
    agent._interrupt_state.activate()

    interrupt_response = {}

    def interrupt_callback(event):
        interrupt_response["response"] = event.interrupt("test_name", "test reason")

    agent.hooks.add_callback(BeforeToolCallEvent, interrupt_callback)

    model.stream.side_effect = [agenerator([{"contentBlockStop": {}}])]

    stream = strands.event_loop.event_loop.event_loop_cycle(agent, invocation_state={})
    events = await alist(stream)

    tru_stop_reason, _, _, _, _, _, _ = events[-1]["stop"]
    exp_stop_reason = "end_turn"
    assert tru_stop_reason == exp_stop_reason

    tru_result_message = agent.messages[-2]
    exp_result_message = {
        "role": "user",
        "content": [
            {
                "toolResult": {
                    "toolUseId": "t2",
                    "status": "success",
                    "content": [{"text": "t2 result"}],
                },
            },
            {
                "toolResult": {
                    "toolUseId": "t1",
                    "status": "success",
                    "content": [{"text": "test input"}],
                },
            },
        ],
    }
    assert tru_result_message == exp_result_message

    tru_response = interrupt_response["response"]
    exp_response = "test response"
    assert tru_response == exp_response

    tru_state = agent._interrupt_state.to_dict()
    exp_state = {
        "activated": False,
        "context": {},
        "interrupts": {},
    }
    assert tru_state == exp_state


@pytest.mark.asyncio
async def test_invalid_tool_names_adds_tool_uses(agent, model, alist):
    model.stream = MockedModelProvider(
        [
            {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": "tool_use_id",
                            "name": "invalid tool",
                            "input": "{}",
                        }
                    }
                ],
            },
            {"role": "assistant", "content": [{"text": "I invoked a tool!"}]},
        ]
    ).stream

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)

    # ensure that we got end_turn and not tool_use
    assert events[-1] == EventLoopStopEvent(
        stop_reason="end_turn",
        message={"content": [{"text": "I invoked a tool!"}], "role": "assistant", "metadata": ANY},
        metrics=ANY,
        request_state={},
    )

    # Ensure that an "invalid tool name" message was added properly
    assert agent.messages[-2] == {
        "content": [
            {
                "toolResult": {
                    "content": [{"text": "Error: tool_name=<invalid tool> | invalid tool name pattern"}],
                    "status": "error",
                    "toolUseId": "tool_use_id",
                }
            }
        ],
        "role": "user",
    }


@pytest.mark.asyncio
async def test_event_loop_metrics_recorded_before_recursion(
    agent,
    model,
    tool,
    agenerator,
    alist,
):
    model.stream.side_effect = [
        agenerator(
            [
                {
                    "contentBlockStart": {
                        "start": {
                            "toolUse": {
                                "toolUseId": "t1",
                                "name": tool.tool_spec["name"],
                            },
                        },
                    },
                },
                {"contentBlockStop": {}},
                {"messageStop": {"stopReason": "tool_use"}},
            ]
        ),
        agenerator(
            [
                {"contentBlockDelta": {"delta": {"text": "test text"}}},
                {"contentBlockStop": {}},
            ]
        ),
    ]

    with unittest.mock.patch.object(agent.event_loop_metrics, "end_cycle") as mock_end_cycle:
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={"request_state": {}},
        )
        events = await alist(stream)

        # Verify end_cycle was called once for tool cycle, once for text cycle
        assert mock_end_cycle.call_count == 2

        # Verify the event loop completed successfully
        tru_stop_reason, _, _, _, _, _, _ = events[-1]["stop"]
        assert tru_stop_reason == "end_turn"


class TestEstimateInputTokens:
    """Tests for _estimate_input_tokens helper."""

    @pytest.mark.asyncio
    async def test_cold_start_estimates_all_messages(self):
        """On cold start (no prior usage metadata), estimates all messages with lazily resolved tool specs."""
        agent = unittest.mock.AsyncMock()
        agent.messages = [{"role": "user", "content": [{"text": "Hi"}]}]
        agent.system_prompt = "You are helpful"
        agent._system_prompt_content = None
        agent.tool_registry = unittest.mock.MagicMock()
        agent.tool_registry.get_all_tool_specs.return_value = [{"name": "tool1"}]
        agent.model.count_tokens = AsyncMock(return_value=42)

        result = await strands.event_loop.event_loop._estimate_input_tokens(agent)

        assert result == 42
        agent.tool_registry.get_all_tool_specs.assert_called_once()
        agent.model.count_tokens.assert_called_once_with(
            agent.messages,
            tool_specs=[{"name": "tool1"}],
            system_prompt="You are helpful",
            system_prompt_content=None,
        )

    @pytest.mark.asyncio
    async def test_baseline_only_no_new_messages(self):
        """When last message is assistant with usage and no new messages after, returns baseline."""
        agent = unittest.mock.AsyncMock()
        agent.messages = [
            {"role": "user", "content": [{"text": "Hi"}]},
            {
                "role": "assistant",
                "content": [{"text": "Hello"}],
                "metadata": {"usage": {"inputTokens": 100, "outputTokens": 20, "totalTokens": 120}},
            },
        ]
        agent.system_prompt = "You are helpful"

        result = await strands.event_loop.event_loop._estimate_input_tokens(agent)

        assert result == 120
        agent.model.count_tokens.assert_not_called()

    @pytest.mark.asyncio
    async def test_baseline_plus_delta(self):
        """When new messages exist after last assistant, adds estimated delta to baseline."""
        agent = unittest.mock.AsyncMock()
        agent.messages = [
            {"role": "user", "content": [{"text": "Hi"}]},
            {
                "role": "assistant",
                "content": [{"text": "Hello"}],
                "metadata": {"usage": {"inputTokens": 100, "outputTokens": 30, "totalTokens": 130}},
            },
            {"role": "user", "content": [{"text": "tool result"}]},
        ]
        agent.system_prompt = "You are helpful"
        agent.model.count_tokens = AsyncMock(return_value=50)

        result = await strands.event_loop.event_loop._estimate_input_tokens(agent)

        # baseline (100+30) + delta (50) = 180
        assert result == 180
        agent.model.count_tokens.assert_called_once()

    @pytest.mark.asyncio
    async def test_error_fallback_returns_none_at_call_site(self):
        """When count_tokens raises, the caller catches and sets projected_input_tokens to None."""
        agent = unittest.mock.AsyncMock()
        agent.messages = [{"role": "user", "content": [{"text": "Hi"}]}]
        agent.system_prompt = "You are helpful"
        agent._system_prompt_content = None
        agent.tool_registry = unittest.mock.MagicMock()
        agent.tool_registry.get_all_tool_specs.return_value = []
        agent.model.count_tokens = AsyncMock(side_effect=Exception("API unavailable"))

        with pytest.raises(Exception, match="API unavailable"):
            await strands.event_loop.event_loop._estimate_input_tokens(agent)


# --- Checkpoint event loop integration (Tasks 9-10) ---


@pytest.mark.asyncio
async def test_event_loop_cycle_checkpoint_after_model(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    """With checkpointing=True, tool_use stop_reason yields after_model checkpoint instead of running tools."""
    agent._checkpointing = True
    agent._checkpoint = None

    model.stream.return_value = agenerator(tool_stream)

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    stop = events[-1]["stop"]
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = stop

    assert tru_stop_reason == "checkpoint"
    assert tru_checkpoint is not None
    assert tru_checkpoint.position == "after_model"
    assert tru_checkpoint.cycle_index == 0


@pytest.mark.asyncio
async def test_event_loop_cycle_checkpoint_after_tools(
    agent,
    model,
    tool,
    tool_stream,
    agenerator,
    alist,
):
    """With checkpointing=True and resume from after_model, tools execute then yield after_tools checkpoint."""
    agent._checkpointing = True
    agent._checkpoint = Checkpoint(position="after_model", cycle_index=0)

    model.stream.return_value = agenerator(tool_stream)

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = events[-1]["stop"]

    assert tru_stop_reason == "checkpoint"
    assert tru_checkpoint is not None
    assert tru_checkpoint.position == "after_tools"
    assert tru_checkpoint.cycle_index == 0


@pytest.mark.asyncio
async def test_event_loop_cycle_checkpoint_resume_after_tools_increments_cycle(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    """Resuming from after_tools sets cycle_index to previous + 1 for the next after_model checkpoint."""
    agent._checkpointing = True
    agent._checkpoint = Checkpoint(position="after_tools", cycle_index=2)

    model.stream.return_value = agenerator(tool_stream)

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = events[-1]["stop"]

    assert tru_stop_reason == "checkpoint"
    assert tru_checkpoint.position == "after_model"
    assert tru_checkpoint.cycle_index == 3


@pytest.mark.asyncio
async def test_event_loop_cycle_cancel_beats_after_model_checkpoint(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    """When a cancel signal is set after model call, cancel wins over after_model checkpoint.

    A user who calls agent.cancel() expects stop_reason="cancelled", not a stray
    "checkpoint" with a snapshot they never asked for. Documented in Agent.cancel().
    """
    agent._checkpointing = True
    agent._checkpoint = None

    # Cancel the agent before invoking. Model streams tool_use — the emission site
    # that would normally fire an after_model checkpoint must yield "cancelled" instead.
    agent._cancel_signal.set()
    model.stream.return_value = agenerator(tool_stream)

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = events[-1]["stop"]

    assert tru_stop_reason == "cancelled"
    assert tru_checkpoint is None


@pytest.mark.asyncio
async def test_event_loop_cycle_cancel_mid_cycle_beats_after_model_checkpoint(
    agent,
    model,
    tool_stream,
    agenerator,
    alist,
):
    """Cancel signal set between model completion and after_model emission yields 'cancelled', not 'checkpoint'."""
    agent._checkpointing = True
    agent._checkpoint = None

    # Stub the model stream so that after it yields its stop event, we simulate
    # a cancel signal arriving between model completion and the after_model check.
    original_stream = agenerator(tool_stream)

    async def stream_with_mid_cycle_cancel():
        async for item in original_stream:
            yield item
        agent._cancel_signal.set()

    model.stream.return_value = stream_with_mid_cycle_cancel()

    stream = strands.event_loop.event_loop.event_loop_cycle(
        agent=agent,
        invocation_state={},
    )
    events = await alist(stream)
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = events[-1]["stop"]

    assert tru_stop_reason == "cancelled"
    assert tru_checkpoint is None


@pytest.mark.asyncio
async def test_event_loop_cycle_cancel_mid_cycle_beats_after_tools_checkpoint(
    agent,
    model,
    tool,
    tool_stream,
    agenerator,
    alist,
):
    """Cancel set after tools complete but before after_tools emission yields 'cancelled'."""
    from strands.experimental.checkpoint import Checkpoint

    agent._checkpointing = True
    agent._checkpoint = Checkpoint(position="after_model", cycle_index=0)

    # Wrap the tool executor so that after tools complete, cancel is signaled.
    # The real gap: cancel arriving between tool completion and checkpoint emission.
    original_execute = agent.tool_executor._execute

    def execute_then_cancel(*args, **kwargs):
        stream = original_execute(*args, **kwargs)

        async def wrapped():
            async for event in stream:
                yield event
            agent._cancel_signal.set()

        return wrapped()

    model.stream.return_value = agenerator(tool_stream)

    with unittest.mock.patch.object(agent.tool_executor, "_execute", side_effect=execute_then_cancel):
        stream = strands.event_loop.event_loop.event_loop_cycle(
            agent=agent,
            invocation_state={},
        )
        events = await alist(stream)
    tru_stop_reason, _, _, _, _, _, tru_checkpoint = events[-1]["stop"]

    assert tru_stop_reason == "cancelled"
    assert tru_checkpoint is None
