"""Ansible-rulebook event source plugin for Azure Service Bus.

This module provides asynchronous consumers for Azure Service Bus to feed
events into Ansible EDA rulebooks.
"""

import asyncio
import contextlib
import json
from typing import Any

from azure.servicebus.aio import ServiceBusClient

DOCUMENTATION = r"""
---
short_description: Receive events from an Azure service bus.
description:
  - An ansible-rulebook event source module for receiving events from an Azure service bus.
  - In order to get the service bus and the connection string, refer to
    https://learn.microsoft.com/en-us/azure/service-bus-messaging/service-bus-python-how-to-use-queues?tabs=passwordless
options:
  conn_str:
    description:
      - The connection string to connect to the Azure service bus.
    type: str
    required: true
  queue_name:
    description:
      - The name of the queue to pull messages from.
    type: str
    required: true
  logging_enable:
    description:
      - Whether to turn on logging.
    type: bool
    default: true
"""

EXAMPLES = r"""
- azure.azcollection.azure_service_bus:
    conn_str: "{{connection_str}}"
    queue_name: "{{queue_name}}"
"""


async def receive_events(
    queue: asyncio.Queue[Any],
    args: dict[str, Any],  # pylint: disable=W0621
) -> None:
    """Receive events from service bus asynchronously."""
    servicebus_client = ServiceBusClient.from_connection_string(
        conn_str=args["conn_str"],
        logging_enable=bool(args.get("logging_enable", True)),
    )

    async with servicebus_client:
        receiver = servicebus_client.get_queue_receiver(queue_name=args["queue_name"])
        async with receiver:
            async for msg in receiver:
                meta = {"message_id": msg.message_id}
                body = str(msg)
                with contextlib.suppress(json.JSONDecodeError):
                    body = json.loads(body)

                await queue.put({"body": body, "meta": meta})
                await receiver.complete_message(msg)


async def main(
    queue: asyncio.Queue[Any],
    args: dict[str, Any],  # pylint: disable=W0621
) -> None:
    """Receive events from service bus in a loop."""
    await receive_events(queue, args)


if __name__ == "__main__":
    # MockQueue if running directly.

    class MockQueue(asyncio.Queue[Any]):
        """A fake queue."""

        async def put(self: "MockQueue", event: dict[str, Any]) -> None:
            """Print the event."""
            print(event)  # noqa: T201

    args = {
        "conn_str": "Endpoint=sb://foo.servicebus.windows.net/",
        "queue_name": "eda-queue",
    }
    asyncio.run(main(MockQueue(), args))
