import asyncio
import aio_pika

from config import RABBITMQ_URI


async def setup_direct_exchange(exchange_name="demo.direct", queue_prefix="demo.direct.queue-"):
    """设置Direct类型交换器并绑定队列"""
    # 建立连接
    connection = await aio_pika.connect_robust(RABBITMQ_URI)
    channel = await connection.channel()

    # 1. 声明Direct类型交换器
    direct_exchange = await channel.declare_exchange(
        exchange_name,
        aio_pika.ExchangeType.DIRECT,
        durable=True  # 交换器持久化
    )

    # 2. 定义队列及对应的绑定键（Direct交换器需要精确匹配）
    # 格式: (队列名称, 绑定键列表)
    queue_bindings = [
        (f"{queue_prefix}error", ["error"]),  # 处理错误级别的消息
        (f"{queue_prefix}warning", ["warning"]),  # 处理警告级别的消息
        (f"{queue_prefix}info", ["info", "debug"])  # 处理信息和调试级别的消息
    ]

    # 3. 循环创建队列并绑定到交换器
    for queue_name, binding_keys in queue_bindings:
        # 声明队列（持久化）
        queue = await channel.declare_queue(
            queue_name,
            durable=True,
            auto_delete=False
        )
        # 绑定多个路由键到同一个队列
        for binding_key in binding_keys:
            await queue.bind(
                direct_exchange,
                routing_key=binding_key  # Direct交换器需要指定绑定键
            )
        print(f"Queue {queue_name} bound to routing keys: {binding_keys}")

    await connection.close()


async def direct_publish(message: str, routing_key: str, exchange_name: str = "demo.direct"):
    """向Direct交换器发送消息"""
    # 建立连接
    connection = await aio_pika.connect_robust(RABBITMQ_URI)
    channel = await connection.channel()

    # 声明交换器（确保交换器存在）
    exchange = await channel.declare_exchange(
        exchange_name,
        aio_pika.ExchangeType.DIRECT,
        durable=True
    )

    # 构建消息对象（持久化）
    message_obj = aio_pika.Message(
        body=message.encode("utf-8"),
        delivery_mode=aio_pika.DeliveryMode.PERSISTENT
    )

    # 发送消息（指定路由键，Direct交换器会根据路由键匹配队列）
    await exchange.publish(
        message_obj,
        routing_key=routing_key  # 路由键决定消息流向哪个队列
    )
    print(f"Message sent: {message} (routing key: {routing_key})")

    await connection.close()