import json
from uuid import uuid4
import aio_pika
from faststream.rabbit import RabbitBroker


def connection_for_broker(broker: RabbitBroker) -> aio_pika.RobustConnection:
    connection = broker._connection
    if connection is None:
        raise Exception("Broker not connected")

    return connection


async def custom_wait_for_reply(broker: RabbitBroker, *args, **kwargs):
    connection = connection_for_broker(broker)

    channel = await connection.channel()

    queue = await channel.declare_queue(
        name="reply" + str(uuid4()),
        exclusive=True,
        auto_delete=True,
    )

    kwargs_updated = {**kwargs}
    kwargs_updated.update(dict(reply_to=queue.name))
    await broker.publish(*args, **kwargs_updated)

    async for message in queue:  # type: ignore
        result = json.loads(message.body.decode())
        break

    return result  # type: ignore
