import aio_pika
from config import RABBITMQ_URI


async def setup_multi_queue_balance(
        exchange_name="demo.direct.multi.queue",
        queue_prefix="task.queue.",
        route_prefix="route.",
        queue_count=3  # 3个队列
):
    """创建多个队列，每个队列绑定不同的路由键"""
    connection = await aio_pika.connect_robust(RABBITMQ_URI)
    channel = await connection.channel()

    # 声明Direct交换器
    await channel.declare_exchange(
        exchange_name,
        aio_pika.ExchangeType.DIRECT,
        durable=True
    )

    # 创建N个队列，每个队列绑定独立的路由键
    for i in range(queue_count):
        queue_name = f"{queue_prefix}{i + 1}"
        route_key = f"{route_prefix}{i + 1}"

        # 声明队列
        queue = await channel.declare_queue(
            queue_name,
            durable=True,
            auto_delete=False
        )

        # 绑定路由键（每个队列对应唯一路由键）
        await queue.bind(exchange_name, routing_key=route_key)
        print(f"Queue {queue_name} bound to routing key: {route_key}")

    await connection.close()


class BalancedProducer:
    def __init__(self, exchange_name="demo.direct.multi.queue", queue_count=3):
        self.exchange_name = exchange_name
        self.queue_count = queue_count
        self.current_index = 0  # 轮询索引

    async def connect(self):
        """建立连接（可复用连接提升性能）"""
        self.connection = await aio_pika.connect_robust(RABBITMQ_URI)
        self.channel = await self.connection.channel()
        self.exchange = await self.channel.declare_exchange(
            self.exchange_name,
            aio_pika.ExchangeType.DIRECT,
            durable=True
        )

    async def publish(self, message: str):
        """轮询选择路由键，发送消息到对应的队列"""
        # 轮询算法：每次发送后切换到下一个路由键
        self.current_index = (self.current_index + 1) % self.queue_count
        route_key = f"route.{self.current_index + 1}"  # 对应 route_1/2/3

        message_obj = aio_pika.Message(
            body=message.encode("utf-8"),
            delivery_mode=aio_pika.DeliveryMode.PERSISTENT
        )

        await self.exchange.publish(message_obj, routing_key=route_key)
        print(f"Message sent: {message} (routed to {route_key})")

    async def close(self):
        """关闭连接"""
        await self.connection.close()