/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.pulsar.broker.service.persistent;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.bookkeeper.mledger.Entry;
import org.apache.bookkeeper.mledger.ManagedCursor;
import org.apache.bookkeeper.mledger.Position;
import org.apache.pulsar.broker.service.BrokerServiceException;
import org.apache.pulsar.broker.service.Consumer;
import org.apache.pulsar.broker.service.EntryBatchSizes;
import org.apache.pulsar.broker.service.HashRangeStickyKeyConsumerSelector;
import org.apache.pulsar.broker.service.SendMessageInfo;
import org.apache.pulsar.broker.service.StickyKeyConsumerSelector;
import org.apache.pulsar.broker.service.Subscription;
import org.apache.pulsar.common.api.proto.PulsarApi.CommandSubscribe.SubType;
import org.apache.pulsar.common.util.Murmur3_32Hash;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDispatcherMultipleConsumers {

    private final StickyKeyConsumerSelector selector;

    PersistentStickyKeyDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) {
        super(topic, cursor, subscription);
        //TODO: Consumer selector Pluggable
        selector = new HashRangeStickyKeyConsumerSelector();
    }

    @Override
    public synchronized void addConsumer(Consumer consumer) throws BrokerServiceException {
        super.addConsumer(consumer);
        selector.addConsumer(consumer);
    }

    @Override
    public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceException {
        super.removeConsumer(consumer);
        selector.removeConsumer(consumer);
    }

    @Override
    protected void sendMessagesToConsumers(ReadType readType, List<Entry> entries) {
        long totalMessagesSent = 0;
        long totalBytesSent = 0;
        if (entries.size() > 0) {
            final Map<Integer, List<Entry>> groupedEntries = new HashMap<>();
            for (Entry entry : entries) {
                int key = Murmur3_32Hash.getInstance().makeHash(peekStickyKey(entry.getDataBuffer()));
                groupedEntries.putIfAbsent(key, new ArrayList<>());
                groupedEntries.get(key).add(entry);
            }
            final Iterator<Map.Entry<Integer, List<Entry>>> iterator = groupedEntries.entrySet().iterator();
            AtomicInteger keyNumbers = new AtomicInteger(groupedEntries.size());
            while (iterator.hasNext() && totalAvailablePermits > 0 && isAtleastOneConsumerAvailable()) {
                final Map.Entry<Integer, List<Entry>> entriesWithSameKey = iterator.next();
                //TODO: None key policy
                Consumer consumer = null;
                if (selector instanceof HashRangeStickyKeyConsumerSelector) {
                    consumer = ((HashRangeStickyKeyConsumerSelector)selector).select(entriesWithSameKey.getKey());
                }
                if (consumer == null) {
                    // Do nothing, cursor will be rewind at reconnection
                    log.info("[{}] rewind because no available consumer found for key {} from total {}", name,
                            entriesWithSameKey.getKey(), consumerList.size());
                    entriesWithSameKey.getValue().forEach(Entry::release);
                    cursor.rewind();
                    return;
                }

                int messagesForC = Math.min(entriesWithSameKey.getValue().size(), consumer.getAvailablePermits());
                if (log.isDebugEnabled()) {
                    log.debug("[{}] select consumer {} for key {} with messages num {}, read type is {}",
                            name, consumer.consumerName(), entriesWithSameKey.getKey(), messagesForC, readType);
                }
                if (messagesForC > 0) {
                    // remove positions first from replay list first : sendMessages recycles entries
                    List<Entry> subList = new ArrayList<>(entriesWithSameKey.getValue().subList(0, messagesForC));
                    if (readType == ReadType.Replay) {
                        subList.forEach(entry -> messagesToRedeliver.remove(entry.getLedgerId(), entry.getEntryId()));
                    }

                    SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal();
                    EntryBatchSizes batchSizes = EntryBatchSizes.get(subList.size());
                    filterEntriesForConsumer(subList, batchSizes, sendMessageInfo);

                    consumer.sendMessages(subList, batchSizes, sendMessageInfo.getTotalMessages(),
                            sendMessageInfo.getTotalBytes(), getRedeliveryTracker()).addListener(future -> {
                                if (future.isSuccess() && keyNumbers.decrementAndGet() == 0) {
                                    readMoreEntries();
                                }
                    });
                    entriesWithSameKey.getValue().removeAll(subList);

                    totalAvailablePermits -= sendMessageInfo.getTotalMessages();
                    totalMessagesSent += sendMessageInfo.getTotalMessages();
                    totalBytesSent += sendMessageInfo.getTotalBytes();

                    if (entriesWithSameKey.getValue().size() == 0) {
                        iterator.remove();
                    }
                }
            }

            // acquire message-dispatch permits for already delivered messages
            if (serviceConfig.isDispatchThrottlingOnNonBacklogConsumerEnabled() || !cursor.isActive()) {
                if (topic.getDispatchRateLimiter().isPresent()) {
                    topic.getDispatchRateLimiter().get().tryDispatchPermit(totalMessagesSent, totalBytesSent);
                }

                if (dispatchRateLimiter.isPresent()) {
                    dispatchRateLimiter.get().tryDispatchPermit(totalMessagesSent, totalBytesSent);
                }
            }

            if (groupedEntries.size() > 0) {
                int laterReplay = 0;
                for (List<Entry> entryList : groupedEntries.values()) {
                    laterReplay += entryList.size();
                    entryList.forEach(entry -> {
                        messagesToRedeliver.add(entry.getLedgerId(), entry.getEntryId());
                        entry.release();
                    });
                }
                if (log.isDebugEnabled()) {
                    log.debug("[{}] No consumers found with available permits, storing {} positions for later replay", name,
                            laterReplay);
                }
            }
        }
    }

    @Override
    public SubType getType() {
        return SubType.Key_Shared;
    }

    @Override
    protected Set<? extends Position> asyncReplayEntries(Set<? extends Position> positions) {
        return cursor.asyncReplayEntries(positions, this, ReadType.Replay, true);
    }

    private static final Logger log = LoggerFactory.getLogger(PersistentStickyKeyDispatcherMultipleConsumers.class);

}
