import json
import time
from typing import TypedDict

import pika
from pika.adapters.blocking_connection import BlockingChannel
from pika.spec import Basic, BasicProperties

conn = pika.BlockingConnection(pika.ConnectionParameters("localhost"))
ch = conn.channel()

# rq will wait for ack or failure before sending more message
# this is 0 = unlimited by default
# not given I use BlockingConnection, this won't run in parallel at all
# but basically that's our impl, as far as q is concerned, it has "sent"
# messages for processing
ch.basic_qos(prefetch_count=1)


class OrderMessage(TypedDict):
    order_id: int
    user_id: str
    amount: float
    should_fail: bool
    retry_count: int


def publish_status(routing_key: str, payload: OrderMessage) -> None:
    ch.basic_publish(
        exchange="orders.x",
        routing_key=routing_key,
        body=json.dumps(payload),
        properties=pika.BasicProperties(
            delivery_mode=2,
            content_type="application/json",
        ),
    )


def on_message(
    channel: BlockingChannel,
    method: Basic.Deliver,
    properties: BasicProperties,
    body: bytes,
) -> None:
    _ = properties
    order: OrderMessage = json.loads(body)
    order_id = order["order_id"]
    retry_count = order.get("retry_count", 0)
    should_fail = order.get("should_fail", False)

    print(f"[worker] got order={order_id} retry={retry_count}")
    time.sleep(1)

    if should_fail and retry_count < 2:
        order["retry_count"] = retry_count + 1
        publish_status("order.created", order)  # this key, this obj
        # ack back on the channel is separate
        # even though this LOOKS optional I feel that's just bad api
        # as it'll fail at rutime if you miss it
        channel.basic_ack(delivery_tag=method.delivery_tag)
        print(
            f"[worker] retrying order={order_id} retry={order['retry_count']}")
        return

    if should_fail and retry_count >= 2:
        publish_status("order.failed", order)
        channel.basic_ack(delivery_tag=method.delivery_tag)
        print(f"[worker] failed order={order_id} after retries")
        return

    publish_status("order.processed", order)
    channel.basic_ack(delivery_tag=method.delivery_tag)
    print(f"[worker] processed order={order_id}")


ch.basic_consume(queue="orders.q",
                 on_message_callback=on_message, auto_ack=False)
print("[worker] waiting for messages...")
ch.start_consuming()