/*
 * 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.rocketmq.broker.lite;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.tuple.Triple;
import org.apache.rocketmq.broker.BrokerController;
import org.apache.rocketmq.common.BrokerConfig;
import org.apache.rocketmq.common.ServiceThread;
import org.apache.rocketmq.common.constant.LoggerName;
import org.apache.rocketmq.common.entity.ClientGroup;
import org.apache.rocketmq.common.lite.LiteSubscription;
import org.apache.rocketmq.common.lite.LiteUtil;
import org.apache.rocketmq.logging.org.slf4j.Logger;
import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;

public class LiteEventDispatcher extends ServiceThread {

    private static final Logger LOGGER = LoggerFactory.getLogger(LoggerName.ROCKETMQ_POP_LITE_LOGGER_NAME);
    private static final Object PRESENT = new Object();
    private static final long CLIENT_INACTIVE_INTERVAL = 10 * 1000; // inactive time when it has unprocessed events
    protected static final long CLIENT_LONG_POLLING_INTERVAL = 30 * 1000 + 5000; // at least a period of long polling as 30s
    protected static final long ACTIVE_CONSUMING_WINDOW = 5000;
    protected static final double LOW_WATER_MARK = 0.2;
    private static final int BLACKLIST_EXPIRE_SECONDS = 10;
    private static final int SCAN_LOG_INTERVAL = 10000;

    private final BrokerController brokerController;
    private final LiteSubscriptionRegistry liteSubscriptionRegistry;
    private final AbstractLiteLifecycleManager liteLifecycleManager;

    protected final ConcurrentMap<String, ClientEventSet> clientEventMap = new ConcurrentHashMap<>();
    protected final ConcurrentSkipListSet<FullDispatchRequest> fullDispatchSet = new ConcurrentSkipListSet<>(COMPARATOR);
    protected final ConcurrentMap<String, Object> fullDispatchMap = new ConcurrentHashMap<>(); // deduplication
    private final Cache<String, Object> blacklist =
        CacheBuilder.newBuilder().expireAfterWrite(BLACKLIST_EXPIRE_SECONDS, TimeUnit.SECONDS).build();
    private final Random random = ThreadLocalRandom.current();
    private long lastLogTime = System.currentTimeMillis();

    public LiteEventDispatcher(BrokerController brokerController,
        LiteSubscriptionRegistry liteSubscriptionRegistry, AbstractLiteLifecycleManager liteLifecycleManager) {
        this.brokerController = brokerController;
        this.liteSubscriptionRegistry = liteSubscriptionRegistry;
        this.liteLifecycleManager = liteLifecycleManager;
    }

    public void init() {
        this.liteSubscriptionRegistry.addListener(new LiteCtlListenerImpl());
    }

    /**
     * Try to dispatch event to one client when message arriving or available.
     * In most cases, there is only one subscriber for a LMQ under a consumer group,
     * but also supports multiple clients consuming in share mode.
     * When group is null, dispatch to all subscribers regardless of their group,
     * when group is specified, only dispatch to subscribers belonging to this group.
     */
    public void dispatch(String group, String lmqName, int queueId, long offset, long msgStoreTime) {
        if (queueId != 0 || !LiteUtil.isLiteTopicQueue(lmqName)) {
            return;
        }
        doDispatch(group, lmqName, null);
    }

    protected void doDispatch(String group, String lmqName, String excludeClientId) {
        SubscriberWrapper wrapper = liteSubscriptionRegistry.getAllSubscriber(group, lmqName);
        if (null == wrapper) {
            return;
        }
        if (wrapper instanceof SubscriberWrapper.ListWrapper) {
            selectAndDispatch(lmqName, wrapper.asListWrapper().getClients(), excludeClientId);
        }
        if (wrapper instanceof SubscriberWrapper.MapWrapper) {
            Map<String, List<ClientGroup>> map = wrapper.asMapWrapper().getGroupMap();
            map.forEach((key, value) -> selectAndDispatch(lmqName, value, excludeClientId));
        }
    }

    /**
     * Select an appropriate client from the client list and try to dispatch the event to it.
     * If there's only one client, dispatch directly to it.
     * If there are multiple clients, randomly select one and consider fallback options
     * Try to avoid dispatching to the excluded one but fallback if no other choice.
     *
     * @param clients         all clients of one group
     * @param excludeClientId the client ID to exclude from selection, probably consuming blocked.
     * @return true if dispatched to one client
     */
    @VisibleForTesting
    public boolean selectAndDispatch(String lmqName, List<ClientGroup> clients, String excludeClientId) {
        if (CollectionUtils.isEmpty(clients)) {
            return true;
        }

        String group = clients.get(0).group;
        boolean isWildcardGroup = LiteMetadataUtil.isWildcardGroup(group, brokerController);
        String selectedClient = null; // the selected one
        int start = random.nextInt(clients.size());
        List<ClientGroup> fallbackList = new ArrayList<>(clients.size());
        for (int i = 0; i < clients.size(); i++) {
            int index = (start + i) % clients.size();
            if (clients.get(index).clientId.equals(excludeClientId)) {
                fallbackList.add(clients.get(index));
                continue;
            }
            if (blacklist.getIfPresent(clients.get(index).clientId) != null) {
                if (!isWildcardGroup) { // prevent iterating twice for large client set
                    fallbackList.add(clients.get(index));
                }
                continue;
            }
            if (tryDispatchToClient(lmqName, clients.get(index).clientId, group, !isWildcardGroup)) {
                selectedClient = clients.get(index).clientId;
                break;
            }
        }
        if (null == selectedClient) {
            for (ClientGroup clientGroup : fallbackList) {
                if (tryDispatchToClient(lmqName, clientGroup.clientId, group, !isWildcardGroup)) {
                    selectedClient = clientGroup.clientId;
                    break;
                }
            }
        }
        if (selectedClient != null) {
            notifyMessageArriving(selectedClient, group);
        } else if (isWildcardGroup) { // no one available in this group, so schedule a full dispatch once
            scheduleFullDispatchForWildcardGroup(group,
                brokerController.getBrokerConfig().getLiteEventFullDispatchDelayTimeForWildcardGroup());
        }
        return selectedClient != null;
    }

    /**
     * Try to dispatch an event to a selected client by adding it to the client's event queue.
     * If the event queue is full, mark a full dispatch for retry later.
     * @param scheduleFullDispatchIfFull schedule full dispatch if full, only false if it's a wildcard group.
     */
    @VisibleForTesting
    public boolean tryDispatchToClient(String lmqName, String clientId, String group, boolean scheduleFullDispatchIfFull) {
        ClientEventSet eventSet = clientEventMap.computeIfAbsent(clientId, key -> new ClientEventSet(group));
        if (eventSet.offer(lmqName)) {
            return true;
        }
        if (scheduleFullDispatchIfFull) {
            long delayTime = brokerController.getBrokerConfig().getLiteEventFullDispatchDelayTime()
                + (blacklist.getIfPresent(clientId) != null ? random.nextInt(15 * 1000) : 0);
            scheduleFullDispatchForClient(clientId, group, delayTime);
        }
        LOGGER.warn("client event set is full. {}", clientId);
        return false;
    }

    /**
     * Get an iterator for iterating over events for a specific client.
     * Returns events from the client's event queue.
     */
    public Iterator<String> getEventIterator(String clientId) {
        return new EventSetIterator(clientEventMap.get(clientId));
    }

    /**
     * Perform a full dispatch for a client which was previously marked for a delayed full dispatch.
     * This always happens when a client's event queue is full or re-dispatching is needed.
     * It iterates through all LMQ topics subscribed by the client and dispatches events for those
     * with available messages.
     */
    public void doFullDispatchForClient(String clientId, String group) {
        LiteSubscription subscription = liteSubscriptionRegistry.getLiteSubscription(clientId);
        if (null == subscription || CollectionUtils.isEmpty(subscription.getLiteTopicSet())) {
            LOGGER.info("client full dispatch, but no subscription. {}", clientId);
            return;
        }
        ClientEventSet eventSet = clientEventMap.computeIfAbsent(clientId, key -> new ClientEventSet(group));
        if (eventSet.maybeBlock()) {
            LOGGER.warn("client may block for a while, wait another period. {}", clientId);
            scheduleFullDispatchForClient(clientId, group,
                brokerController.getBrokerConfig().getLiteEventFullDispatchDelayTime() + random.nextInt(15 * 1000));
            return;
        }
        boolean isActiveConsuming = eventSet.isActiveConsuming();
        if (!eventSet.isLowWaterMark()) {
            LOGGER.warn("client event set high water mark, wait another period. {}, {}", clientId, isActiveConsuming);
            scheduleFullDispatchForClient(clientId, group, brokerController.getBrokerConfig().getLiteEventFullDispatchDelayTime()
                + (isActiveConsuming ? 0 : random.nextInt(10 * 1000)));
            return;
        }
        LOGGER.info("client full dispatch, {}, total:{}", clientId, subscription.getLiteTopicSet().size());
        int count = 0;
        for (String lmqName : subscription.getLiteTopicSet()) {
            long maxOffset = liteLifecycleManager.getMaxOffsetInQueue(lmqName);
            if (maxOffset <= 0) {
                continue;
            }
            long consumerOffset = brokerController.getConsumerOffsetManager().queryOffset(group, lmqName, 0);
            if (consumerOffset >= maxOffset) {
                continue;
            }
            if (eventSet.offer(lmqName)) {
                if (count++ % 10 == 0) {
                    notifyMessageArriving(clientId, group);
                }
            } else {
                LOGGER.warn("client event set full again, wait another period. {}, {}", clientId, isActiveConsuming);
                scheduleFullDispatchForClient(clientId, group, brokerController.getBrokerConfig().getLiteEventFullDispatchDelayTime()
                    + (isActiveConsuming ? 0 : random.nextInt(15 * 1000)));
                break;
            }
        }
        notifyMessageArriving(clientId, group);
        LOGGER.info("client full dispatch finish. {}, dispatch:{}", clientId, count);
    }

    private void notifyMessageArriving(String clientId, String group) {
        brokerController.getPopLiteMessageProcessor()
            .getPopLiteLongPollingService()
            .notifyMessageArriving(clientId, true, 0, group);
        brokerController.getNotificationProcessor()
            .getPopLiteLongPollingService()
            .notifyMessageArriving(clientId, true, 0, group);
    }

    /**
     * Check whether a client has any events in the event queue.
     *
     * @param clientId the client ID to check
     * @return true if the client has events, false otherwise
     */
    public boolean hasEvents(String clientId) {
        ClientEventSet eventSet = clientEventMap.get(clientId);
        return eventSet != null && eventSet.size() > 0;
    }

    /**
     * Perform a full dispatch for wildcard group which was previously marked for a delayed full dispatch.
     * It iterates through all LMQ topics in CQ table, so it may be a heavy work.
     */
    public void doFullDispatchForWildcardGroup(String group) {
        String parentTopic = LiteMetadataUtil.getLiteBindTopic(group, brokerController);
        if (null == parentTopic || !LiteMetadataUtil.isWildcardGroup(group, brokerController)) {
            return;
        }
        List<ClientGroup> clients =  liteSubscriptionRegistry.getWildcardSubscriber(group, parentTopic).getClients();
        if (CollectionUtils.isEmpty(clients)) {
            return;
        }
        AtomicInteger count = new AtomicInteger();
        Function<Triple<String, Long, Long>, Boolean> function = triple -> {
            String lmqName = triple.getLeft();
            long maxOffset = triple.getMiddle();
            if (!LiteUtil.belongsTo(lmqName, parentTopic)) {
                return true;
            }
            if (maxOffset <= 0) {
                return true;
            }
            long consumerOffset = brokerController.getConsumerOffsetManager().queryOffset(group, lmqName, 0);
            if (consumerOffset >= maxOffset) {
                return true;
            }
            if (selectAndDispatch(lmqName, clients, null)) {
                count.incrementAndGet();
            } else {
                LOGGER.warn("doFullDispatchForWildcardGroup, wait another period. {}", group);
                return false;
            }
            return true;
        };
        liteLifecycleManager.forEachLiteTopic(function);
        LOGGER.info("doFullDispatchForWildcardGroup finish. {}, dispatch:{}", group, count);
    }

    /**
     * Perform a full dispatch for all clients under a specific group, only invoked by admin for now.
     */
    public void doFullDispatchByGroup(String group) {
        if (LiteMetadataUtil.isWildcardGroup(group, brokerController)) {
            doFullDispatchForWildcardGroup(group);
        } else {
            List<String> clientIds = liteSubscriptionRegistry.getAllClientIdByGroup(group);
            LOGGER.info("do full dispatch by group, {}, size:{}", group, clientIds.size());
            for (String clientId : clientIds) {
                doFullDispatchForClient(clientId, group);
            }
        }
    }

    public void scheduleFullDispatchForClient(String clientId, String group, long delayTime) {
        if (fullDispatchMap.putIfAbsent(clientId, PRESENT) != null) {
            return;
        }
        if (delayTime < 50) {
            delayTime = 50;
        }
        fullDispatchSet.add(new FullDispatchRequest(clientId, group, delayTime));
    }

    public void scheduleFullDispatchForWildcardGroup(String group, long delayTime) {
        scheduleFullDispatchForClient("$" + group + "$", group, delayTime); // $group$ as clientId, no conflict expected
    }

    /**
     * Get the last access time of a client's event set.
     *
     * @param clientId the client id
     * @return the last access time in milliseconds, or -1 if client not found
     */
    public long getClientLastAccessTime(String clientId) {
        ClientEventSet eventSet = clientEventMap.get(clientId);
        if (eventSet != null) {
            return eventSet.lastAccessTime;
        }
        return -1;
    }

    @Override
    public String getServiceName() {
        if (brokerController.getBrokerConfig().isInBrokerContainer()) {
            return brokerController.getBrokerIdentity().getIdentifier() + LiteEventDispatcher.class.getSimpleName();
        }
        return LiteEventDispatcher.class.getSimpleName();
    }

    @Override
    public void run() {
        while (!this.isStopped()) {
            long checkInterval = brokerController.getBrokerConfig().getLiteEventCheckInterval();
            this.waitForRunning(checkInterval);
            try {
                scan();
            } catch (Exception e) {
                LOGGER.error("LiteEventDispatcher-scan error.", e);
            }
        }
    }

    /**
     * Due to the event pre-allocation mechanism, it is necessary to perform
     * two main tasks to check inactive event queues and do full dispatch to reduce potential delivery latency.
     * 1. Check client event set for inactive clients and re-dispatches their events
     * 2. Process delayed full dispatch requests that are ready to be executed
     */
    public void scan() {
        boolean needLog = System.currentTimeMillis() - lastLogTime > SCAN_LOG_INTERVAL;

        // 1. check all client event set
        if (needLog) {
            LOGGER.info("Check client event set. size:{}", clientEventMap.size());
            lastLogTime = System.currentTimeMillis();
        }
        Iterator<Map.Entry<String, ClientEventSet>> iterator = clientEventMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, ClientEventSet> entry = iterator.next();
            ClientEventSet eventSet = entry.getValue();
            if (!eventSet.maybeBlock()) {
                continue;
            }
            String clientId = entry.getKey();
            LOGGER.warn("remove inactive client and re-dispatch. {}, {}", clientId, eventSet.events.size());
            iterator.remove();
            blacklist.put(clientId, PRESENT);
            String event;
            while ((event = eventSet.poll()) != null) {
                doDispatch(eventSet.group, event, clientId); // may still dispatch to current client
            }
        }

        // 2. perform full dispatch
        if (needLog) {
            LOGGER.info("Begin to trigger full dispatch. size:{}, mapSize:{}", fullDispatchSet.size(), fullDispatchMap.size());
            lastLogTime = System.currentTimeMillis();
        }
        FullDispatchRequest request;
        while ((request = fullDispatchSet.pollFirst()) != null) {
            if (request.timestamp > System.currentTimeMillis()) {
                fullDispatchSet.add(request);
                break;
            }
            fullDispatchMap.remove(request.clientId);
            if (LiteMetadataUtil.isWildcardGroup(request.group, brokerController)) {
                doFullDispatchForWildcardGroup(request.group);
            } else {
                doFullDispatchForClient(request.clientId, request.group);
            }
        }
    }

    public int getEventMapSize() {
        return clientEventMap.size();
    }

    /**
     * We use dual data structure to maintain the event queue for each client
     * and ensure event deduplication to avoid duplicate events, although it
     * has a bit more memory usage than a single concurrent set.
     */
    protected class ClientEventSet {
        private final BlockingQueue<String> events;
        private final ConcurrentMap<String, Object> map = new ConcurrentHashMap<>();
        private final String group;
        private volatile long lastAccessTime = System.currentTimeMillis();
        private volatile long lastConsumeTime = System.currentTimeMillis();
        /**
         * Cache resolved max capacity to avoid per-offer SubscriptionGroupConfig lookup + attribute
         * parsing on the hot dispatch path. Soft-cap semantics tolerate a short staleness window,
         * so refresh lazily by TTL {@link BrokerConfig#getLiteEventCapacityCacheTtlMs()}.
         */
        private volatile int maxCapacityCache;
        private volatile long capacityRefreshTime = System.currentTimeMillis();

        public ClientEventSet(String group) {
            this.group = group;
            // Use a large bounded queue as a hard ceiling; the effective capacity is enforced
            // dynamically via soft-cap in offer() so that maxClientEventCount can be changed
            // at runtime without restart.
            this.events = new LinkedBlockingQueue<>(100_000);
            this.maxCapacityCache = LiteMetadataUtil.getMaxClientEventCount(group, brokerController);
        }

        private int getMaxCapacity() {
            long now = System.currentTimeMillis();
            long ttl = brokerController.getBrokerConfig().getLiteEventCapacityCacheTtlMs();
            if (now - capacityRefreshTime > ttl) {
                maxCapacityCache = LiteMetadataUtil.getMaxClientEventCount(group, brokerController);
                capacityRefreshTime = now;
            }
            return maxCapacityCache;
        }

        // return false if and only if the queue is full, has race condition with poll(), but no side effect.
        public boolean offer(String event) {
            if (events.size() >= getMaxCapacity()) {
                return false;
            }
            boolean rst;
            if (map.putIfAbsent(event, PRESENT) == null) {
                rst = events.offer(event);
                if (!rst) {
                    map.remove(event);
                }
            } else {
                rst = true;
            }
            return rst;
        }

        public String poll() {
            lastAccessTime = System.currentTimeMillis();
            String event = events.poll();
            if (event != null) {
                map.remove(event);
                lastConsumeTime = System.currentTimeMillis();
            }
            return event;
        }

        public boolean maybeBlock() {
            long inactiveTime = System.currentTimeMillis() - lastAccessTime;
            return inactiveTime > CLIENT_LONG_POLLING_INTERVAL
                || !events.isEmpty() && inactiveTime > CLIENT_INACTIVE_INTERVAL;
        }

        public boolean isLowWaterMark() {
            int used = events.size();
            int maxCapacity = getMaxCapacity();
            return maxCapacity <= 0 || (double) used / maxCapacity < LOW_WATER_MARK;
        }

        public boolean isActiveConsuming() {
            return System.currentTimeMillis() - lastAccessTime < ACTIVE_CONSUMING_WINDOW;
        }

        public int size() {
            return events.size();
        }
    }

    class LiteCtlListenerImpl implements LiteCtlListener {

        @Override
        public void onRegister(String clientId, String group, String lmqName) {
            if (LiteMetadataUtil.isWildcardGroup(group, brokerController)) {
                scheduleFullDispatchForWildcardGroup(group, 5000);
            } else {
                if (liteLifecycleManager.isLmqExist(lmqName)) {
                    doDispatch(group, lmqName, null);
                }
            }
        }

        @Override
        public void onUnregister(String clientId, String group, String lmqName) {
        }

        /**
         * Mostly triggered when client channel closed, ensure that lite subscriptions is cleared before.
         */
        @Override
        public void onRemoveAll(String clientId, String group) {
            ClientEventSet eventSet = clientEventMap.remove(clientId);
            if (null == eventSet) {
                return;
            }
            LOGGER.warn("Maybe client offline. {}", clientId);
            String event;
            while ((event = eventSet.poll()) != null) {
                doDispatch(eventSet.group, event, clientId);
            }
        }
    }

    static class EventSetIterator implements Iterator<String> {
        private final ClientEventSet eventSet;

        public EventSetIterator(ClientEventSet eventSet) {
            this.eventSet = eventSet;
        }

        @Override
        public boolean hasNext() {
            return eventSet != null && !eventSet.events.isEmpty();
        }

        @Override
        public String next() {
            return eventSet.poll();
        }
    }

    protected static class FullDispatchRequest {
        private final String clientId;
        private final String group;
        private final long timestamp;

        public FullDispatchRequest(String clientId, String group, long delayMillis) {
            this.clientId = clientId;
            this.group = group;
            this.timestamp = System.currentTimeMillis() + delayMillis;
        }
    }

    // no need to compare group
    static final Comparator<FullDispatchRequest> COMPARATOR = (r1, r2) -> {
        if (null == r1 || null == r2 || null == r1.clientId || null == r2.clientId) {
            return 0;
        }
        if (r1.clientId.equals(r2.clientId)) {
            return 0;
        }
        int ret = Long.compare(r1.timestamp, r2.timestamp);
        if (ret != 0) {
            return ret;
        }
        return r1.clientId.compareTo(r2.clientId);
    };
}
