# Checkpointing with Azure Event Hubs

Patterns for reliable event processing with checkpoint stores.

## Why Checkpointing?

Checkpointing tracks which events have been processed, enabling:
- **Resume after failure** — Pick up where you left off
- **Scalable consumers** — Multiple consumers share work without duplication
- **At-least-once delivery** — Ensure no events are lost

## Blob Checkpoint Store (Recommended)

```python
from azure.eventhub import EventHubConsumerClient
from azure.eventhub.extensions.checkpointstoreblob import BlobCheckpointStore
from azure.identity import DefaultAzureCredential

credential = DefaultAzureCredential()

# Create checkpoint store
checkpoint_store = BlobCheckpointStore(
    blob_account_url="https://<account>.blob.core.windows.net",
    container_name="checkpoints",
    credential=credential
)

# Consumer with checkpoint store
consumer = EventHubConsumerClient(
    fully_qualified_namespace="<namespace>.servicebus.windows.net",
    eventhub_name="my-eventhub",
    consumer_group="$Default",
    credential=credential,
    checkpoint_store=checkpoint_store
)
```

## Async Blob Checkpoint Store

```python
from azure.eventhub.aio import EventHubConsumerClient
from azure.eventhub.extensions.checkpointstoreblob.aio import BlobCheckpointStore
from azure.identity.aio import DefaultAzureCredential

async def create_consumer():
    credential = DefaultAzureCredential()
    
    checkpoint_store = BlobCheckpointStore(
        blob_account_url="https://<account>.blob.core.windows.net",
        container_name="checkpoints",
        credential=credential
    )
    
    consumer = EventHubConsumerClient(
        fully_qualified_namespace="<namespace>.servicebus.windows.net",
        eventhub_name="my-eventhub",
        consumer_group="$Default",
        credential=credential,
        checkpoint_store=checkpoint_store
    )
    
    return consumer
```

## Checkpoint Strategies

### After Every Event (Most Reliable)

```python
async def on_event(partition_context, event):
    """Checkpoint after every event - highest reliability, highest overhead."""
    try:
        await process_event(event)
        await partition_context.update_checkpoint(event)
    except Exception as e:
        # Don't checkpoint on failure - event will be reprocessed
        print(f"Processing failed: {e}")
```

### Batch Checkpointing (Balanced)

```python
class BatchCheckpointer:
    def __init__(self, batch_size: int = 100):
        self.batch_size = batch_size
        self.counts: dict[str, int] = {}
    
    async def on_event(self, partition_context, event):
        partition_id = partition_context.partition_id
        
        await process_event(event)
        
        self.counts[partition_id] = self.counts.get(partition_id, 0) + 1
        
        if self.counts[partition_id] >= self.batch_size:
            await partition_context.update_checkpoint(event)
            self.counts[partition_id] = 0
            print(f"Checkpointed partition {partition_id}")
```

### Time-Based Checkpointing

```python
import asyncio
from datetime import datetime, timedelta

class TimedCheckpointer:
    def __init__(self, interval_seconds: float = 30.0):
        self.interval = interval_seconds
        self.last_checkpoint: dict[str, datetime] = {}
        self.last_event: dict[str, any] = {}
    
    async def on_event(self, partition_context, event):
        partition_id = partition_context.partition_id
        now = datetime.utcnow()
        
        await process_event(event)
        self.last_event[partition_id] = event
        
        last = self.last_checkpoint.get(partition_id)
        if not last or (now - last).total_seconds() >= self.interval:
            await partition_context.update_checkpoint(event)
            self.last_checkpoint[partition_id] = now
            print(f"Timed checkpoint for partition {partition_id}")
```

### Hybrid Checkpointing (Batch + Time)

```python
class HybridCheckpointer:
    def __init__(self, batch_size: int = 100, interval_seconds: float = 30.0):
        self.batch_size = batch_size
        self.interval = interval_seconds
        self.counts: dict[str, int] = {}
        self.last_checkpoint: dict[str, datetime] = {}
    
    async def on_event(self, partition_context, event):
        partition_id = partition_context.partition_id
        now = datetime.utcnow()
        
        await process_event(event)
        
        self.counts[partition_id] = self.counts.get(partition_id, 0) + 1
        last = self.last_checkpoint.get(partition_id)
        
        should_checkpoint = (
            self.counts[partition_id] >= self.batch_size or
            (last and (now - last).total_seconds() >= self.interval)
        )
        
        if should_checkpoint:
            await partition_context.update_checkpoint(event)
            self.counts[partition_id] = 0
            self.last_checkpoint[partition_id] = now
```

## Checkpoint on Batch Complete

```python
async def on_event_batch(partition_context, events):
    """Process batch and checkpoint once at the end."""
    if not events:
        return
    
    for event in events:
        await process_event(event)
    
    # Checkpoint only the last event
    await partition_context.update_checkpoint(events[-1])
    print(f"Processed {len(events)} events, checkpointed partition {partition_context.partition_id}")

async with consumer:
    await consumer.receive_batch(
        on_event_batch=on_event_batch,
        max_batch_size=100,
        max_wait_time=5.0
    )
```

## Manual Checkpoint Management

```python
from azure.eventhub.extensions.checkpointstoreblob import BlobCheckpointStore

async def inspect_checkpoints(checkpoint_store: BlobCheckpointStore):
    """List all checkpoints for debugging."""
    checkpoints = await checkpoint_store.list_checkpoints(
        fully_qualified_namespace="<namespace>.servicebus.windows.net",
        eventhub_name="my-eventhub",
        consumer_group="$Default"
    )
    
    for cp in checkpoints:
        print(f"Partition: {cp['partition_id']}")
        print(f"  Offset: {cp['offset']}")
        print(f"  Sequence: {cp['sequence_number']}")
```

## Checkpoint Data Structure

Blob checkpoint stores data in this format:

```
Container: checkpoints
└── <namespace>/<eventhub>/<consumer-group>/checkpoint/
    ├── 0  (partition 0 checkpoint)
    ├── 1  (partition 1 checkpoint)
    └── 2  (partition 2 checkpoint)
```

Each checkpoint blob contains:
```json
{
    "offset": "12345",
    "sequence_number": 100
}
```

## Graceful Shutdown with Final Checkpoint

```python
import signal
import asyncio

class GracefulConsumer:
    def __init__(self, consumer):
        self.consumer = consumer
        self.running = True
        self.last_events: dict[str, any] = {}
        self.partition_contexts: dict[str, any] = {}
    
    async def on_event(self, partition_context, event):
        if not self.running:
            return
        
        await process_event(event)
        
        # Track for final checkpoint
        self.last_events[partition_context.partition_id] = event
        self.partition_contexts[partition_context.partition_id] = partition_context
    
    async def shutdown(self):
        """Checkpoint all partitions before stopping."""
        self.running = False
        
        for partition_id, event in self.last_events.items():
            context = self.partition_contexts[partition_id]
            await context.update_checkpoint(event)
            print(f"Final checkpoint for partition {partition_id}")
        
        await self.consumer.close()

# Usage
consumer = GracefulConsumer(client)

def signal_handler(sig, frame):
    asyncio.create_task(consumer.shutdown())

signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler)
```

## Error Handling and Retry

```python
async def on_event_with_retry(partition_context, event, max_retries: int = 3):
    """Process with retries, only checkpoint on success."""
    for attempt in range(max_retries):
        try:
            await process_event(event)
            await partition_context.update_checkpoint(event)
            return
        except TransientError as e:
            if attempt < max_retries - 1:
                await asyncio.sleep(2 ** attempt)  # Exponential backoff
                continue
            raise
        except PermanentError as e:
            # Log and checkpoint to skip - event cannot be processed
            print(f"Permanent failure, skipping: {e}")
            await partition_context.update_checkpoint(event)
            return
```

## Checkpoint Store Configuration

| Option | Description | Default |
|--------|-------------|---------|
| `blob_account_url` | Storage account URL | Required |
| `container_name` | Blob container for checkpoints | Required |
| `credential` | Auth credential | Required |
| `api_version` | Blob storage API version | Latest |

## Best Practices

1. **Use dedicated container** — Separate checkpoint data from application data
2. **Match consumer groups** — Each consumer group needs its own checkpoints
3. **Balance frequency** — Too frequent = overhead, too rare = reprocessing on failure
4. **Handle idempotency** — Design for at-least-once delivery (events may replay)
5. **Monitor checkpoint lag** — Track time between event enqueue and checkpoint
6. **Test failure scenarios** — Verify resume works correctly after crashes
7. **Use async store** — The async BlobCheckpointStore is more performant
8. **Checkpoint on success only** — Don't checkpoint failed events

## Checkpoint Frequency Guidelines

| Scenario | Strategy | Trade-off |
|----------|----------|-----------|
| Critical data, low volume | Every event | Max reliability, higher latency |
| Normal processing | Every 100 events | Good balance |
| High throughput | Every 30 seconds | Lower overhead, more reprocessing |
| Batch processing | End of batch | Natural fit for batch workloads |

# Partition Management with Azure Event Hubs

Patterns for working with partitions, load balancing, and event ordering.

## Understanding Partitions

Partitions enable:
- **Parallel processing** — Multiple consumers process different partitions
- **Ordered delivery** — Events within a partition maintain order
- **Scalability** — More partitions = higher throughput

## Get Partition Information

```python
from azure.eventhub import EventHubProducerClient
from azure.identity import DefaultAzureCredential

producer = EventHubProducerClient(
    fully_qualified_namespace="<namespace>.servicebus.windows.net",
    eventhub_name="my-eventhub",
    credential=DefaultAzureCredential()
)

with producer:
    # Get Event Hub properties
    eh_props = producer.get_eventhub_properties()
    print(f"Event Hub: {eh_props['name']}")
    print(f"Partitions: {eh_props['partition_ids']}")
    print(f"Created: {eh_props['created_at']}")
    
    # Get individual partition properties
    for partition_id in eh_props['partition_ids']:
        props = producer.get_partition_properties(partition_id)
        print(f"\nPartition {partition_id}:")
        print(f"  First sequence: {props['beginning_sequence_number']}")
        print(f"  Last sequence: {props['last_enqueued_sequence_number']}")
        print(f"  Last offset: {props['last_enqueued_offset']}")
        print(f"  Last enqueued: {props['last_enqueued_time_utc']}")
        print(f"  Is empty: {props['is_empty']}")
```

## Sending to Specific Partitions

### By Partition ID (Direct)

```python
# Send to specific partition
batch = producer.create_batch(partition_id="0")
batch.add(EventData("Goes to partition 0"))
producer.send_batch(batch)
```

### By Partition Key (Consistent Hashing)

```python
# Same key always goes to same partition
batch = producer.create_batch(partition_key="user-123")
batch.add(EventData("All user-123 events in same partition"))
producer.send_batch(batch)

# Different keys may go to different partitions
for user_id in ["user-1", "user-2", "user-3"]:
    batch = producer.create_batch(partition_key=user_id)
    batch.add(EventData(f"Event for {user_id}"))
    producer.send_batch(batch)
```

### Round-Robin (Default)

```python
# No partition_id or partition_key = round-robin distribution
batch = producer.create_batch()
batch.add(EventData("Distributed across partitions"))
producer.send_batch(batch)
```

## Partition Selection Strategies

| Strategy | Use Case | Code |
|----------|----------|------|
| Round-robin | Even distribution, no ordering needs | `create_batch()` |
| Partition key | Related events in same partition | `create_batch(partition_key="...")` |
| Explicit partition | Direct control | `create_batch(partition_id="0")` |

## Receiving from Partitions

### All Partitions (Load Balanced)

```python
from azure.eventhub import EventHubConsumerClient
from azure.eventhub.extensions.checkpointstoreblob import BlobCheckpointStore

checkpoint_store = BlobCheckpointStore(...)

consumer = EventHubConsumerClient(
    fully_qualified_namespace="<namespace>.servicebus.windows.net",
    eventhub_name="my-eventhub",
    consumer_group="$Default",
    credential=DefaultAzureCredential(),
    checkpoint_store=checkpoint_store  # Required for load balancing
)

async def on_event(partition_context, event):
    print(f"Partition {partition_context.partition_id}: {event.body_as_str()}")
    await partition_context.update_checkpoint(event)

async with consumer:
    # Automatically distributes partitions across consumers
    await consumer.receive(on_event=on_event)
```

### Specific Partition

```python
async with consumer:
    await consumer.receive(
        on_event=on_event,
        partition_id="0"  # Only receive from partition 0
    )
```

## Load Balancing Multiple Consumers

When multiple consumers share a consumer group with a checkpoint store, partitions are automatically distributed.

### Example: 3 Consumers, 8 Partitions

```
Consumer 1: Partitions 0, 1, 2
Consumer 2: Partitions 3, 4, 5  
Consumer 3: Partitions 6, 7
```

### Partition Ownership Events

```python
async def on_partition_initialize(partition_context):
    """Called when consumer claims a partition."""
    print(f"Initialized partition {partition_context.partition_id}")

async def on_partition_close(partition_context, reason):
    """Called when consumer releases a partition."""
    print(f"Closed partition {partition_context.partition_id}: {reason}")

async with consumer:
    await consumer.receive(
        on_event=on_event,
        on_partition_initialize=on_partition_initialize,
        on_partition_close=on_partition_close
    )
```

### Close Reasons

| Reason | Description |
|--------|-------------|
| `SHUTDOWN` | Consumer is closing normally |
| `OWNERSHIP_LOST` | Another consumer claimed this partition |

## Partition Context Properties

```python
async def on_event(partition_context, event):
    # Partition info
    print(f"Partition ID: {partition_context.partition_id}")
    print(f"Consumer Group: {partition_context.consumer_group}")
    print(f"Event Hub: {partition_context.eventhub_name}")
    print(f"Namespace: {partition_context.fully_qualified_namespace}")
    
    # Last enqueued event (if tracking enabled)
    if partition_context.last_enqueued_event_properties:
        props = partition_context.last_enqueued_event_properties
        print(f"Last sequence: {props.sequence_number}")
        print(f"Last offset: {props.offset}")
        print(f"Last enqueued: {props.enqueued_time}")
```

## Track Partition Health

```python
async def on_event(partition_context, event):
    props = partition_context.last_enqueued_event_properties
    
    if props:
        # Calculate lag
        current_seq = event.sequence_number
        last_seq = props.sequence_number
        lag = last_seq - current_seq
        
        if lag > 1000:
            print(f"WARNING: Partition {partition_context.partition_id} lag: {lag}")
```

## Parallel Partition Processing

```python
import asyncio
from azure.eventhub.aio import EventHubConsumerClient

async def process_partition(consumer, partition_id: str):
    """Process a single partition."""
    async def on_event(partition_context, event):
        await process_event(event)
        await partition_context.update_checkpoint(event)
    
    await consumer.receive(
        on_event=on_event,
        partition_id=partition_id
    )

async def parallel_consume():
    consumer = EventHubConsumerClient(...)
    
    async with consumer:
        # Get all partitions
        props = await consumer.get_eventhub_properties()
        
        # Start parallel tasks for each partition
        tasks = [
            process_partition(consumer, pid)
            for pid in props['partition_ids']
        ]
        
        await asyncio.gather(*tasks)
```

## Partition-Aware Batching

```python
class PartitionBatcher:
    """Batch events by partition for efficient sending."""
    
    def __init__(self, producer):
        self.producer = producer
        self.batches: dict[str, any] = {}
    
    async def add(self, event: EventData, partition_key: str):
        if partition_key not in self.batches:
            self.batches[partition_key] = await self.producer.create_batch(
                partition_key=partition_key
            )
        
        try:
            self.batches[partition_key].add(event)
        except ValueError:
            # Batch full, send and create new
            await self.producer.send_batch(self.batches[partition_key])
            self.batches[partition_key] = await self.producer.create_batch(
                partition_key=partition_key
            )
            self.batches[partition_key].add(event)
    
    async def flush(self):
        for batch in self.batches.values():
            if batch:
                await self.producer.send_batch(batch)
        self.batches.clear()
```

## Starting Positions

```python
from azure.eventhub import EventHubConsumerClient

# From beginning
await consumer.receive(
    on_event=on_event,
    starting_position="-1"  # Beginning of stream
)

# From end (new events only)
await consumer.receive(
    on_event=on_event,
    starting_position="@latest"
)

# From specific offset
await consumer.receive(
    on_event=on_event,
    starting_position="12345"  # Specific offset
)

# From specific time
from datetime import datetime, timezone
start_time = datetime(2024, 1, 1, tzinfo=timezone.utc)
await consumer.receive(
    on_event=on_event,
    starting_position=start_time
)

# Different positions per partition
await consumer.receive(
    on_event=on_event,
    starting_position={
        "0": "-1",      # Partition 0 from beginning
        "1": "@latest", # Partition 1 from end
        "2": "5000"     # Partition 2 from offset 5000
    }
)
```

## Monitoring Partition Distribution

```python
class PartitionMonitor:
    def __init__(self):
        self.owned_partitions: set[str] = set()
        self.event_counts: dict[str, int] = {}
    
    async def on_partition_initialize(self, partition_context):
        self.owned_partitions.add(partition_context.partition_id)
        print(f"Now own partitions: {self.owned_partitions}")
    
    async def on_partition_close(self, partition_context, reason):
        self.owned_partitions.discard(partition_context.partition_id)
        print(f"Released {partition_context.partition_id}: {reason}")
    
    async def on_event(self, partition_context, event):
        pid = partition_context.partition_id
        self.event_counts[pid] = self.event_counts.get(pid, 0) + 1
        
        # Log distribution every 1000 events
        total = sum(self.event_counts.values())
        if total % 1000 == 0:
            print(f"Event distribution: {self.event_counts}")
```

## Best Practices

1. **Use partition keys** for related events that need ordering
2. **Avoid explicit partition IDs** unless you have a specific reason
3. **Scale consumers** with partitions — aim for 1 consumer per partition max
4. **Monitor lag** — track difference between enqueued and processed events
5. **Handle ownership changes** — design for partitions moving between consumers
6. **Use checkpoint stores** for automatic load balancing
7. **Start from checkpoint** — let checkpoint store manage starting position
8. **Partition count is fixed** — plan capacity upfront (cannot change after creation)

## Partition Limits

| Limit | Value |
|-------|-------|
| Min partitions | 1 |
| Max partitions (Standard) | 32 |
| Max partitions (Premium/Dedicated) | 100+ |
| Max throughput per partition | ~1 MB/s or ~1000 events/s |

#!/usr/bin/env python3
"""
CLI tool for Azure Event Hubs consumer setup and monitoring.

Usage:
    # Show Event Hub info
    python setup_consumer.py info --namespace mynamespace --eventhub myeventhub
    
    # Show partition details
    python setup_consumer.py partitions --namespace mynamespace --eventhub myeventhub
    
    # Receive events (simple)
    python setup_consumer.py receive --namespace mynamespace --eventhub myeventhub
    
    # Receive with checkpointing
    python setup_consumer.py receive --namespace mynamespace --eventhub myeventhub \
        --storage-account mystorageaccount --checkpoint-container checkpoints
    
    # Receive from specific partition
    python setup_consumer.py receive --namespace mynamespace --eventhub myeventhub \
        --partition 0 --starting-position earliest
    
    # Send test events
    python setup_consumer.py send --namespace mynamespace --eventhub myeventhub \
        --message "Hello World" --count 10

Environment Variables:
    EVENT_HUB_FULLY_QUALIFIED_NAMESPACE: <namespace>.servicebus.windows.net
    EVENT_HUB_NAME: Event Hub name
    STORAGE_ACCOUNT_URL: https://<account>.blob.core.windows.net
    CHECKPOINT_CONTAINER: Checkpoint container name
"""

import argparse
import asyncio
import json
import os
import sys
from datetime import datetime, timezone
from typing import Optional

from azure.eventhub import EventData
from azure.eventhub.aio import EventHubConsumerClient, EventHubProducerClient
from azure.identity.aio import DefaultAzureCredential


async def get_eventhub_info(namespace: str, eventhub: str):
    """Display Event Hub information."""
    credential = DefaultAzureCredential()

    async with EventHubProducerClient(
        fully_qualified_namespace=namespace,
        eventhub_name=eventhub,
        credential=credential,
    ) as producer:
        props = await producer.get_eventhub_properties()

        print(f"Event Hub: {props['name']}")
        print(f"Created: {props['created_at']}")
        print(
            f"Partitions: {len(props['partition_ids'])} ({', '.join(props['partition_ids'])})"
        )


async def get_partition_info(namespace: str, eventhub: str):
    """Display detailed partition information."""
    credential = DefaultAzureCredential()

    async with EventHubProducerClient(
        fully_qualified_namespace=namespace,
        eventhub_name=eventhub,
        credential=credential,
    ) as producer:
        props = await producer.get_eventhub_properties()

        print(f"Event Hub: {props['name']}")
        print(f"Total Partitions: {len(props['partition_ids'])}")
        print("-" * 60)

        total_events = 0
        for partition_id in props["partition_ids"]:
            p_props = await producer.get_partition_properties(partition_id)

            begin_seq = p_props["beginning_sequence_number"]
            last_seq = p_props["last_enqueued_sequence_number"]
            event_count = last_seq - begin_seq if not p_props["is_empty"] else 0
            total_events += event_count

            print(f"\nPartition {partition_id}:")
            print(f"  Empty: {p_props['is_empty']}")
            print(f"  Sequence Range: {begin_seq} - {last_seq}")
            print(f"  Event Count (approx): {event_count}")
            print(f"  Last Offset: {p_props['last_enqueued_offset']}")
            print(f"  Last Enqueued: {p_props['last_enqueued_time_utc']}")

        print("-" * 60)
        print(f"Total Events (approx): {total_events}")


async def receive_events(
    namespace: str,
    eventhub: str,
    consumer_group: str = "$Default",
    partition_id: Optional[str] = None,
    starting_position: str = "latest",
    storage_account: Optional[str] = None,
    checkpoint_container: Optional[str] = None,
    max_events: int = 100,
    max_wait_time: float = 30.0,
):
    """Receive events from Event Hub."""
    credential = DefaultAzureCredential()
    checkpoint_store = None

    # Setup checkpoint store if provided
    if storage_account and checkpoint_container:
        from azure.eventhub.extensions.checkpointstoreblob.aio import (
            BlobCheckpointStore,
        )

        storage_url = f"https://{storage_account}.blob.core.windows.net"
        checkpoint_store = BlobCheckpointStore(
            blob_account_url=storage_url,
            container_name=checkpoint_container,
            credential=credential,
        )
        print(f"Using checkpoint store: {storage_url}/{checkpoint_container}")

    # Parse starting position
    if starting_position == "earliest":
        start_pos = "-1"
    elif starting_position == "latest":
        start_pos = "@latest"
    else:
        start_pos = starting_position

    event_count = 0

    async def on_event(partition_context, event):
        nonlocal event_count

        if event:
            event_count += 1
            print(
                f"\n[Partition {partition_context.partition_id}] Event {event_count}:"
            )
            print(f"  Sequence: {event.sequence_number}")
            print(f"  Offset: {event.offset}")
            print(f"  Enqueued: {event.enqueued_time}")

            body = event.body_as_str()
            if len(body) > 200:
                body = body[:200] + "..."
            print(f"  Body: {body}")

            if event.properties:
                print(f"  Properties: {event.properties}")

            # Checkpoint if store available
            if checkpoint_store:
                await partition_context.update_checkpoint(event)

        if event_count >= max_events:
            raise StopIteration("Max events reached")

    async def on_error(partition_context, error):
        if partition_context:
            print(f"Error in partition {partition_context.partition_id}: {error}")
        else:
            print(f"Error: {error}")

    consumer = EventHubConsumerClient(
        fully_qualified_namespace=namespace,
        eventhub_name=eventhub,
        consumer_group=consumer_group,
        credential=credential,
        checkpoint_store=checkpoint_store,
    )

    print(f"Receiving from Event Hub: {eventhub}")
    print(f"Consumer Group: {consumer_group}")
    print(f"Starting Position: {starting_position}")
    if partition_id:
        print(f"Partition: {partition_id}")
    print(f"Max Events: {max_events}")
    print("-" * 60)

    try:
        async with consumer:
            if partition_id:
                await consumer.receive(
                    on_event=on_event,
                    on_error=on_error,
                    partition_id=partition_id,
                    starting_position=start_pos,
                    max_wait_time=max_wait_time,
                )
            else:
                await consumer.receive(
                    on_event=on_event,
                    on_error=on_error,
                    starting_position=start_pos,
                    max_wait_time=max_wait_time,
                )
    except StopIteration:
        pass
    except KeyboardInterrupt:
        print("\nStopped by user")

    print(f"\n-" * 60)
    print(f"Total events received: {event_count}")


async def send_events(
    namespace: str,
    eventhub: str,
    message: str,
    count: int = 1,
    partition_key: Optional[str] = None,
    partition_id: Optional[str] = None,
):
    """Send test events to Event Hub."""
    credential = DefaultAzureCredential()

    async with EventHubProducerClient(
        fully_qualified_namespace=namespace,
        eventhub_name=eventhub,
        credential=credential,
    ) as producer:
        # Create batch with optional partition targeting
        batch_kwargs = {}
        if partition_id:
            batch_kwargs["partition_id"] = partition_id
            print(f"Sending to partition: {partition_id}")
        elif partition_key:
            batch_kwargs["partition_key"] = partition_key
            print(f"Using partition key: {partition_key}")

        batch = await producer.create_batch(**batch_kwargs)

        sent = 0
        for i in range(count):
            event_body = f"{message} #{i + 1}" if count > 1 else message
            event = EventData(event_body)
            event.properties = {
                "index": i,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

            try:
                batch.add(event)
            except ValueError:
                # Batch full, send and create new
                await producer.send_batch(batch)
                sent += batch.size_in_bytes
                batch = await producer.create_batch(**batch_kwargs)
                batch.add(event)

        # Send remaining
        if batch:
            await producer.send_batch(batch)

        print(f"Sent {count} event(s) to {eventhub}")


def main():
    parser = argparse.ArgumentParser(
        description="Azure Event Hubs consumer setup and monitoring",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Common arguments
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument(
        "--namespace",
        "-n",
        default=os.environ.get("EVENT_HUB_FULLY_QUALIFIED_NAMESPACE"),
        help="Event Hub namespace (e.g., mynamespace.servicebus.windows.net)",
    )
    common.add_argument(
        "--eventhub",
        "-e",
        default=os.environ.get("EVENT_HUB_NAME"),
        help="Event Hub name",
    )

    # Info command
    info_parser = subparsers.add_parser(
        "info", parents=[common], help="Show Event Hub info"
    )

    # Partitions command
    partitions_parser = subparsers.add_parser(
        "partitions", parents=[common], help="Show partition details"
    )

    # Receive command
    receive_parser = subparsers.add_parser(
        "receive", parents=[common], help="Receive events"
    )
    receive_parser.add_argument(
        "--consumer-group",
        "-g",
        default="$Default",
        help="Consumer group (default: $Default)",
    )
    receive_parser.add_argument(
        "--partition", "-p", help="Specific partition to receive from"
    )
    receive_parser.add_argument(
        "--starting-position",
        choices=["earliest", "latest"],
        default="latest",
        help="Starting position (default: latest)",
    )
    receive_parser.add_argument(
        "--storage-account",
        default=os.environ.get("STORAGE_ACCOUNT_URL", "")
        .replace("https://", "")
        .replace(".blob.core.windows.net", ""),
        help="Storage account for checkpointing",
    )
    receive_parser.add_argument(
        "--checkpoint-container",
        default=os.environ.get("CHECKPOINT_CONTAINER"),
        help="Container for checkpoints",
    )
    receive_parser.add_argument(
        "--max-events",
        type=int,
        default=100,
        help="Maximum events to receive (default: 100)",
    )
    receive_parser.add_argument(
        "--max-wait-time",
        type=float,
        default=30.0,
        help="Max wait time in seconds (default: 30)",
    )

    # Send command
    send_parser = subparsers.add_parser(
        "send", parents=[common], help="Send test events"
    )
    send_parser.add_argument(
        "--message", "-m", default="Test event", help="Message to send"
    )
    send_parser.add_argument(
        "--count",
        "-c",
        type=int,
        default=1,
        help="Number of events to send (default: 1)",
    )
    send_parser.add_argument(
        "--partition-key", help="Partition key for consistent routing"
    )
    send_parser.add_argument("--partition-id", help="Specific partition ID to send to")

    args = parser.parse_args()

    # Validate common args
    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Add .servicebus.windows.net if not present
    namespace = args.namespace
    if namespace and not namespace.endswith(".servicebus.windows.net"):
        namespace = f"{namespace}.servicebus.windows.net"

    if not namespace or not args.eventhub:
        print("Error: --namespace and --eventhub are required")
        sys.exit(1)

    # Run command
    try:
        if args.command == "info":
            asyncio.run(get_eventhub_info(namespace, args.eventhub))

        elif args.command == "partitions":
            asyncio.run(get_partition_info(namespace, args.eventhub))

        elif args.command == "receive":
            asyncio.run(
                receive_events(
                    namespace=namespace,
                    eventhub=args.eventhub,
                    consumer_group=args.consumer_group,
                    partition_id=args.partition,
                    starting_position=args.starting_position,
                    storage_account=args.storage_account
                    if args.storage_account
                    else None,
                    checkpoint_container=args.checkpoint_container,
                    max_events=args.max_events,
                    max_wait_time=args.max_wait_time,
                )
            )

        elif args.command == "send":
            asyncio.run(
                send_events(
                    namespace=namespace,
                    eventhub=args.eventhub,
                    message=args.message,
                    count=args.count,
                    partition_key=args.partition_key,
                    partition_id=args.partition_id,
                )
            )

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()