/*
 * 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.store.queue;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import org.apache.rocketmq.common.BoundaryType;
import org.apache.rocketmq.common.Pair;
import org.apache.rocketmq.common.attribute.CQType;
import org.apache.rocketmq.common.constant.LoggerName;
import org.apache.rocketmq.common.message.MessageExtBrokerInner;
import org.apache.rocketmq.logging.org.slf4j.Logger;
import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
import org.apache.rocketmq.store.ConsumeQueue;
import org.apache.rocketmq.store.DispatchRequest;
import org.apache.rocketmq.store.MessageFilter;
import org.apache.rocketmq.store.config.MessageStoreConfig;
import org.rocksdb.RocksDBException;

public class RocksDBConsumeQueue implements ConsumeQueueInterface {
    private static final Logger log = LoggerFactory.getLogger(LoggerName.STORE_LOGGER_NAME);
    private static final Logger ERROR_LOG = LoggerFactory.getLogger(LoggerName.STORE_ERROR_LOGGER_NAME);

    private final MessageStoreConfig messageStoreConfig;
    private final RocksDBConsumeQueueStore consumeQueueStore;
    private final String topic;
    private final int queueId;

    public RocksDBConsumeQueue(final MessageStoreConfig messageStoreConfig,
        final RocksDBConsumeQueueStore consumeQueueStore,
        final String topic, final int queueId) {
        this.messageStoreConfig = messageStoreConfig;
        this.consumeQueueStore = consumeQueueStore;
        this.topic = topic;
        this.queueId = queueId;
    }

    /**
     * Only used to pass parameters when calling the destroy method
     *
     * @see RocksDBConsumeQueueStore#destroy(ConsumeQueueInterface)
     */
    public RocksDBConsumeQueue(final String topic, final int queueId) {
        this(null, null, topic, queueId);
    }

    @Override
    public boolean load() {
        return true;
    }

    @Override
    public void recover() {
        // ignore
    }

    @Override
    public void checkSelf() {
        // ignore
    }

    @Override
    public boolean flush(final int flushLeastPages) {
        return true;
    }

    @Override
    public void destroy() {
        // ignore
    }

    @Override
    public void truncateDirtyLogicFiles(long maxCommitLogPos) {
        // ignored
    }

    @Override
    public int deleteExpiredFile(long minCommitLogPos) {
        return 0;
    }

    @Override
    public long rollNextFile(long nextBeginOffset) {
        return 0;
    }

    @Override
    public boolean isFirstFileAvailable() {
        return true;
    }

    @Override
    public boolean isFirstFileExist() {
        return true;
    }

    @Override
    public void swapMap(int reserveNum, long forceSwapIntervalMs, long normalSwapIntervalMs) {
        // ignore
    }

    @Override
    public void cleanSwappedMap(long forceCleanSwapIntervalMs) {
        // ignore
    }

    @Override
    public long getMaxOffsetInQueue() {
        try {
            return this.consumeQueueStore.getMaxOffsetInQueue(topic, queueId);
        } catch (RocksDBException e) {
            ERROR_LOG.error("getMaxOffsetInQueue Failed. topic: {}, queueId: {}", topic, queueId, e);
            return 0;
        }
    }

    @Override
    public long getMessageTotalInQueue() {
        try {
            long maxOffsetInQueue = this.consumeQueueStore.getMaxOffsetInQueue(topic, queueId);
            long minOffsetInQueue = this.consumeQueueStore.getMinOffsetInQueue(topic, queueId);
            return maxOffsetInQueue - minOffsetInQueue;
        } catch (RocksDBException e) {
            ERROR_LOG.error("getMessageTotalInQueue Failed. topic: {}, queueId: {}, {}", topic, queueId, e);
        }
        return -1;
    }

    /**
     * We already implement it in RocksDBConsumeQueueStore.
     * @see RocksDBConsumeQueueStore#getOffsetInQueueByTime
     * @param timestamp timestamp
     * @return
     */
    @Override
    public long getOffsetInQueueByTime(long timestamp) {
        return 0;
    }

    /**
     * We already implement it in RocksDBConsumeQueueStore.
     * @see RocksDBConsumeQueueStore#getOffsetInQueueByTime
     * @param timestamp    timestamp
     * @param boundaryType Lower or Upper
     * @return
     */
    @Override
    public long getOffsetInQueueByTime(long timestamp, BoundaryType boundaryType) {
        return 0;
    }

    @Override
    public long getMaxPhysicOffset() {
        Long maxPhyOffset = this.consumeQueueStore.getMaxPhyOffsetInConsumeQueue(topic, queueId);
        return maxPhyOffset == null ? -1 : maxPhyOffset;
    }

    @Override
    public long getMinLogicOffset() {
        return 0;
    }

    @Override
    public CQType getCQType() {
        return CQType.RocksDBCQ;
    }

    @Override
    public long getTotalSize() {
        // ignored
        return 0;
    }

    @Override
    public int getUnitSize() {
        // attention: unitSize should equal to 'ConsumeQueue.CQ_STORE_UNIT_SIZE'
        return ConsumeQueue.CQ_STORE_UNIT_SIZE;
    }

    /**
     * Ignored, we already implement this method
     * @see org.apache.rocketmq.store.queue.RocksDBConsumeQueueOffsetTable#getMinCqOffset(String, int)
     */
    @Override
    public void correctMinOffset(long minCommitLogOffset) {

    }

    /**
     * Ignored, in rocksdb mode, we build cq in RocksDBConsumeQueueStore
     */
    @Override
    public void putMessagePositionInfoWrapper(DispatchRequest request) {

    }

    @Override
    public void assignQueueOffset(QueueOffsetOperator queueOffsetOperator, MessageExtBrokerInner msg) throws RocksDBException {
        String topicQueueKey = getTopic() + "-" + getQueueId();
        Long queueOffset = queueOffsetOperator.getTopicQueueNextOffset(topicQueueKey);
        if (queueOffset == null) {
            // we will recover topic queue table from rocksdb when we use it.
            queueOffset = this.consumeQueueStore.getMaxOffsetInQueue(topic, queueId);
            queueOffsetOperator.updateQueueOffset(topicQueueKey, queueOffset);
        }
        msg.setQueueOffset(queueOffset);
    }

    @Override
    public void increaseQueueOffset(QueueOffsetOperator queueOffsetOperator, MessageExtBrokerInner msg, short messageNum) {
        String topicQueueKey = getTopic() + "-" + getQueueId();
        queueOffsetOperator.increaseQueueOffset(topicQueueKey, messageNum);
    }

    /**
     * It is CPU-intensive with many offline group
     * Optimize by caching their estimated info
     */
    @Override
    public long estimateMessageCount(long from, long to, MessageFilter filter) {

        Pair<CqUnit, Long> fromUnit = getCqUnitAndStoreTime(from);
        if (fromUnit == null || from >= to) {
            return -1L;
        }

        if (to > getMaxOffsetInQueue()) {
            to = getMaxOffsetInQueue();
        }

        int sampleCount = 0;
        int sampleTotal = Math.min((int) (to - from), messageStoreConfig.getMaxConsumeQueueScan());

        int matchCount = 0;
        int matchTotal = messageStoreConfig.getSampleCountThreshold();

        try {
            ReferredIterator<CqUnit> iterator = this.iterateFrom(from, matchTotal);
            while (iterator != null && iterator.hasNext() && sampleCount++ < sampleTotal) {
                CqUnit cqUnit = iterator.next();
                if (filter.isMatchedByConsumeQueue(
                    cqUnit.getTagsCode(), cqUnit.getCqExtUnit())) {
                    if (++matchCount > matchTotal) {
                        sampleTotal = sampleCount;
                        break;
                    }
                }
            }
        } catch (Throwable t) {
            log.error("EstimateLag error, from={}, to={}", from, to, t);
        }

        long result = sampleTotal == 0 ? 0 :
            (long) ((to - from) * (matchCount / (sampleTotal * 1.0)));

        if (log.isTraceEnabled()) {
            log.trace("EstimateLag, topic={}, queueId={}, offset={}-{}, total={}, hit rate={}/{}({}%), result={}",
                topic, queueId, from, to, to - from,
                matchCount, sampleCount, String.format("%.1f", (double) matchCount * 100.0 / sampleCount), result);
        }

        return result;
    }


    @Override
    public long getMinOffsetInQueue() {
        try {
            return this.consumeQueueStore.getMinOffsetInQueue(topic, queueId);
        } catch (RocksDBException e) {
            ERROR_LOG.error("getMinOffsetInQueue Failed. topic: {}, queueId: {}", topic, queueId, e);
            return -1;
        }
    }

    private int pullNum(long cqOffset, long maxCqOffset) {
        long diffLong = maxCqOffset - cqOffset;
        if (diffLong < Integer.MAX_VALUE) {
            return (int) diffLong;
        }
        return Integer.MAX_VALUE;
    }

    @Override
    public ReferredIterator<CqUnit> iterateFrom(final long startIndex) {
        long maxCqOffset = getMaxOffsetInQueue();
        if (startIndex < maxCqOffset && startIndex >= 0) {
            int num = pullNum(startIndex, maxCqOffset);
            return new LargeRocksDBConsumeQueueIterator(startIndex, num);
        }
        return null;
    }

    @Override
    public ReferredIterator<CqUnit> iterateFrom(long startIndex, int count) throws RocksDBException {
        long maxCqOffset = getMaxOffsetInQueue();
        if (startIndex < maxCqOffset) {
            int num = Math.min((int)(maxCqOffset - startIndex), count);
            return iterateFrom0(startIndex, num, maxCqOffset);
        }
        return null;
    }

    @Override
    public CqUnit get(long index) {
        Pair<CqUnit, Long> pair = getCqUnitAndStoreTime(index);
        return pair == null ? null : pair.getObject1();
    }

    @Override
    public Pair<CqUnit, Long> getCqUnitAndStoreTime(long index) {
        ByteBuffer byteBuffer;
        try {
            byteBuffer = this.consumeQueueStore.get(topic, queueId, index);
        } catch (RocksDBException e) {
            ERROR_LOG.error("getUnitAndStoreTime Failed. topic: {}, queueId: {}", topic, queueId, e);
            return null;
        }
        if (byteBuffer == null || byteBuffer.remaining() < RocksDBConsumeQueueTable.CQ_UNIT_SIZE) {
            return null;
        }
        long phyOffset = byteBuffer.getLong();
        int size = byteBuffer.getInt();
        long tagCode = byteBuffer.getLong();
        long messageStoreTime = byteBuffer.getLong();
        return new Pair<>(new CqUnit(index, phyOffset, size, tagCode), messageStoreTime);
    }

    @Override
    public Pair<CqUnit, Long> getEarliestUnitAndStoreTime() {
        try {
            long minOffset = this.consumeQueueStore.getMinOffsetInQueue(topic, queueId);
            return getCqUnitAndStoreTime(minOffset);
        } catch (RocksDBException e) {
            ERROR_LOG.error("getEarliestUnitAndStoreTime Failed. topic: {}, queueId: {}", topic, queueId, e);
        }
        return null;
    }

    @Override
    public CqUnit getEarliestUnit() {
        Pair<CqUnit, Long> pair = getEarliestUnitAndStoreTime();
        return pair == null ? null : pair.getObject1();
    }

    @Override
    public CqUnit getLatestUnit() {
        try {
            long maxOffset = this.consumeQueueStore.getMaxOffsetInQueue(topic, queueId);
            return get(maxOffset > 0 ? maxOffset - 1 : maxOffset);
        } catch (RocksDBException e) {
            ERROR_LOG.error("getLatestUnit Failed. topic: {}, queueId: {}, {}", topic, queueId, e.getMessage());
        }
        return null;
    }

    @Override
    public long getLastOffset() {
        return getMaxPhysicOffset();
    }

    private ReferredIterator<CqUnit> iterateFrom0(
        final long startIndex, final int count, final long maxOffset) throws RocksDBException {

        if (messageStoreConfig.isIteratorWhenUseRocksdbConsumeQueue()) {
            return new RocksDBReusableIterator(topic, queueId, startIndex, count, maxOffset);
        }

        List<ByteBuffer> byteBufferList = this.consumeQueueStore.rangeQuery(topic, queueId, startIndex, count);
        if (byteBufferList == null || byteBufferList.isEmpty()) {
            if (this.messageStoreConfig.isEnableRocksDBLog()) {
                log.warn("iterateFrom0 - find nothing, startIndex:{}, count:{}", startIndex, count);
            }
            return null;
        }
        return new RocksDBConsumeQueueIterator(byteBufferList, startIndex);
    }

    @Override
    public String getTopic() {
        return topic;
    }

    @Override
    public int getQueueId() {
        return queueId;
    }

    private class RocksDBReusableIterator implements ReferredIterator<CqUnit> {

        private final String topic;
        private final int queueId;
        private long offset;
        private final int count;
        private final long maxOffset;

        private int bufferIndex;
        private List<ByteBuffer> buffers;

        // offset + count <= max offset
        public RocksDBReusableIterator(String topic, int queueId, long offset, int count, long maxOffset) {
            this.topic = topic;
            this.queueId = queueId;
            this.offset = offset;
            this.count = count;
            this.maxOffset = maxOffset;

            this.bufferIndex = 0;
            this.buffers = new ArrayList<>(count);
        }

        @Override
        public void release() {
        }

        @Override
        public CqUnit nextAndRelease() {
            try {
                return next();
            } finally {
                release();
            }
        }

        @Override
        public boolean hasNext() {
            return offset < maxOffset;
        }

        @Override
        public CqUnit next() {
            try {
                if (buffers.isEmpty() || bufferIndex >= buffers.size()) {
                    int batchSize = (int) Math.min(count, maxOffset - offset);
                    if (batchSize == 0) {
                        return null;
                    } else {
                        bufferIndex = 0;
                        buffers = consumeQueueStore.rangeQuery(topic, queueId, offset, batchSize);
                    }
                }
                if (bufferIndex < buffers.size()) {
                    ByteBuffer buffer = buffers.get(bufferIndex++);
                    return new CqUnit(offset++, buffer.getLong(), buffer.getInt(), buffer.getLong());
                }
            } catch (Throwable t) {
                log.error("RocksDB reusable iterator search error, " +
                    "topic={}, queueId={}, offset={}, count={}", topic, queueId, offset, count, maxOffset, t);
            }
            return null;
        }
    }

    private class RocksDBConsumeQueueIterator implements ReferredIterator<CqUnit> {
        private final List<ByteBuffer> byteBufferList;
        private final long startIndex;
        private final int totalCount;
        private int currentIndex;

        public RocksDBConsumeQueueIterator(final List<ByteBuffer> byteBufferList, final long startIndex) {
            this.byteBufferList = byteBufferList;
            this.startIndex = startIndex;
            this.totalCount = byteBufferList.size();
            this.currentIndex = 0;
        }

        @Override
        public boolean hasNext() {
            return this.currentIndex < this.totalCount;
        }

        @Override
        public CqUnit next() {
            if (!hasNext()) {
                return null;
            }
            final int currentIndex = this.currentIndex;
            final ByteBuffer byteBuffer = this.byteBufferList.get(currentIndex);
            CqUnit cqUnit = new CqUnit(this.startIndex + currentIndex, byteBuffer.getLong(), byteBuffer.getInt(), byteBuffer.getLong());
            this.currentIndex++;
            return cqUnit;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("remove");
        }

        @Override
        public void release() {
        }

        @Override
        public CqUnit nextAndRelease() {
            try {
                return next();
            } finally {
                release();
            }
        }
    }

    private class LargeRocksDBConsumeQueueIterator implements ReferredIterator<CqUnit> {
        private final long startIndex;
        private final int totalCount;
        private int currentIndex;

        public LargeRocksDBConsumeQueueIterator(final long startIndex, final int num) {
            this.startIndex = startIndex;
            this.totalCount = num;
            this.currentIndex = 0;
        }

        @Override
        public boolean hasNext() {
            return this.currentIndex < this.totalCount;
        }


        @Override
        public CqUnit next() {
            if (!hasNext()) {
                return null;
            }

            final ByteBuffer byteBuffer;
            try {
                byteBuffer = consumeQueueStore.get(topic, queueId, startIndex + currentIndex);
            } catch (RocksDBException e) {
                ERROR_LOG.error("get cq from rocksdb failed. topic: {}, queueId: {}", topic, queueId, e);
                return null;
            }
            if (byteBuffer == null || byteBuffer.remaining() < RocksDBConsumeQueueTable.CQ_UNIT_SIZE) {
                return null;
            }
            CqUnit cqUnit = new CqUnit(this.startIndex + currentIndex, byteBuffer.getLong(), byteBuffer.getInt(), byteBuffer.getLong());
            this.currentIndex++;
            return cqUnit;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("remove");
        }

        @Override
        public void release() {
        }

        @Override
        public CqUnit nextAndRelease() {
            try {
                return next();
            } finally {
                release();
            }
        }
    }

    public void initializeWithOffset(long offset, long minPhyOffset) {
        log.info("RocksDBConsumeQueue initializeWithOffset topic={}, queueId={}, offset={}, oldMax={}, oldMin={}",
            topic, queueId, offset, getMaxOffsetInQueue(), getMinOffsetInQueue());
        try {
            // clean the expired cqUnit and offset
            consumeQueueStore.cleanExpired(minPhyOffset);

            // update the max and min offset
            this.consumeQueueStore.updateCqOffset(topic, queueId, 0L, offset - 1, true);
            // set phyOffset to 0, min cq offset will be lazy corrected to max cq Offset + 1
            this.consumeQueueStore.updateCqOffset(topic, queueId, 0L, offset, false);
        } catch (RocksDBException e) {
            ERROR_LOG.error("RocksDBConsumeQueue initializeWithOffset Failed. topic={}, queueId={}, offset={}", topic, queueId, offset, e);
        }
    }

    @Override
    public boolean shutdown() {
        return true;
    }
}
