import asyncio
import logging
from unittest.mock import AsyncMock, MagicMock

import pytest
from a2a.server.events.event_queue import EventQueue
from a2a.types import TaskState, TaskStatusUpdateEvent

from codex_a2a.execution.executor import CodexAgentExecutor
from codex_a2a.upstream.client import CodexClient
from tests.support.context import configure_mock_client_runtime, make_request_context_mock


@pytest.mark.asyncio
async def test_cancel_interrupts_running_execute_and_keeps_queue_open():
    client = AsyncMock(spec=CodexClient)
    send_started = asyncio.Event()
    send_cancelled = asyncio.Event()

    async def send_message(
        session_id,
        _text,
        *,
        directory=None,  # noqa: ARG001
        execution_options=None,  # noqa: ARG001
        timeout_override=None,  # noqa: ARG001
    ):
        send_started.set()
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            send_cancelled.set()
            raise
        response = MagicMock()
        response.text = "Codex response"
        response.session_id = session_id
        response.message_id = "msg-1"
        return response

    client.create_session.return_value = "session-1"
    client.send_message.side_effect = send_message
    configure_mock_client_runtime(client)

    executor = CodexAgentExecutor(client, streaming_enabled=False)

    execute_context = make_request_context_mock(
        task_id="task-1",
        context_id="context-A",
        identity="user-1",
        user_input="hello",
    )
    execute_queue = AsyncMock(spec=EventQueue)

    execute_task = asyncio.create_task(executor.execute(execute_context, execute_queue))
    await asyncio.wait_for(send_started.wait(), timeout=1.0)

    cancel_context = make_request_context_mock(
        task_id="task-1",
        context_id="context-A",
        call_context_enabled=False,
    )
    cancel_queue = AsyncMock(spec=EventQueue)

    await asyncio.wait_for(executor.cancel(cancel_context, cancel_queue), timeout=1.0)

    cancel_events = [call.args[0] for call in cancel_queue.enqueue_event.call_args_list]
    assert any(
        isinstance(event, TaskStatusUpdateEvent) and event.status.state == TaskState.canceled
        for event in cancel_events
    )
    cancel_queue.close.assert_not_called()

    with pytest.raises(asyncio.CancelledError):
        await asyncio.wait_for(execute_task, timeout=1.0)

    assert send_cancelled.is_set()
    runtime = executor._session_runtime
    assert await runtime.bound_session_for(identity="user-1", context_id="context-A") is None
    assert (
        await runtime.running_execution_snapshot(task_id="task-1", context_id="context-A") is None
    )


@pytest.mark.asyncio
async def test_cancel_does_not_block_with_real_event_queue() -> None:
    executor = CodexAgentExecutor(MagicMock(), streaming_enabled=False)
    context = make_request_context_mock(
        task_id=None,
        context_id=None,
        call_context_enabled=False,
    )
    queue = EventQueue()

    await asyncio.wait_for(executor.cancel(context, queue), timeout=0.5)


@pytest.mark.asyncio
async def test_cancel_logs_abort_timeout_when_cleanup_does_not_finish(
    caplog: pytest.LogCaptureFixture,
) -> None:
    client = AsyncMock(spec=CodexClient)
    send_started = asyncio.Event()
    release_send = asyncio.Event()

    async def send_message(
        session_id,
        _text,
        *,
        directory=None,  # noqa: ARG001
        execution_options=None,  # noqa: ARG001
        timeout_override=None,  # noqa: ARG001
    ):
        send_started.set()
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            await release_send.wait()
        response = MagicMock()
        response.text = "Codex response"
        response.session_id = session_id
        response.message_id = "msg-1"
        return response

    client.create_session.return_value = "session-1"
    client.send_message.side_effect = send_message
    configure_mock_client_runtime(client)

    executor = CodexAgentExecutor(
        client,
        streaming_enabled=False,
        cancel_abort_timeout_seconds=0.01,
    )

    execute_context = make_request_context_mock(
        task_id="task-timeout",
        context_id="context-timeout",
        identity="user-1",
        user_input="hello",
    )
    execute_queue = AsyncMock(spec=EventQueue)
    execute_task = asyncio.create_task(executor.execute(execute_context, execute_queue))
    await asyncio.wait_for(send_started.wait(), timeout=1.0)

    cancel_context = make_request_context_mock(
        task_id="task-timeout",
        context_id="context-timeout",
        call_context_enabled=False,
    )
    cancel_queue = AsyncMock(spec=EventQueue)
    caplog.set_level(logging.WARNING, logger="codex_a2a.execution.executor")

    await asyncio.wait_for(executor.cancel(cancel_context, cancel_queue), timeout=0.5)

    assert any("Cancel abort timeout exceeded" in record.getMessage() for record in caplog.records)

    release_send.set()
    await asyncio.wait_for(execute_task, timeout=1.0)
