// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.messaging.eventhubs;

import com.azure.core.amqp.exception.AmqpException;
import com.azure.core.amqp.implementation.AmqpReceiveLink;
import com.azure.core.amqp.implementation.MessageSerializer;
import com.azure.core.amqp.implementation.RequestResponseChannelClosedException;
import com.azure.core.amqp.implementation.RetryUtil;
import com.azure.core.amqp.implementation.StringUtil;
import com.azure.core.annotation.ReturnType;
import com.azure.core.annotation.ServiceClient;
import com.azure.core.annotation.ServiceMethod;
import com.azure.core.util.logging.ClientLogger;
import com.azure.messaging.eventhubs.implementation.AmqpReceiveLinkProcessor;
import com.azure.messaging.eventhubs.implementation.EventHubConnectionProcessor;
import com.azure.messaging.eventhubs.implementation.EventHubManagementNode;
import com.azure.messaging.eventhubs.implementation.instrumentation.EventHubsConsumerInstrumentation;
import com.azure.messaging.eventhubs.models.EventPosition;
import com.azure.messaging.eventhubs.models.PartitionEvent;
import com.azure.messaging.eventhubs.models.ReceiveOptions;
import reactor.core.publisher.BaseSubscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;

import java.io.Closeable;
import java.util.Locale;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

import static com.azure.core.util.FluxUtil.fluxError;
import static com.azure.core.util.FluxUtil.monoError;
import static com.azure.messaging.eventhubs.implementation.ClientConstants.CONNECTION_ID_KEY;
import static com.azure.messaging.eventhubs.implementation.ClientConstants.LINK_NAME_KEY;
import static com.azure.messaging.eventhubs.implementation.ClientConstants.PARTITION_ID_KEY;
import static com.azure.messaging.eventhubs.implementation.ClientConstants.SIGNAL_TYPE_KEY;

/**
 * An <b>asynchronous</b> consumer responsible for reading {@link EventData} from either a specific Event Hub partition
 * or all partitions in the context of a specific consumer group.
 *
 * <p>The examples shown in this document use a credential object named DefaultAzureCredential for authentication,
 * which is appropriate for most scenarios, including local development and production environments. Additionally, we
 * recommend using
 * <a href="https://learn.microsoft.com/azure/active-directory/managed-identities-azure-resources/">managed identity</a>
 * for authentication in production environments. You can find more information on different ways of authenticating and
 * their corresponding credential types in the
 * <a href="https://learn.microsoft.com/java/api/overview/azure/identity-readme">Azure Identity documentation"</a>.
 * </p>
 *
 * <p><strong>Sample: Creating an {@link EventHubConsumerAsyncClient}</strong></p>
 *
 * <p>The following code sample demonstrates the creation of the asynchronous client
 * {@link EventHubConsumerAsyncClient}.  The {@code fullyQualifiedNamespace} is the Event Hubs Namespace's host name.
 * It is listed under the "Essentials" panel after navigating to the Event Hubs Namespace via Azure Portal. The
 * {@code consumerGroup} is found by navigating to the Event Hub instance, and selecting "Consumer groups" under the
 * "Entities" panel.  The {@link EventHubClientBuilder#consumerGroup(String)} is required for creating consumer clients.
 * The credential used is {@code DefaultAzureCredential} because it combines commonly used credentials in deployment
 * and development and chooses the credential to used based on its running environment.</p>
 *
 * <!-- src_embed com.azure.messaging.eventhubs.eventhubproducerasyncclient.construct -->
 * <pre>
 * TokenCredential credential = new DefaultAzureCredentialBuilder&#40;&#41;.build&#40;&#41;;
 *
 * &#47;&#47; &quot;&lt;&lt;fully-qualified-namespace&gt;&gt;&quot; will look similar to &quot;&#123;your-namespace&#125;.servicebus.windows.net&quot;
 * &#47;&#47; &quot;&lt;&lt;event-hub-name&gt;&gt;&quot; will be the name of the Event Hub instance you created inside the Event Hubs namespace.
 * EventHubProducerAsyncClient producer = new EventHubClientBuilder&#40;&#41;
 *     .credential&#40;&quot;&lt;&lt;fully-qualified-namespace&gt;&gt;&quot;, &quot;&lt;&lt;event-hub-name&gt;&gt;&quot;,
 *         credential&#41;
 *     .buildAsyncProducerClient&#40;&#41;;
 * </pre>
 * <!-- end com.azure.messaging.eventhubs.eventhubproducerasyncclient.construct -->
 *
 * <p><strong>Sample: Consuming events a single partition from Event Hub</strong></p>
 *
 * <p>The code sample below demonstrates receiving events from partition "0" of an Event Hub starting from
 * {@link EventPosition#latest()}.  {@link EventPosition#latest()} points to the end of the partition stream.  The
 * consumer receives events enqueued <i>after</i> it started subscribing for events.</p>
 *
 * <p>{@link #receiveFromPartition(String, EventPosition)} is a non-blocking call.  After setting up the operation,
 * its async representation is returned. The {@code Flux<PartitionEvent>} must be subscribed to, like the sample below,
 * to start receiving events.</p>
 *
 * <!-- src_embed com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#string-eventposition -->
 * <pre>
 * EventHubConsumerAsyncClient consumer = new EventHubClientBuilder&#40;&#41;
 *     .credential&#40;&quot;&lt;&lt;fully-qualified-namespace&gt;&gt;&quot;, &quot;&lt;&lt;event-hub-name&gt;&gt;&quot;,
 *         new DefaultAzureCredentialBuilder&#40;&#41;.build&#40;&#41;&#41;
 *     .consumerGroup&#40;EventHubClientBuilder.DEFAULT_CONSUMER_GROUP_NAME&#41;
 *     .buildAsyncConsumerClient&#40;&#41;;
 *
 * &#47;&#47; Obtain partitionId from EventHubConsumerAsyncClient.getPartitionIds&#40;&#41;
 * String partitionId = &quot;0&quot;;
 * EventPosition startingPosition = EventPosition.latest&#40;&#41;;
 *
 * &#47;&#47; Keep a reference to `subscription`. When the program is finished receiving events, call
 * &#47;&#47; subscription.dispose&#40;&#41;. This will stop fetching events from the Event Hub.
 * &#47;&#47;
 * &#47;&#47; NOTE: This is a non-blocking call and will move to the next line of code after setting up the async
 * &#47;&#47; operation.  If the program ends after this, or the class is immediately disposed, no events will be
 * &#47;&#47; received.
 * Disposable subscription = consumer.receiveFromPartition&#40;partitionId, startingPosition&#41;
 *     .subscribe&#40;partitionEvent -&gt; &#123;
 *         PartitionContext partitionContext = partitionEvent.getPartitionContext&#40;&#41;;
 *         EventData event = partitionEvent.getData&#40;&#41;;
 *
 *         System.out.printf&#40;&quot;Received event from partition '%s'%n&quot;, partitionContext.getPartitionId&#40;&#41;&#41;;
 *         System.out.printf&#40;&quot;Contents of event as string: '%s'%n&quot;, event.getBodyAsString&#40;&#41;&#41;;
 *     &#125;, error -&gt; &#123;
 *         &#47;&#47; This is a terminal signal.  No more events will be received from the same Flux object.
 *         System.err.print&#40;&quot;An error occurred:&quot; + error&#41;;
 *     &#125;, &#40;&#41; -&gt; &#123;
 *         &#47;&#47; This is a terminal signal.  No more events will be received from the same Flux object.
 *         System.out.print&#40;&quot;Stream has ended.&quot;&#41;;
 *     &#125;&#41;;
 * </pre>
 * <!-- end com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#string-eventposition -->
 *
 * <p><strong>Sample: Including latest partition information in received events</strong></p>
 *
 * <p>{@link EventData} can be decorated with the latest partition information and sent to consumers. Enable this by
 * setting {@link ReceiveOptions#setTrackLastEnqueuedEventProperties(boolean)} to {@code true}. As events come in,
 * explore the {@link PartitionEvent} object.  This is useful in scenarios where customers want to constant up-to-date
 * information about their Event Hub. This does take a performance hit as the extra partition information must be sent
 * over the wire with every event.</p>
 *
 * <p>{@link #receiveFromPartition(String, EventPosition, ReceiveOptions)} is a non-blocking call.  After setting up the
 * operation, its async representation is returned. The {@code Flux<PartitionEvent>} must be subscribed to, like
 * sample below, to start receiving events.</p>
 *
 * <!-- src_embed com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receiveFromPartition#string-eventposition-receiveoptions -->
 * <pre>
 * &#47;&#47; Set `setTrackLastEnqueuedEventProperties` to true to get the last enqueued information from the partition for
 * &#47;&#47; each event that is received.
 * ReceiveOptions receiveOptions = new ReceiveOptions&#40;&#41;
 *     .setTrackLastEnqueuedEventProperties&#40;true&#41;;
 * EventPosition startingPosition = EventPosition.earliest&#40;&#41;;
 *
 * &#47;&#47; Receives events from partition &quot;0&quot; starting at the beginning of the stream.
 * &#47;&#47; Keep a reference to `subscription`. When the program is finished receiving events, call
 * &#47;&#47; subscription.dispose&#40;&#41;. This will stop fetching events from the Event Hub.
 * Disposable subscription = consumer.receiveFromPartition&#40;&quot;0&quot;, startingPosition, receiveOptions&#41;
 *     .subscribe&#40;partitionEvent -&gt; &#123;
 *         LastEnqueuedEventProperties properties = partitionEvent.getLastEnqueuedEventProperties&#40;&#41;;
 *         System.out.printf&#40;&quot;Information received at %s. Last enqueued sequence number: %s%n&quot;,
 *             properties.getRetrievalTime&#40;&#41;,
 *             properties.getSequenceNumber&#40;&#41;&#41;;
 *     &#125;&#41;;
 * </pre>
 * <!-- end com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receiveFromPartition#string-eventposition-receiveoptions -->
 *
 * <p><strong>Sample: Rate limiting consumption of events from Event Hub</strong></p>
 *
 * <p>For event consumers that need to limit the number of events they receive at a given time, they can use
 * {@link BaseSubscriber#request(long)}.  Using a custom subscriber allows developers more granular control over the
 * rate at which they receive events.</p>
 *
 * <p>{@link #receiveFromPartition(String, EventPosition)} is a non-blocking call.  After setting up the operation,
 * its async representation is returned. The {@code Flux<PartitionEvent>} must be subscribed to, like the sample below,
 * to start receiving events.</p>
 *
 * <!-- src_embed com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#string-eventposition-basesubscriber -->
 * <pre>
 * consumer.receiveFromPartition&#40;partitionId, EventPosition.latest&#40;&#41;&#41;.subscribe&#40;new BaseSubscriber&lt;PartitionEvent&gt;&#40;&#41; &#123;
 *     private static final int NUMBER_OF_EVENTS = 5;
 *     private final AtomicInteger currentNumberOfEvents = new AtomicInteger&#40;&#41;;
 *
 *     &#64;Override
 *     protected void hookOnSubscribe&#40;Subscription subscription&#41; &#123;
 *         &#47;&#47; Tell the Publisher we only want 5 events at a time.
 *         request&#40;NUMBER_OF_EVENTS&#41;;
 *     &#125;
 *
 *     &#64;Override
 *     protected void hookOnNext&#40;PartitionEvent value&#41; &#123;
 *         &#47;&#47; Process the EventData
 *
 *         &#47;&#47; If the number of events we have currently received is a multiple of 5, that means we have reached the
 *         &#47;&#47; last event the Publisher will provide to us. Invoking request&#40;long&#41; here, tells the Publisher that
 *         &#47;&#47; the subscriber is ready to get more events from upstream.
 *         if &#40;currentNumberOfEvents.incrementAndGet&#40;&#41; % 5 == 0&#41; &#123;
 *             request&#40;NUMBER_OF_EVENTS&#41;;
 *         &#125;
 *     &#125;
 * &#125;&#41;;
 * </pre>
 * <!-- end com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#string-eventposition-basesubscriber -->
 *
 * <p><strong>Sample: Receiving from all partitions</strong></p>
 *
 * <p>The code sample below demonstrates receiving events from all partitions of an Event Hub starting the beginning of
 * each partition's stream.  This is valuable for demo purposes but <i>is not</i> intended for production scenarios.
 * For production scenarios, consider using {@link EventProcessorClient}.</p>
 *
 * <p>{@link #receive(boolean)} is a non-blocking call.  After setting up the operation, its async representation is
 * returned. The {@code Flux<PartitionEvent>} must be subscribed to, like the sample below, to start receiving events.
 * </p>
 *
 * <!-- src_embed com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#boolean -->
 * <pre>
 * &#47;&#47; Keep a reference to `subscription`. When the program is finished receiving events, call
 * &#47;&#47; subscription.dispose&#40;&#41;. This will stop fetching events from the Event Hub.
 * Disposable subscription = consumer.receive&#40;true&#41;
 *     .subscribe&#40;partitionEvent -&gt; &#123;
 *         PartitionContext context = partitionEvent.getPartitionContext&#40;&#41;;
 *         EventData event = partitionEvent.getData&#40;&#41;;
 *
 *         System.out.printf&#40;&quot;Event %s is from partition %s%n.&quot;, event.getSequenceNumber&#40;&#41;,
 *             context.getPartitionId&#40;&#41;&#41;;
 *     &#125;, error -&gt; &#123;
 *         &#47;&#47; This is a terminal signal.  No more events will be received from the same Flux object.
 *         System.err.print&#40;&quot;An error occurred:&quot; + error&#41;;
 *     &#125;, &#40;&#41; -&gt; &#123;
 *         &#47;&#47; This is a terminal signal.  No more events will be received from the same Flux object.
 *         System.out.print&#40;&quot;Stream has ended.&quot;&#41;;
 *     &#125;&#41;;
 * </pre>
 * <!-- end com.azure.messaging.eventhubs.eventhubconsumerasyncclient.receive#boolean -->
 *
 * @see com.azure.messaging.eventhubs
 * @see EventHubClientBuilder
 */
@ServiceClient(builder = EventHubClientBuilder.class, isAsync = true)
public class EventHubConsumerAsyncClient implements Closeable {
    private static final String RECEIVER_ENTITY_PATH_FORMAT = "%s/ConsumerGroups/%s/Partitions/%s";
    private static final ClientLogger LOGGER = new ClientLogger(EventHubConsumerAsyncClient.class);

    private final AtomicBoolean isDisposed = new AtomicBoolean();
    private final ReceiveOptions defaultReceiveOptions = new ReceiveOptions();
    private final String fullyQualifiedNamespace;
    private final String eventHubName;
    private final EventHubConnectionProcessor connectionProcessor;
    private final MessageSerializer messageSerializer;
    private final String consumerGroup;
    private final int prefetchCount;
    private final boolean isSharedConnection;
    private final Runnable onClientClosed;
    private final String identifier;
    private final EventHubsConsumerInstrumentation instrumentation;

    /**
     * Keeps track of the open partition consumers keyed by linkName. The link name is generated as: {@code
     * "partitionId_GUID"}. For receiving from all partitions, links are prefixed with {@code "all-GUID-partitionId"}.
     */
    private final ConcurrentHashMap<String, EventHubPartitionAsyncConsumer> openPartitionConsumers =
        new ConcurrentHashMap<>();

    EventHubConsumerAsyncClient(String fullyQualifiedNamespace, String eventHubName,
        EventHubConnectionProcessor connectionProcessor, MessageSerializer messageSerializer, String consumerGroup,
        int prefetchCount, boolean isSharedConnection, Runnable onClientClosed, String identifier,
        EventHubsConsumerInstrumentation instrumentation) {
        this.fullyQualifiedNamespace = fullyQualifiedNamespace;
        this.eventHubName = eventHubName;
        this.connectionProcessor = connectionProcessor;
        this.messageSerializer = messageSerializer;
        this.consumerGroup = consumerGroup;
        this.prefetchCount = prefetchCount;
        this.isSharedConnection = isSharedConnection;
        this.onClientClosed = onClientClosed;
        this.identifier = identifier;
        this.instrumentation = instrumentation;
    }

    /**
     * Gets the fully qualified Event Hubs namespace that the connection is associated with. This is likely similar to
     * {@code {yournamespace}.servicebus.windows.net}.
     *
     * @return The fully qualified Event Hubs namespace that the connection is associated with
     */
    public String getFullyQualifiedNamespace() {
        return fullyQualifiedNamespace;
    }

    /**
     * Gets the Event Hub name this client interacts with.
     *
     * @return The Event Hub name this client interacts with.
     */
    public String getEventHubName() {
        return eventHubName;
    }

    /**
     * Gets the consumer group this consumer is reading events as a part of.
     *
     * @return The consumer group this consumer is reading events as a part of.
     */
    public String getConsumerGroup() {
        return consumerGroup;
    }

    /**
     * Retrieves information about an Event Hub, including the number of partitions present and their identifiers.
     *
     * @return The set of information for the Event Hub that this client is associated with.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Mono<EventHubProperties> getEventHubProperties() {
        return instrumentation.getTracer().traceMono(connectionProcessor.flatMap(connection -> connection.getManagementNode())
                .flatMap(EventHubManagementNode::getEventHubProperties),
            "EventHubs.getEventHubProperties");
    }

    /**
     * Retrieves the identifiers for the partitions of an Event Hub.
     *
     * @return A Flux of identifiers for the partitions of an Event Hub.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<String> getPartitionIds() {
        return getEventHubProperties().flatMapMany(properties -> Flux.fromIterable(properties.getPartitionIds()));
    }

    /**
     * Retrieves information about a specific partition for an Event Hub, including elements that describe the available
     * events in the partition event stream.
     *
     * @param partitionId The unique identifier of a partition associated with the Event Hub.
     *
     * @return The set of information for the requested partition under the Event Hub this client is associated with.
     *
     * @throws NullPointerException if {@code partitionId} is null.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Mono<PartitionProperties> getPartitionProperties(String partitionId) {
        if (Objects.isNull(partitionId)) {
            return monoError(LOGGER, new NullPointerException("'partitionId' cannot be null."));
        } else if (partitionId.isEmpty()) {
            return monoError(LOGGER, new IllegalArgumentException("'partitionId' cannot be an empty string."));
        }

        return instrumentation.getTracer().traceMono(connectionProcessor.flatMap(connection -> connection.getManagementNode())
                .flatMap(node -> node.getPartitionProperties(partitionId)),
            "EventHubs.getPartitionProperties");
    }

    /**
     * Consumes events from a single partition starting at {@code startingPosition}.
     *
     * @param partitionId Identifier of the partition to read events from.
     * @param startingPosition Position within the Event Hub partition to begin consuming events.
     *
     * @return A stream of events for this partition starting from {@code startingPosition}.
     *
     * @throws NullPointerException if {@code partitionId}, or {@code startingPosition} is null.
     * @throws IllegalArgumentException if {@code partitionId} is an empty string.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<PartitionEvent> receiveFromPartition(String partitionId, EventPosition startingPosition) {
        return receiveFromPartition(partitionId, startingPosition, defaultReceiveOptions);
    }

    /**
     * Consumes events from a single partition starting at {@code startingPosition} with a set of {@link ReceiveOptions
     * receive options}.
     *
     * <ul>
     * <li>If receive is invoked where {@link ReceiveOptions#getOwnerLevel()} has a value, then Event Hubs service will
     * guarantee only one active consumer exists per partitionId and consumer group combination. This receive operation
     * is sometimes referred to as an "Epoch Consumer".</li>
     * <li>Multiple consumers per partitionId and consumer group combination can be created by not setting
     * {@link ReceiveOptions#getOwnerLevel()} when invoking receive operations. This non-exclusive consumer is sometimes
     * referred to as a "Non-Epoch Consumer."</li>
     * </ul>
     *
     * @param partitionId Identifier of the partition to read events from.
     * @param startingPosition Position within the Event Hub partition to begin consuming events.
     * @param receiveOptions Options when receiving events from the partition.
     *
     * @return A stream of events for this partition. If a stream for the events was opened before, the same position
     *     within that partition is returned. Otherwise, events are read starting from {@code startingPosition}.
     *
     * @throws NullPointerException if {@code partitionId}, {@code startingPosition}, {@code receiveOptions} is
     *     null.
     * @throws IllegalArgumentException if {@code partitionId} is an empty string.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<PartitionEvent> receiveFromPartition(String partitionId, EventPosition startingPosition,
        ReceiveOptions receiveOptions) {
        if (Objects.isNull(partitionId)) {
            return fluxError(LOGGER, new NullPointerException("'partitionId' cannot be null."));
        } else if (partitionId.isEmpty()) {
            return fluxError(LOGGER, new IllegalArgumentException("'partitionId' cannot be an empty string."));
        }

        if (Objects.isNull(startingPosition)) {
            return fluxError(LOGGER, new NullPointerException("'startingPosition' cannot be null."));
        }

        final String linkName = StringUtil.getRandomString(partitionId);
        return createConsumer(linkName, partitionId, startingPosition, receiveOptions);
    }

    /**
     * Consumes events from all partitions starting from the beginning of each partition.
     *
     * <p>This method is <b>not</b> recommended for production use; the {@link EventProcessorClient} should be used for
     * reading events from all partitions in a production scenario, as it offers a much more robust experience with
     * higher throughput.
     *
     * It is important to note that this method does not guarantee fairness amongst the partitions. Depending on service
     * communication, there may be a clustering of events per partition and/or there may be a noticeable bias for a
     * given partition or subset of partitions.</p>
     *
     *
     * @return A stream of events for every partition in the Event Hub starting from the beginning of each partition.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<PartitionEvent> receive() {
        return receive(true, defaultReceiveOptions);
    }

    /**
     * Consumes events from all partitions.
     *
     * <p>This method is <b>not</b> recommended for production use; the {@link EventProcessorClient} should be used for
     * reading events from all partitions in a production scenario, as it offers a much more robust experience with
     * higher throughput.
     *
     * It is important to note that this method does not guarantee fairness amongst the partitions. Depending on service
     * communication, there may be a clustering of events per partition and/or there may be a noticeable bias for a
     * given partition or subset of partitions.</p>
     *
     * @param startReadingAtEarliestEvent {@code true} to begin reading at the first events available in each
     *     partition; otherwise, reading will begin at the end of each partition seeing only new events as they are
     *     published.
     *
     * @return A stream of events for every partition in the Event Hub.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<PartitionEvent> receive(boolean startReadingAtEarliestEvent) {
        return receive(startReadingAtEarliestEvent, defaultReceiveOptions);
    }

    /**
     * Consumes events from all partitions configured with a set of {@code receiveOptions}.
     *
     * <p>This method is <b>not</b> recommended for production use; the {@link EventProcessorClient} should be used for
     * reading events from all partitions in a production scenario, as it offers a much more robust experience with
     * higher throughput.
     *
     * It is important to note that this method does not guarantee fairness amongst the partitions. Depending on service
     * communication, there may be a clustering of events per partition and/or there may be a noticeable bias for a
     * given partition or subset of partitions.</p>
     *
     * <ul>
     * <li>If receive is invoked where {@link ReceiveOptions#getOwnerLevel()} has a value, then Event Hubs service will
     * guarantee only one active consumer exists per partitionId and consumer group combination. This receive operation
     * is sometimes referred to as an "Epoch Consumer".</li>
     * <li>Multiple consumers per partitionId and consumer group combination can be created by not setting
     * {@link ReceiveOptions#getOwnerLevel()} when invoking receive operations. This non-exclusive consumer is sometimes
     * referred to as a "Non-Epoch Consumer."</li>
     * </ul>
     *
     * @param startReadingAtEarliestEvent {@code true} to begin reading at the first events available in each
     *     partition; otherwise, reading will begin at the end of each partition seeing only new events as they are
     *     published.
     * @param receiveOptions Options when receiving events from each Event Hub partition.
     *
     * @return A stream of events for every partition in the Event Hub.
     *
     * @throws NullPointerException if {@code receiveOptions} is null.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public Flux<PartitionEvent> receive(boolean startReadingAtEarliestEvent, ReceiveOptions receiveOptions) {
        if (Objects.isNull(receiveOptions)) {
            return fluxError(LOGGER, new NullPointerException("'receiveOptions' cannot be null."));
        }

        final EventPosition startingPosition = startReadingAtEarliestEvent
            ? EventPosition.earliest()
            : EventPosition.latest();
        final String prefix = StringUtil.getRandomString("all");
        final Flux<PartitionEvent> allPartitionEvents = getPartitionIds().flatMap(partitionId -> {
            final String linkName = prefix + "-" + partitionId;
            return createConsumer(linkName, partitionId, startingPosition, receiveOptions);
        });

        return Flux.merge(allPartitionEvents);
    }

    /**
     * Disposes of the consumer by closing the underlying connection to the service.
     */
    @Override
    public void close() {
        if (isDisposed.getAndSet(true)) {
            return;
        }
        openPartitionConsumers.forEach((key, value) -> value.close());
        openPartitionConsumers.clear();

        if (isSharedConnection) {
            onClientClosed.run();
        } else {
            connectionProcessor.dispose();
        }
    }

    private Flux<PartitionEvent> createConsumer(String linkName, String partitionId, EventPosition startingPosition,
        ReceiveOptions receiveOptions) {
        return openPartitionConsumers
            .computeIfAbsent(linkName,
                name -> createPartitionConsumer(name, partitionId, startingPosition, receiveOptions))
            .receive()
            .doFinally(signal -> removeLink(linkName, partitionId, signal));
    }

    private void removeLink(String linkName, String partitionId, SignalType signalType) {
        LOGGER.atInfo()
            .addKeyValue(LINK_NAME_KEY, linkName)
            .addKeyValue(PARTITION_ID_KEY, partitionId)
            .addKeyValue(SIGNAL_TYPE_KEY, signalType)
            .log("Receiving completed.");

        final EventHubPartitionAsyncConsumer consumer = openPartitionConsumers.remove(linkName);

        if (consumer != null) {
            consumer.close();
        }
    }

    private EventHubPartitionAsyncConsumer createPartitionConsumer(String linkName, String partitionId,
        EventPosition startingPosition, ReceiveOptions receiveOptions) {
        final String entityPath = String.format(Locale.US, RECEIVER_ENTITY_PATH_FORMAT,
            getEventHubName(), consumerGroup, partitionId);

        final AtomicReference<Supplier<EventPosition>> initialPosition = new AtomicReference<>(() -> startingPosition);

        // The Mono, when subscribed, creates a AmqpReceiveLink in the AmqpConnection emitted by the connectionProcessor
        //
        final Mono<AmqpReceiveLink> receiveLinkMono = connectionProcessor
            .flatMap(connection -> {
                LOGGER.atInfo()
                    .addKeyValue(LINK_NAME_KEY, linkName)
                    .addKeyValue(PARTITION_ID_KEY, partitionId)
                    .addKeyValue(CONNECTION_ID_KEY, connection.getId())
                    .log("Creating receive consumer for partition.");
                return connection.createReceiveLink(linkName, entityPath, initialPosition.get().get(), receiveOptions, identifier);
            });

        // A Mono that resubscribes to 'receiveLinkMono' to retry the creation of AmqpReceiveLink.
        //
        // The scenarios where this retry helps are -
        // [1]. When we try to create a link on a session being disposed but connection is healthy, the retry can
        //      eventually create a new session then the link.
        // [2]. When we try to create a new session (to host the new link) but on a connection being disposed,
        //      the retry can eventually receive a new connection and then proceed with creating session and link.
        //
        final Mono<AmqpReceiveLink> retryableReceiveLinkMono = RetryUtil.withRetry(receiveLinkMono.onErrorMap(
                RequestResponseChannelClosedException.class,
                e -> {
                    // When the current connection is being disposed, the connectionProcessor can produce
                    // a new connection if downstream request.
                    // In this context, treat RequestResponseChannelClosedException from the RequestResponseChannel scoped
                    // to the current connection being disposed as retry-able so that retry can obtain new connection.
                    return new AmqpException(true, e.getMessage(), e, null);
                }),
            connectionProcessor.getRetryOptions(),
            "Failed to create receive link " + linkName,
            true);

        // A Flux that produces a new AmqpReceiveLink each time it receives a request from the below
        // 'AmqpReceiveLinkProcessor'. Obviously, the processor requests a link when there is a downstream subscriber.
        // It also requests a new link (i.e. retry) when the current link it holds gets terminated
        // (e.g., when the service decides to close that link).
        //
        final Flux<AmqpReceiveLink> receiveLinkFlux = retryableReceiveLinkMono
            .repeat()
            // The re-subscribe nature of 'MonoRepeat' following the emission of a link will cause the 'MonoFlatmap' (the
            // upstream of repeat operator) to cache the same link. When AmqpReceiveLinkProcessor later requests a new link,
            // the corresponding request from the 'MonoRepeat' will be answered with the cached (and closed) link by the
            // 'MonoFlatmap'. We'll filter out these cached (closed) links to avoid AmqpReceiveLinkProcessor from doing
            // unusable work (creating subscriptions and attempting to place the credit) on those links and associated logging.
            // See the PR description (https://github.com/Azure/azure-sdk-for-java/pull/33204) for more details.
            .filter(link -> !link.isDisposed());

        final AmqpReceiveLinkProcessor linkMessageProcessor = receiveLinkFlux.subscribeWith(
            new AmqpReceiveLinkProcessor(entityPath, prefetchCount, partitionId, connectionProcessor, instrumentation));

        return new EventHubPartitionAsyncConsumer(linkMessageProcessor, messageSerializer, getFullyQualifiedNamespace(),
            getEventHubName(), consumerGroup, partitionId, initialPosition,
            receiveOptions.getTrackLastEnqueuedEventProperties());
    }

    boolean isConnectionClosed() {
        return this.connectionProcessor.isChannelClosed();
    }

    EventHubsConsumerInstrumentation getInstrumentation() {
        return instrumentation;
    }

    /**
     * Gets the client identifier.
     *
     * @return The unique identifier string for current client.
     */
    public String getIdentifier() {
        return identifier;
    }
}
