"""Ansible-rulebook event source plugin for Azure Event Hub.

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

import asyncio
import json
import logging
from typing import Any

from azure.eventhub.aio import (EventHubConsumerClient, PartitionContext)
from azure.eventhub import EventData
from azure.identity.aio import ClientSecretCredential

DOCUMENTATION = r"""
---
short_description: Receive events via a Azure Event Hub
description:
  - An ansible-rulebook event source plugin for receiving events
    via Azure Event Hub
options:
  azure_tenant_id:
    description:
      - The azure tenant id
    type: str
    required: true
  azure_client_id:
    description:
      - The azure client id
    type: str
    required: true
  azure_client_secret:
    description:
      - The azure client secret
    type: str
    required: true
  azure_namespace:
    description:
      - The azure event hub namespace which includes the host name
    type: str
    example: "test.servicebus.windows.net"
    required: true
  azure_event_hub_name:
    description:
      - The azure event hub name
    type: str
    required: true
  azure_starting_position:
    description:
      - The starting position
    type: str
    default: "-1"
  azure_consumer_group:
    description:
      - The name of the consumer group
    type: str
    default: "$Default"
"""

EXAMPLES = r"""
- azure.azcollection.azure_event_hub:
    "azure_tenant_id": "your_tenant_id"
    "azure_client_id": "your_client_id"
    "azure_client_secret": "your_client_secret"
    "azure_namespace": "example.servicebus.windows.net"
    "azure_event_hub_name": "your_hub_name"
    "azure_starting_position": "-1"
"""

logger = logging.getLogger()


REQUIRED_ARGS = [
    "azure_tenant_id",
    "azure_client_id",
    "azure_client_secret",
    "azure_namespace",
    "azure_event_hub_name",
]


class AzureHubConsumer:
    """Azure Hub Consumer."""

    def __init__(self, queue: asyncio.Queue[Any], args: dict[str, Any]) -> None:
        """Initialize Hub Consumer."""
        self.queue = queue
        tenant_id = args.get("azure_tenant_id")
        client_id = args.get("azure_client_id")
        client_secret = args.get("azure_client_secret")

        self.event_hub_namespace = args.get("azure_namespace")
        self.event_hub_name = args.get("azure_event_hub_name")

        self.consumer_group = args.get("azure_consumer_group", "$Default")
        self.starting_position = int(args.get("azure_starting_position", "-1"))

        self.credential = ClientSecretCredential(
            tenant_id=tenant_id,
            client_id=client_id,
            client_secret=client_secret,
        )

    async def on_event(self, partition_context: PartitionContext, event: EventData) -> None:
        """Receiving event data."""
        if event:
            await partition_context.update_checkpoint(event)
            meta = {}
            # Process message body
            try:
                value = event.body_as_str()
                logger.debug(
                    "Received event from partition %s: %s",
                    str(partition_context.partition_id),
                    value,
                )
            except UnicodeError:
                logger.exception("Unicode error while decoding message body")
                data = None
            else:
                try:
                    data = json.loads(value)
                except json.decoder.JSONDecodeError:
                    logger.info("JSON decode error, storing raw value")
                    data = value

            # Add data to the event and put it into the queue
            if data:
                await self.queue.put({"body": data, "meta": meta})

        await asyncio.sleep(0)

    async def start_receiving(self) -> None:
        """Start receiving data."""
        client = EventHubConsumerClient(
            fully_qualified_namespace=self.event_hub_namespace,
            eventhub_name=self.event_hub_name,
            consumer_group=self.consumer_group,
            credential=self.credential,
        )
        async with client:
            await client.receive(self.on_event)


# Usage
async def main(  # pylint: disable=R0914
    queue: asyncio.Queue[Any],
    args: dict[str, Any],
) -> None:
    """Entry Point."""
    for key in REQUIRED_ARGS:
        if key not in args:
            msg = f"Please provide {key} it is a required argument."
            raise ValueError(msg)

    consumer = AzureHubConsumer(queue, args)
    await consumer.start_receiving()


if __name__ == "__main__":

    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

    test_args = {
        "azure_tenant_id": "your_tenant_id",
        "azure_client_id": "your_client_id",
        "azure_client_secret": "your_client_secret",
        "azure_namespace": "example.servicebus.windows.net",
        "azure_event_hub_name": "your_hub_name",
        "azure_starting_position": "-1",
    }

    asyncio.run(
        main(
            MockQueue(),
            test_args,
        ),
    )
