# requires: aio_pika
"""RabbitMQ 延迟取消订单：消息先进带 TTL 的等待队列，过期后经死信交换机转到处理队列。

运行：AMQP_URL=amqp://guest:guest@localhost/ python delay_cancel.py
没有设置 AMQP_URL 时打印提示并正常退出。会创建并删除名字带随机后缀的交换机和队列。
"""
import asyncio
import json
import os
import sys
import time
import uuid

import aio_pika

DELAY_MS = 1000  # 演示用 1 秒；真实业务常见是 15~30 分钟


class OrderStore:
    """模拟数据库里的订单状态（真实项目里用条件 UPDATE 保证幂等）。"""

    def __init__(self) -> None:
        self.status: dict[str, str] = {}

    def cancel_if_unpaid(self, order_id: str) -> bool:
        # 等价于 UPDATE orders SET status='cancelled' WHERE id=? AND status='pending'
        if self.status.get(order_id) == "pending":
            self.status[order_id] = "cancelled"
            return True
        return False


async def main() -> int:
    url = os.environ.get("AMQP_URL")
    if not url:
        print("未设置 AMQP_URL，跳过（示例：amqp://guest:guest@localhost/）")
        return 0
    suffix = uuid.uuid4().hex[:8]
    conn = await aio_pika.connect_robust(url)
    async with conn:
        ch = await conn.channel()
        await ch.set_qos(prefetch_count=10)
        dlx = await ch.declare_exchange(f"order.dlx.{suffix}", aio_pika.ExchangeType.DIRECT, auto_delete=True)
        # 处理队列：真正执行“超时取消”的消费者监听这里
        work_q = await ch.declare_queue(f"order.cancel.{suffix}", durable=True, auto_delete=True)
        await work_q.bind(dlx, routing_key="cancel")
        # 等待队列：没有消费者，消息在这里过期后转到死信交换机
        wait_q = await ch.declare_queue(f"order.wait.{suffix}", durable=True, auto_delete=True, arguments={
            "x-message-ttl": DELAY_MS,
            "x-dead-letter-exchange": dlx.name,
            "x-dead-letter-routing-key": "cancel",
        })

        store = OrderStore()
        store.status = {"o1": "pending", "o2": "pending"}
        start = time.monotonic()
        for oid in ("o1", "o2", "o1"):  # o1 故意发两次，模拟重复消息
            await ch.default_exchange.publish(
                aio_pika.Message(json.dumps({"order_id": oid}).encode(),
                                 delivery_mode=aio_pika.DeliveryMode.PERSISTENT, message_id=f"cancel-{oid}"),
                routing_key=wait_q.name)
        store.status["o2"] = "paid"  # o2 在超时前支付了

        results: list[tuple[str, bool]] = []
        async with work_q.iterator() as it:
            async for msg in it:
                async with msg.process():  # 正常结束自动 ack，抛异常则 nack
                    oid = json.loads(msg.body)["order_id"]
                    results.append((oid, store.cancel_if_unpaid(oid)))
                if len(results) == 3:
                    break
        elapsed = time.monotonic() - start
        print("处理结果：", results, f"耗时 {elapsed:.1f}s")
        assert elapsed >= DELAY_MS / 1000 * 0.9, "消息应该在延迟之后才被处理"
        assert store.status == {"o1": "cancelled", "o2": "paid"}, store.status
        assert sorted(results) == [("o1", False), ("o1", True), ("o2", False)]  # 重复消息不会重复取消
        await wait_q.delete(if_unused=False, if_empty=False)
        await work_q.delete(if_unused=False, if_empty=False)
        print("延迟取消正确：未支付的被取消，已支付的不受影响，重复消息幂等")
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))