"""CloudTrips: send interleaved orders and verify processing within each session."""
import argparse
import json
import uuid

from azure.identity import AzureCliCredential
from azure.servicebus import ServiceBusClient, ServiceBusMessage, TransportType


def run(client, queue):
    run_id = uuid.uuid4().hex[:12]
    sessions = {order: f"{run_id}-{order}" for order in ("CT-3001", "CT-3002")}
    stages = ("Created", "Paid", "Shipped")
    print(f"Run: {run_id}", flush=True)
    with client.get_queue_sender(queue_name=queue) as sender:
        for step, stage in enumerate(stages, start=1):
            for order, session_id in sessions.items():
                body = {"orderId": order, "step": step, "stage": stage}
                sender.send_messages(ServiceBusMessage(
                    json.dumps(body), session_id=session_id,
                    message_id=f"{session_id}-{step}", content_type="application/json",
                ))
                print(f"SENT {order} {step} {stage}", flush=True)

    # Deliberately select the second order first: ordering is per session.
    for order in reversed(sessions):
        session_id = sessions[order]
        with client.get_queue_receiver(
            queue_name=queue, session_id=session_id, prefetch_count=0,
            max_wait_time=20,
        ) as receiver:
            for expected_step, stage in enumerate(stages, start=1):
                messages = receiver.receive_messages(max_message_count=1, max_wait_time=20)
                if not messages:
                    raise RuntimeError(f"Missing step {expected_step} for {order}")
                message = messages[0]
                actual = json.loads(str(message))
                expected = {"orderId": order, "step": expected_step, "stage": stage}
                if actual != expected:
                    raise RuntimeError(f"Order check failed: expected {expected}, received {actual}")
                # Checking the body represents the work in this lab. Complete it
                # before requesting another message; no concurrent handlers here.
                receiver.complete_message(message)
                print(f"PROCESSED {order} {expected_step} {stage}", flush=True)
            print(f"PASS {order}: 1 -> 2 -> 3", flush=True)
    print("PASS: both sessions processed in order; six messages completed.", flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--namespace", required=True, help="Namespace name, without the DNS suffix")
    parser.add_argument("--queue", default="orders-sessions")
    args = parser.parse_args()
    with AzureCliCredential() as credential:
        with ServiceBusClient(
            fully_qualified_namespace=f"{args.namespace}.servicebus.windows.net",
            credential=credential, transport_type=TransportType.AmqpOverWebsocket,
        ) as client:
            run(client, args.queue)


if __name__ == "__main__":
    main()