import asyncio
import itertools
import logging
from contextlib import suppress
from functools import partial
from typing import Callable, List, Type, Optional
from unittest.mock import patch

import aiomisc
import aiormq.exceptions
import pytest
import shortuuid
from aiomisc_pytest import TCPProxy  # type: ignore
from yarl import URL

import aio_pika
from aio_pika.abc import AbstractRobustChannel, AbstractRobustConnection
from aio_pika.connection import Connection
from aio_pika.exceptions import QueueEmpty, CONNECTION_EXCEPTIONS
from aio_pika.message import Message
from aio_pika.robust_channel import RobustChannel
from aio_pika.robust_connection import RobustConnection
from aio_pika.robust_queue import RobustQueue
from tests import get_random_name


@pytest.fixture
async def proxy(tcp_proxy: Type[TCPProxy], amqp_direct_url: URL):
    p = tcp_proxy(
        amqp_direct_url.host,
        amqp_direct_url.port,
        buffered=False,
    )

    await p.start()
    try:
        yield p
    finally:
        await p.close()


@pytest.fixture
def amqp_url(request, amqp_direct_url, proxy: TCPProxy):
    return (
        amqp_direct_url.with_host(
            proxy.proxy_host,
        )
        .with_port(
            proxy.proxy_port,
        )
        .update_query(
            name=request.node.nodeid,
            reconnect_interval=1,
            heartbeat=1,
        )
    )


@pytest.fixture
def proxy_port(aiomisc_unused_port_factory) -> int:
    return aiomisc_unused_port_factory()


@pytest.fixture(scope="module")
def connection_fabric():
    return aio_pika.connect_robust


@pytest.fixture
def create_direct_connection(request, event_loop, amqp_direct_url):
    return partial(
        aio_pika.connect,
        amqp_direct_url.update_query(
            name=request.node.nodeid + "::direct",
            heartbeat=30,
        ),
        loop=event_loop,
    )


@pytest.fixture
def create_connection(connection_fabric, event_loop, amqp_url):
    return partial(connection_fabric, amqp_url, loop=event_loop)


@pytest.fixture
async def direct_connection(  # type: ignore
    create_direct_connection,
) -> aio_pika.Connection:
    async with await create_direct_connection() as conn:
        yield conn


async def test_channel_fixture(channel: aio_pika.RobustChannel):
    assert isinstance(channel, aio_pika.RobustChannel)


async def test_connection_fixture(connection: aio_pika.RobustConnection):
    assert isinstance(connection, aio_pika.RobustConnection)


def test_amqp_url_is_not_direct(amqp_url, amqp_direct_url):
    assert amqp_url != amqp_direct_url


async def test_set_qos(channel: aio_pika.Channel):
    await channel.set_qos(prefetch_count=1)


async def test_revive_passive_queue_on_reconnect(
    create_connection,
    direct_connection,
    proxy: TCPProxy,
):
    client = await create_connection()
    assert isinstance(client, RobustConnection)

    reconnect_event = asyncio.Event()
    reconnect_count = 0

    def reconnect_callback(conn: Optional[AbstractRobustConnection]):
        nonlocal reconnect_count
        reconnect_count += 1
        reconnect_event.set()
        reconnect_event.clear()

    client.reconnect_callbacks.add(reconnect_callback)

    queue_name = get_random_name()
    channel = await client.channel()
    assert isinstance(channel, RobustChannel)

    direct_channel = await direct_connection.channel()

    direct_queue = await direct_channel.declare_queue(
        queue_name,
        auto_delete=True,
        passive=False,
    )

    queue2 = await channel.declare_queue(
        direct_queue.name,
        passive=True,
        auto_delete=False,
    )
    assert isinstance(queue2, RobustQueue)

    await proxy.disconnect_all()
    await reconnect_event.wait()

    assert reconnect_count == 1

    with suppress(asyncio.TimeoutError):
        await asyncio.wait_for(
            reconnect_event.wait(),
            client.reconnect_interval * 2,
        )

    assert reconnect_count == 1


@aiomisc.timeout(30)
async def test_robust_reconnect(
    create_connection,
    direct_connection,
    proxy: TCPProxy,
    event_loop,
    add_cleanup: Callable,
):
    read_conn = await create_connection()  # type: aio_pika.RobustConnection

    reconnect_event = asyncio.Event()
    read_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    assert isinstance(read_conn, aio_pika.RobustConnection)
    write_channel = await direct_connection.channel()

    async with read_conn:
        read_channel = await read_conn.channel()

        assert isinstance(read_channel, aio_pika.RobustChannel)

        qname = get_random_name("robust", "proxy", "shared")

        async with read_channel:
            shared = []

            # Declaring temporary queue
            queue = await write_channel.declare_queue(
                qname,
                auto_delete=False,
                durable=True,
            )

            consumer_event = asyncio.Event()

            async def reader(queue_name):
                try:
                    queue = await read_channel.declare_queue(
                        name=queue_name,
                        passive=True,
                    )

                    async with queue.iterator() as q:
                        event_loop.call_soon(consumer_event.set)

                        async for message in q:
                            shared.append(message)
                            await message.ack()
                finally:
                    logging.info("Exit reader task")

            try:
                reader_task = event_loop.create_task(reader(queue.name))

                await consumer_event.wait()
                logging.info("Disconnect all clients")
                with proxy.slowdown(1, 1):
                    for i in range(5):
                        await write_channel.default_exchange.publish(
                            Message(str(i).encode()),
                            queue.name,
                        )

                    await proxy.disconnect_all()

                    # noinspection PyTypeChecker
                    with pytest.raises(aiormq.AMQPError):
                        await read_conn.channel()

                logging.info("Waiting reconnect")
                await reconnect_event.wait()

                logging.info("Waiting connections")
                await asyncio.wait_for(read_conn.ready(), timeout=20)

                for i in range(5, 10):
                    await write_channel.default_exchange.publish(
                        Message(str(i).encode()),
                        queue.name,
                    )

                while len(shared) < 10:
                    await asyncio.sleep(0.1)

                assert len(shared) == 10

                reader_task.cancel()
                await asyncio.gather(reader_task, return_exceptions=True)

                with pytest.raises(QueueEmpty):
                    await queue.get(timeout=0.5)
            finally:
                await queue.purge()

    # Waiting for rabbitmq queue not in use
    await asyncio.sleep(1)
    await queue.delete()


async def test_channel_locked_resource2(connection: aio_pika.RobustConnection):
    ch1: AbstractRobustChannel = await connection.channel()  # type: ignore
    ch2: AbstractRobustChannel = await connection.channel()  # type: ignore

    qname = get_random_name("channel", "locked", "resource")

    q1: aio_pika.abc.AbstractRobustQueue = await ch1.declare_queue(
        qname,
        exclusive=True,
        robust=False,
    )

    await q1.consume(print, exclusive=True)

    with pytest.raises(aiormq.exceptions.ChannelAccessRefused):
        q2 = await ch2.declare_queue(qname, exclusive=True, robust=False)
        await q2.consume(print, exclusive=True)


async def test_channel_close_when_exclusive_queue(
    create_connection,
    create_direct_connection,
    proxy: TCPProxy,
    event_loop,
):
    logging.info("Creating connections")
    direct_conn, proxy_conn = await asyncio.gather(
        create_direct_connection(),
        create_connection(),
    )

    logging.info("Creating channels")
    direct_channel, proxy_channel = await asyncio.gather(
        direct_conn.channel(),
        proxy_conn.channel(),
    )

    reconnect_event = asyncio.Event()
    proxy_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
        weak=False,
    )

    qname = get_random_name("robust", "exclusive", "queue")

    logging.info("Declaring exclusing queue: %s", qname)
    proxy_queue = await proxy_channel.declare_queue(
        qname,
        exclusive=True,
        durable=True,
    )

    logging.info("Disconnecting all proxy connections")
    await proxy.disconnect_all()
    await asyncio.sleep(0.5)

    logging.info("Declaring exclusive queue through direct channel")
    await direct_channel.declare_queue(
        qname,
        exclusive=True,
        durable=True,
    )

    async def close_after(delay, closer):
        await asyncio.sleep(delay)
        logging.info("Disconnecting direct connection")
        await closer()
        logging.info("Closed")

    await event_loop.create_task(close_after(5, direct_conn.close))

    # reconnect fired
    await reconnect_event.wait()

    # Wait method ready
    await proxy_conn.connected.wait()
    await proxy_queue.delete()


async def test_context_process_abrupt_channel_close(
    connection: aio_pika.RobustConnection,
    declare_exchange: Callable,
    declare_queue: Callable,
):
    # https://github.com/mosquito/aio-pika/issues/302
    queue_name = get_random_name("test_connection")
    routing_key = get_random_name("rounting_key")

    channel = await connection.channel()
    exchange = await declare_exchange(
        "direct",
        auto_delete=True,
        channel=channel,
    )
    queue = await declare_queue(queue_name, auto_delete=True, channel=channel)

    await queue.bind(exchange, routing_key)
    body = bytes(shortuuid.uuid(), "utf-8")

    await exchange.publish(
        Message(body, content_type="text/plain", headers={"foo": "bar"}),
        routing_key,
    )

    incoming_message = await queue.get(timeout=5)
    # close aiormq channel to emulate abrupt connection/channel close
    underlay_channel = await channel.get_underlay_channel()
    await underlay_channel.close()

    with pytest.raises(aiormq.exceptions.ChannelInvalidStateError):
        async with incoming_message.process():
            # emulate some activity on closed channel
            await channel.get_underlay_channel()

    # emulate connection/channel restoration of connect_robust
    await channel.reopen()

    # cleanup queue
    incoming_message = await queue.get(timeout=5)
    async with incoming_message.process():
        pass
    await queue.unbind(exchange, routing_key)


@aiomisc.timeout(10)
async def test_robust_duplicate_queue(
    connection: aio_pika.RobustConnection,
    direct_connection: aio_pika.Connection,
    declare_exchange: Callable,
    declare_queue: Callable,
    proxy: TCPProxy,
    create_task: Callable,
):
    queue_name = get_random_name("test")

    channel = await connection.channel()
    direct_channel = await direct_connection.channel()

    reconnect_event = asyncio.Event()
    shared_condition = asyncio.Condition()

    connection.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    shared = {}

    # noinspection PyShadowingNames
    async def reader(queue: aio_pika.Queue):
        async with queue.iterator() as q:
            async for message in q:
                # https://www.rabbitmq.com/confirms.html#automatic-requeueing
                async with shared_condition:
                    shared[message.message_id] = message
                    shared_condition.notify_all()
                    await message.ack()

    queue = await declare_queue(
        queue_name,
        channel=channel,
        cleanup=False,
    )

    create_task(reader(queue))

    for x in range(5):
        await direct_channel.default_exchange.publish(
            aio_pika.Message(b"1234567890", message_id=f"0-{x}"),
            queue_name,
        )

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 5),
            timeout=5,
        )

    logging.info("Disconnect all clients")
    await proxy.disconnect_all()

    assert len(shared) == 5, shared

    for x in range(5):
        await direct_channel.default_exchange.publish(
            Message(b"1234567890", message_id=f"1-{x}"),
            queue_name,
        )

    await asyncio.wait_for(reconnect_event.wait(), timeout=5)

    logging.info("Waiting connections")

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 10),
            timeout=5,
        )

    assert len(shared) == 10


@aiomisc.timeout(10)
async def test_channel_restore(
    connection_fabric,
    event_loop,
    amqp_url,
    proxy: TCPProxy,
    add_cleanup: Callable,
):
    heartbeat = 10
    amqp_url = amqp_url.update_query(heartbeat=heartbeat)

    on_reopen = asyncio.Event()

    conn = await connection_fabric(amqp_url, loop=event_loop)
    assert isinstance(conn, aio_pika.RobustConnection)

    async with conn:
        channel: AbstractRobustChannel = await conn.channel()  # type: ignore
        channel.reopen_callbacks.add(
            lambda *_: on_reopen.set(),
            weak=False,
        )
        assert isinstance(channel, aio_pika.RobustChannel)

        async with channel:
            await channel.set_qos(0)
            await channel.set_qos(1)

            with pytest.raises(asyncio.TimeoutError):
                with proxy.slowdown(1, 1):
                    await channel.set_qos(0, timeout=0.5)

            await on_reopen.wait()
            await channel.set_qos(0)
            await channel.set_qos(1)


@aiomisc.timeout(20)
async def test_channel_reconnect(
    connection_fabric,
    event_loop,
    amqp_url,
    proxy: TCPProxy,
    add_cleanup: Callable,
):
    on_reconnect = asyncio.Event()

    conn = await connection_fabric(amqp_url, loop=event_loop)
    assert isinstance(conn, aio_pika.RobustConnection)

    conn.reconnect_callbacks.add(lambda *_: on_reconnect.set(), weak=False)

    async with conn:
        channel = await conn.channel()
        assert isinstance(channel, aio_pika.RobustChannel)

        async with channel:
            await channel.set_qos(0)
            await channel.set_qos(1)

            await proxy.disconnect_all()

            await on_reconnect.wait()
            await channel.set_qos(0)
            await channel.set_qos(1)


class BadNetwork5KB:
    def __init__(self, proxy):
        self.proxy = proxy
        self.num_bytes = 0
        self.loop = asyncio.get_event_loop()
        self.lock = asyncio.Lock()

        proxy.set_content_processors(
            self.client_to_server,
            self.server_to_client,
        )

    async def disconnect(self):
        async with self.lock:
            await self.proxy.disconnect_all()
            self.num_bytes = 0

    async def server_to_client(self, chunk: bytes) -> bytes:
        async with self.lock:
            self.num_bytes += len(chunk)
            if self.num_bytes < 5000:
                return chunk
            self.loop.create_task(self.disconnect())
            return chunk

    @staticmethod
    def client_to_server(chunk: bytes) -> bytes:
        return chunk


@aiomisc.timeout(15)
@pytest.mark.parametrize(
    "reconnect_timeout",
    ["0", "1", "0.5", "0.1", "0.05", "0.025"],
)
async def test_channel_reconnect_after_5kb(
    reconnect_timeout,
    amqp_url,
    amqp_direct_url,
    connection_fabric,
    event_loop: asyncio.AbstractEventLoop,
    proxy: TCPProxy,
    add_cleanup: Callable,
):
    connection = await aio_pika.connect_robust(
        amqp_url.update_query(reconnect_interval=reconnect_timeout),
        loop=event_loop,
    )
    direct_connection = await aio_pika.connect(
        amqp_direct_url,
        loop=event_loop,
    )

    on_reconnect = asyncio.Event()
    connection.reconnect_callbacks.add(
        lambda *_: on_reconnect.set(),
        weak=False,
    )

    BadNetwork5KB(proxy)

    messages_to_exchange = 50
    async with connection.channel() as channel:
        await channel.set_qos(prefetch_count=5)
        queue = await channel.declare_queue(auto_delete=False)

        async with direct_connection.channel() as publish_channel:
            for _ in range(messages_to_exchange):
                await publish_channel.default_exchange.publish(
                    aio_pika.Message(body=b"Hello world " * 100),
                    routing_key=queue.name,
                )

        messages = []
        async for message in queue.iterator():
            messages.append(message)
            if len(messages) == messages_to_exchange:
                break

        assert messages

    assert on_reconnect.is_set()
    await connection.close()
    await direct_connection.close()


class BadNetwork:
    def __init__(self, proxy, stair: int, disconnect_time: float):
        self.proxy = proxy
        self.stair = stair
        self.disconnect_time = disconnect_time
        self.num_bytes = 0
        self.loop = asyncio.get_event_loop()
        self.lock = asyncio.Lock()

        proxy.set_content_processors(
            self.client_to_server,
            self.server_to_client,
        )

    async def disconnect(self):
        async with self.lock:
            await asyncio.sleep(self.disconnect_time)
            await self.proxy.disconnect_all()
            self.stair *= 2
            self.num_bytes = 0

    async def server_to_client(self, chunk: bytes) -> bytes:
        async with self.lock:
            self.num_bytes += len(chunk)
            if self.num_bytes < self.stair:
                return chunk
            self.loop.create_task(self.disconnect())
            return chunk

    @staticmethod
    def client_to_server(chunk: bytes) -> bytes:
        return chunk


DISCONNECT_OFFSETS = [2 << i for i in range(5, 12)]
STAIR_STEPS = list(itertools.product([0.1, 0.0], DISCONNECT_OFFSETS))
STAIR_STEPS_IDS = [
    f"[{i // len(DISCONNECT_OFFSETS)}] {t}-{s}"
    for i, (t, s) in enumerate(STAIR_STEPS)
]


@aiomisc.timeout(30)
async def test_connection_stuck_reconnect(
    connection_fabric,
    event_loop: asyncio.AbstractEventLoop,
    amqp_url: URL,
    proxy: TCPProxy,
):
    """Test that RobustConnection reconnects after detecting a stuck connection.

    This test simulates a scenario where the server stops responding (no
    heartbeats received) but the connection is not explicitly closed. The
    connection should be detected as "stuck" after the heartbeat grace timeout
    and automatically reconnect.

    Related issues: #577, #563, #588, #620
    """
    # Use a short heartbeat for faster test execution
    # heartbeat=2 means grace timeout = (2+1) * 3 = 9 seconds
    heartbeat = 2
    test_url = amqp_url.update_query(heartbeat=heartbeat)

    connection = await connection_fabric(test_url, loop=event_loop)
    assert isinstance(connection, RobustConnection)

    reconnect_event = asyncio.Event()
    reconnect_count = 0

    def on_reconnect(*_):
        nonlocal reconnect_count
        reconnect_count += 1
        reconnect_event.set()

    connection.reconnect_callbacks.add(on_reconnect, weak=False)

    async with connection:
        channel = await connection.channel()
        assert isinstance(channel, RobustChannel)

        # Verify connection is working
        await channel.set_qos(prefetch_count=1)

        # Simulate stuck connection by using slowdown context manager
        # The slowdown makes traffic very slow, simulating a stuck connection
        # After a brief period, disconnect to trigger the stuck detection
        # and allow reconnection to proceed
        grace_timeout = (heartbeat + 1) * 3

        # Use slowdown to delay traffic, making heartbeats fail
        with proxy.slowdown(read_delay=grace_timeout + 2, write_delay=0):
            # Wait long enough for the connection to be detected as stuck
            # but disconnect before the full delay so reconnection can start
            await asyncio.sleep(grace_timeout + 1)
            # Disconnect to force the stuck connection to close
            await proxy.disconnect_all()

        # Wait for reconnection to complete
        await asyncio.wait_for(reconnect_event.wait(), timeout=10)

        # Wait for the connection to be fully ready again
        await asyncio.wait_for(connection.ready(), timeout=10)

        # Verify reconnection happened
        assert reconnect_count >= 1, "Connection should have reconnected"

        # Verify connection is working after reconnection
        await channel.set_qos(prefetch_count=1)


@aiomisc.timeout(30)
async def test_queue_iterator_reconnect_after_disconnect(
    connection_fabric,
    event_loop: asyncio.AbstractEventLoop,
    amqp_url: URL,
    amqp_direct_url: URL,
    proxy: TCPProxy,
):
    """Test that queue.iterator() continues working after connection reconnect.

    This test simulates the scenario from issue reports where users have a
    consumer using queue.iterator(timeout=X) in a loop and experience
    connection issues that should trigger automatic reconnection.

    Related issues: #577, #563, #588, #620
    """
    connection = await connection_fabric(amqp_url, loop=event_loop)
    direct_connection = await aio_pika.connect(amqp_direct_url, loop=event_loop)

    assert isinstance(connection, RobustConnection)

    reconnect_event = asyncio.Event()
    connection.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
        weak=False,
    )

    messages_received: List[bytes] = []
    consumer_ready = asyncio.Event()
    consumer_error: Optional[Exception] = None

    async with connection, direct_connection:
        channel = await connection.channel()
        direct_channel = await direct_connection.channel()

        assert isinstance(channel, RobustChannel)

        queue_name = get_random_name("iterator", "reconnect")
        queue = await channel.declare_queue(queue_name, auto_delete=True)

        async def consumer():
            nonlocal consumer_error
            try:
                consumer_ready.set()
                async with queue.iterator() as queue_iter:
                    async for message in queue_iter:
                        messages_received.append(message.body)
                        await message.ack()
            except Exception as e:
                consumer_error = e
                raise

        consumer_task = event_loop.create_task(consumer())

        try:
            await consumer_ready.wait()

            # Publish some messages before disconnect
            for i in range(3):
                await direct_channel.default_exchange.publish(
                    Message(f"before-{i}".encode()),
                    routing_key=queue_name,
                )

            # Wait for messages to be received
            while len(messages_received) < 3:
                await asyncio.sleep(0.1)

            assert len(messages_received) == 3

            # Disconnect and trigger reconnection
            await proxy.disconnect_all()
            await asyncio.wait_for(reconnect_event.wait(), timeout=10)
            await asyncio.wait_for(connection.ready(), timeout=10)

            # Publish more messages after reconnection
            for i in range(3):
                await direct_channel.default_exchange.publish(
                    Message(f"after-{i}".encode()),
                    routing_key=queue_name,
                )

            # Wait for all messages to be received
            while len(messages_received) < 6:
                await asyncio.sleep(0.1)

            assert len(messages_received) == 6
            assert consumer_error is None

        finally:
            consumer_task.cancel()
            with suppress(asyncio.CancelledError):
                await consumer_task


@aiomisc.timeout(30)
@pytest.mark.parametrize(
    "reconnect_timeout,stair",
    STAIR_STEPS,
    ids=STAIR_STEPS_IDS,
)
async def test_channel_reconnect_stairway(
    reconnect_timeout: float,
    stair: int,
    amqp_url: URL,
    amqp_direct_url: URL,
    connection_fabric,
    event_loop: asyncio.AbstractEventLoop,
    proxy: TCPProxy,
    add_cleanup: Callable,
):

    event_loop.set_debug(True)

    connection = await aio_pika.connect_robust(
        amqp_url.update_query(
            reconnect_interval=f"{reconnect_timeout:.2f}",
            name="proxy",
        ),
        loop=event_loop,
    )
    direct_connection = await aio_pika.connect(
        amqp_direct_url.update_query("name=direct"),
        loop=event_loop,
    )

    on_reconnect = asyncio.Event()
    connection.reconnect_callbacks.add(
        lambda *_: on_reconnect.set(),
        weak=False,
    )

    BadNetwork(proxy, stair, reconnect_timeout)

    messages_to_exchange = 100
    body = b"Hello world " * 1000

    async with connection.channel() as channel:
        queue = await channel.declare_queue(auto_delete=False)

        async with direct_connection.channel() as publish_channel:
            for _ in range(messages_to_exchange):
                await publish_channel.default_exchange.publish(
                    aio_pika.Message(body=body),
                    routing_key=queue.name,
                )

        messages: List[aio_pika.abc.AbstractIncomingMessage] = []

        while True:
            try:
                await channel.set_qos(prefetch_count=1)
                break
            except CONNECTION_EXCEPTIONS:
                await asyncio.sleep(0.1)
                continue

        while len(messages) < messages_to_exchange:
            try:
                message: aio_pika.abc.AbstractIncomingMessage
                async for message in queue.iterator():
                    # noinspection PyBroadException
                    try:
                        await message.ack()
                    except Exception:
                        continue

                    messages.append(message)

                    if len(messages) >= messages_to_exchange:
                        break
            except Exception:
                continue

        assert messages

    assert on_reconnect.is_set()
    await connection.close()
    await direct_connection.close()


@aiomisc.timeout(60)
async def test_iterator_survives_reconnection(
    create_connection,
    direct_connection,
    proxy: TCPProxy,
):
    """Test that RobustQueueIterator survives reconnection without exiting.

    This test verifies the fix for issue #540 where consumers appeared
    connected but were stuck with messages in unack state after reconnection.
    """
    connection = await create_connection()
    assert isinstance(connection, RobustConnection)

    reconnect_event = asyncio.Event()
    connection.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    async with connection:
        channel = await connection.channel()
        assert isinstance(channel, RobustChannel)

        queue = await channel.declare_queue(auto_delete=False)
        assert isinstance(queue, RobustQueue)

        direct_channel = await direct_connection.channel()

        # Start consuming with iterator
        iterator = queue.iterator()
        await iterator.__aenter__()

        # Publish first message
        await direct_channel.default_exchange.publish(
            Message(b"test message 1"),
            routing_key=queue.name,
        )

        # Get first message
        msg1 = await asyncio.wait_for(iterator.__anext__(), timeout=10)
        await msg1.ack()
        assert msg1.body == b"test message 1"

        # Small delay to ensure ack is processed before disconnect
        await asyncio.sleep(0.1)

        # Disconnect (simulate network issue)
        await proxy.disconnect_all()

        # Wait for reconnection
        await asyncio.wait_for(reconnect_event.wait(), timeout=30)
        reconnect_event.clear()

        # Wait for channel to be ready after reconnection
        await asyncio.wait_for(channel.ready(), timeout=10)

        # Publish another message after reconnection
        await direct_channel.default_exchange.publish(
            Message(b"test message 2"),
            routing_key=queue.name,
        )

        # Iterator should still work (not raise StopAsyncIteration)
        # This is the key assertion - before the fix, this would fail
        # because the iterator would have exited during reconnection
        received_messages = []
        for _ in range(2):  # Try to receive up to 2 messages
            try:
                msg = await asyncio.wait_for(iterator.__anext__(), timeout=5)
                received_messages.append(msg.body)
                await msg.ack()
                if msg.body == b"test message 2":
                    break
            except asyncio.TimeoutError:
                break

        # The important thing is that we received message 2 after reconnection
        # Message 1 might be redelivered if ack didn't complete before
        # disconnect
        assert b"test message 2" in received_messages, (
            f"Did not receive message 2. Got: {received_messages}"
        )

        await iterator.__aexit__(None, None, None)

        # Clean up queue
        await queue.delete()


@aiomisc.timeout(30)
async def test_consumer_resumes_after_connection_close(
    create_connection,
    direct_connection,
    proxy: TCPProxy,
):
    """Test that callback-based consumers resume after connection close.

    This test verifies the fix for issue #588 where consumers would stop
    processing messages after reconnection. The issue was that
    RobustChannel._on_close() prematurely set _closed.set_result(True)
    when receiving CancelledError, which could interfere with restoration.

    This test uses queue.consume(callback) instead of the iterator pattern
    tested in test_iterator_survives_reconnection.
    """
    connection = await create_connection()
    assert isinstance(connection, RobustConnection)

    reconnect_event = asyncio.Event()
    reconnect_count = 0

    def on_reconnect(*_):
        nonlocal reconnect_count
        reconnect_count += 1
        reconnect_event.set()

    connection.reconnect_callbacks.add(on_reconnect, weak=False)

    async with connection:
        channel = await connection.channel()
        assert isinstance(channel, RobustChannel)

        queue = await channel.declare_queue(auto_delete=False)
        assert isinstance(queue, RobustQueue)
        direct_channel = await direct_connection.channel()

        # Set up callback-based consumer with condition for precise waiting
        received_messages = []
        message_condition = asyncio.Condition()

        async def consumer(message: aio_pika.abc.AbstractIncomingMessage):
            async with message_condition:
                received_messages.append(message.body)
                message_condition.notify_all()
            # Ack after notification to avoid race conditions
            with suppress(Exception):
                await message.ack()

        await queue.consume(consumer)

        # Publish and receive first message to confirm consumer works
        await direct_channel.default_exchange.publish(
            Message(b"before disconnect"),
            routing_key=queue.name,
        )

        async with message_condition:
            await asyncio.wait_for(
                message_condition.wait_for(
                    lambda: b"before disconnect" in received_messages,
                ),
                timeout=10,
            )

        # Disconnect to trigger reconnection
        await proxy.disconnect_all()

        # Wait for reconnection
        await asyncio.wait_for(reconnect_event.wait(), timeout=20)
        reconnect_event.clear()

        # Wait for channel to be fully restored
        await asyncio.wait_for(channel.ready(), timeout=10)

        # Publish message after reconnection
        await direct_channel.default_exchange.publish(
            Message(b"after disconnect"),
            routing_key=queue.name,
        )

        # Wait specifically for the "after disconnect" message
        async with message_condition:
            await asyncio.wait_for(
                message_condition.wait_for(
                    lambda: b"after disconnect" in received_messages,
                ),
                timeout=10,
            )

        # Verify we had at least one reconnection
        assert reconnect_count >= 1

        # Clean up - suppress errors during cleanup
        with suppress(Exception):
            await queue.delete()


async def test_close_does_not_hang_during_reconnect(event_loop):
    """Test that close() does not hang when reconnection loop is active.

    Reproduces the race condition where close() cancels __reconnection_task
    while __connection_factory is inside Connection.connect(). Without the
    fix (_close_called set before cancel), CancelledError would be swallowed
    by the except block and close() would hang waiting for the task to finish.
    """
    connect_entered = asyncio.Event()

    async def slow_connect(self, timeout=None):
        connect_entered.set()
        # Simulate a long connection attempt that will be cancelled
        await asyncio.sleep(3600)

    url = URL(
        "amqp://guest:guest@localhost/?reconnect_interval=0.1&fail_fast=0"
    )
    conn = RobustConnection(url, loop=event_loop)

    with patch.object(Connection, "connect", slow_connect):
        # Start the connection (launches __connection_factory task).
        # connect() will await connected.wait() which never resolves
        # since Connection.connect is mocked, so run it as a task.
        connect_task = event_loop.create_task(conn.connect())

        # Wait until __connection_factory enters Connection.connect()
        await asyncio.wait_for(connect_entered.wait(), timeout=5)

        # Now close() should cancel the reconnection task and return
        # promptly. Before the fix this would hang because CancelledError
        # was swallowed (_close_called was not yet set when cancel arrived).
        await asyncio.wait_for(conn.close(), timeout=5)

    # Cancel connect_task which is stuck on connected.wait()
    connect_task.cancel()
    with suppress(asyncio.CancelledError):
        await connect_task
