"""CloudTrips Event Hubs lab: send ten events, read them, or replay a run."""
import argparse
import asyncio
import contextlib
import json
import uuid

from azure.eventhub import EventData, TransportType
from azure.eventhub.aio import EventHubConsumerClient, EventHubProducerClient
from azure.identity.aio import AzureCliCredential

COUNT = 10


async def send_events(options, credential, run_id):
    async with EventHubProducerClient(
        fully_qualified_namespace=f"{options.namespace}.servicebus.windows.net",
        eventhub_name=options.hub, credential=credential,
        transport_type=TransportType.AmqpOverWebsocket,
    ) as producer:
        batch = await producer.create_batch(partition_key="sensor-01")
        for number in range(1, COUNT + 1):
            batch.add(EventData(json.dumps({
                "runId": run_id, "deviceId": "sensor-01",
                "reading": number, "temperature": 20 + number / 10,
            })))
        await producer.send_batch(batch)
    print(f"SENT {COUNT} events; run ID: {run_id}", flush=True)


async def read_events(options, credential, run_id):
    seen = set()
    finished = asyncio.Event()
    errors = []
    scanned = 0
    ignored = 0

    async def on_event(partition, event):
        nonlocal scanned, ignored
        if event is None:
            print(f"WAIT partition={partition.partition_id}; "
                  f"scanned={scanned}; matched={len(seen)}/{COUNT}", flush=True)
            return
        scanned += 1
        try:
            body = json.loads(event.body_as_str())
        except (ValueError, UnicodeDecodeError):
            return
        if not isinstance(body, dict) or body.get("runId") != run_id:
            ignored += 1
            if ignored <= 3:
                print(f"SKIP event from another run on partition={partition.partition_id}", flush=True)
            return
        number = body.get("reading")
        if (type(number) is not int or number not in range(1, COUNT + 1)
                or body.get("deviceId") != "sensor-01"
                or body.get("temperature") != 20 + number / 10):
            errors.append(RuntimeError(f"Unexpected event for this run: {body}"))
            finished.set()
            return
        if number not in seen:
            seen.add(number)
            print(f"READ partition={partition.partition_id} "
                  f"sequence={event.sequence_number} reading={number} "
                  f"temperature={body['temperature']}", flush=True)
        if len(seen) == COUNT:
            finished.set()

    async def on_partition_initialize(partition):
        print(f"START receiver partition={partition.partition_id}; "
              f"consumer group={options.consumer_group}; position=earliest", flush=True)

    async def on_error(partition, error):
        print(f"RECEIVE ERROR: {type(error).__name__}: {error}", flush=True)
        errors.append(error)
        finished.set()

    async with EventHubConsumerClient(
        fully_qualified_namespace=f"{options.namespace}.servicebus.windows.net",
        eventhub_name=options.hub, consumer_group=options.consumer_group,
        credential=credential, transport_type=TransportType.AmqpOverWebsocket,
    ) as consumer:
        print(f"READING {options.namespace}/{options.hub}; "
              f"consumer group={options.consumer_group}; run={run_id}", flush=True)
        try:
            partition_ids = await asyncio.wait_for(consumer.get_partition_ids(), timeout=30)
            for partition_id in partition_ids:
                properties = await asyncio.wait_for(
                    consumer.get_partition_properties(partition_id), timeout=30,
                )
                print(f"PARTITION {partition_id}: empty={properties['is_empty']}; "
                      f"last sequence={properties['last_enqueued_sequence_number']}", flush=True)
        except asyncio.TimeoutError as error:
            raise RuntimeError("Timed out reading partition metadata. Check network access and login.") from error
        receiving = asyncio.create_task(consumer.receive(
            on_event=on_event, on_error=on_error, starting_position="-1",
            on_partition_initialize=on_partition_initialize, max_wait_time=10,
        ))
        waiting = asyncio.create_task(finished.wait())
        try:
            done, _ = await asyncio.wait(
                {receiving, waiting}, timeout=90,
                return_when=asyncio.FIRST_COMPLETED,
            )
            if errors:
                raise RuntimeError(f"Receive failed: {errors[0]}")
            if receiving in done:
                await receiving  # Surface any connection/authentication error.
            if len(seen) != COUNT:
                raise RuntimeError(
                    f"Read {len(seen)}/{COUNT} matching events; scanned {scanned} total events. "
                    f"Consumer group: {options.consumer_group}. "
                    "If scanned=0, verify the consumer group exists and your identity has "
                    "receive access. "
                    "If scanned>0, check the run ID and retention."
                )
        finally:
            for task in (receiving, waiting):
                task.cancel()
            for task in (receiving, waiting):
                with contextlib.suppress(asyncio.CancelledError):
                    await task
    print(f"PASS: read all {COUNT} events for run {run_id}", flush=True)


async def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--namespace", required=True)
    parser.add_argument("--hub", default="telemetry")
    parser.add_argument("--consumer-group", default="$Default")
    parser.add_argument("--mode", choices=("both", "send", "read"), default="both")
    parser.add_argument("--run-id")
    options = parser.parse_args()
    if options.mode == "read" and not options.run_id:
        parser.error("--run-id is required for --mode read")
    run_id = options.run_id or uuid.uuid4().hex
    print(f"Run ID: {run_id}", flush=True)
    async with AzureCliCredential() as credential:
        if options.mode in ("both", "send"):
            await send_events(options, credential, run_id)
        if options.mode in ("both", "read"):
            await read_events(options, credential, run_id)


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