/*
 *
 * 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.bookkeeper.bookie.storage.ldb;

import static com.google.common.base.Preconditions.checkArgument;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.Unpooled;
import java.io.Closeable;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;
import java.util.concurrent.locks.ReentrantLock;
import org.apache.bookkeeper.common.util.MathUtils;
import org.apache.bookkeeper.util.collections.ConcurrentLongHashSet;
import org.apache.bookkeeper.util.collections.ConcurrentLongLongHashMap;
import org.apache.bookkeeper.util.collections.ConcurrentLongLongPairHashMap;
import org.apache.bookkeeper.util.collections.ConcurrentLongLongPairHashMap.LongPair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Write cache implementation.
 *
 * <p>The write cache will allocate the requested size from direct memory and it
 * will break it down into multiple segments.
 *
 * <p>The entries are appended in a common buffer and indexed though a hashmap,
 * until the cache is cleared.
 *
 * <p>There is the possibility to iterate through the stored entries in an ordered
 * way, by (ledgerId, entry).
 */
public class WriteCache implements Closeable {

    /**
     * Consumer that is used to scan the entire write cache.
     */
    public interface EntryConsumer {
        void accept(long ledgerId, long entryId, ByteBuf entry) throws IOException;
    }

    private final ConcurrentLongLongPairHashMap index = ConcurrentLongLongPairHashMap.newBuilder()
            .expectedItems(4096)
            .concurrencyLevel(2 * Runtime.getRuntime().availableProcessors())
            .build();

    private final ConcurrentLongLongHashMap lastEntryMap = ConcurrentLongLongHashMap.newBuilder()
            .expectedItems(4096)
            .concurrencyLevel(2 * Runtime.getRuntime().availableProcessors())
            .build();

    private final ByteBuf[] cacheSegments;
    private final int segmentsCount;

    private final long maxCacheSize;
    private final int maxSegmentSize;
    private final long segmentOffsetMask;
    private final long segmentOffsetBits;

    private final AtomicLong cacheSize = new AtomicLong(0);
    private final AtomicLong cacheOffset = new AtomicLong(0);
    private final LongAdder cacheCount = new LongAdder();

    private final ConcurrentLongHashSet deletedLedgers = ConcurrentLongHashSet.newBuilder().build();

    private final ByteBufAllocator allocator;

    public WriteCache(ByteBufAllocator allocator, long maxCacheSize) {
        // Default maxSegmentSize set to 1Gb
        this(allocator, maxCacheSize, 1 * 1024 * 1024 * 1024);
    }

    public WriteCache(ByteBufAllocator allocator, long maxCacheSize, int maxSegmentSize) {
        checkArgument(maxSegmentSize > 0);

        long alignedMaxSegmentSize = alignToPowerOfTwo(maxSegmentSize);
        checkArgument(maxSegmentSize == alignedMaxSegmentSize, "Max segment size needs to be in form of 2^n");

        this.allocator = allocator;
        this.maxCacheSize = maxCacheSize;
        this.maxSegmentSize = (int) maxSegmentSize;
        this.segmentOffsetMask = maxSegmentSize - 1;
        this.segmentOffsetBits = 63 - Long.numberOfLeadingZeros(maxSegmentSize);

        this.segmentsCount = 1 + (int) (maxCacheSize / maxSegmentSize);

        this.cacheSegments = new ByteBuf[segmentsCount];

        for (int i = 0; i < segmentsCount - 1; i++) {
            // All intermediate segments will be full-size
            cacheSegments[i] = Unpooled.directBuffer(maxSegmentSize, maxSegmentSize);
        }

        int lastSegmentSize = (int) (maxCacheSize % maxSegmentSize);
        cacheSegments[segmentsCount - 1] = Unpooled.directBuffer(lastSegmentSize, lastSegmentSize);
    }

    public void clear() {
        cacheSize.set(0L);
        cacheOffset.set(0L);
        cacheCount.reset();
        index.clear();
        lastEntryMap.clear();
        deletedLedgers.clear();
    }

    @Override
    public void close() {
        for (ByteBuf buf : cacheSegments) {
            buf.release();
        }
    }

    public boolean put(long ledgerId, long entryId, ByteBuf entry) {
        int size = entry.readableBytes();

        // Align to 64 bytes so that different threads will not contend the same L1
        // cache line
        int alignedSize = align64(size);

        long offset;
        int localOffset;
        int segmentIdx;

        while (true) {
            offset = cacheOffset.getAndAdd(alignedSize);
            localOffset = (int) (offset & segmentOffsetMask);
            segmentIdx = (int) (offset >>> segmentOffsetBits);

            if ((offset + size) > maxCacheSize) {
                // Cache is full
                return false;
            } else if (maxSegmentSize - localOffset < size) {
                // If an entry is at the end of a segment, we need to get a new offset and try
                // again in next segment
                continue;
            } else {
                // Found a good offset
                break;
            }
        }

        cacheSegments[segmentIdx].setBytes(localOffset, entry, entry.readerIndex(), entry.readableBytes());

        // Update last entryId for ledger. This logic is to handle writes for the same
        // ledger coming out of order and from different thread, though in practice it
        // should not happen and the compareAndSet should be always uncontended.
        while (true) {
            long currentLastEntryId = lastEntryMap.get(ledgerId);
            if (currentLastEntryId > entryId) {
                // A newer entry is already there
                break;
            }

            if (lastEntryMap.compareAndSet(ledgerId, currentLastEntryId, entryId)) {
                break;
            }
        }

        index.put(ledgerId, entryId, offset, size);
        cacheCount.increment();
        cacheSize.addAndGet(size);
        return true;
    }

    public ByteBuf get(long ledgerId, long entryId) {
        LongPair result = index.get(ledgerId, entryId);
        if (result == null) {
            return null;
        }

        long offset = result.first;
        int size = (int) result.second;
        ByteBuf entry = allocator.buffer(size, size);

        int localOffset = (int) (offset & segmentOffsetMask);
        int segmentIdx = (int) (offset >>> segmentOffsetBits);
        entry.writeBytes(cacheSegments[segmentIdx], localOffset, size);
        return entry;
    }

    public boolean hasEntry(long ledgerId, long entryId) {
        return index.get(ledgerId, entryId) != null;
    }

    public ByteBuf getLastEntry(long ledgerId) {
        long lastEntryId = lastEntryMap.get(ledgerId);
        if (lastEntryId == -1) {
            // Ledger not found in write cache
            return null;
        } else {
            return get(ledgerId, lastEntryId);
        }
    }

    public void deleteLedger(long ledgerId) {
        deletedLedgers.add(ledgerId);
    }

    public void forEach(EntryConsumer consumer) throws IOException {
        sortedEntriesLock.lock();

        try {
            int entriesToSort = (int) index.size();
            int arrayLen = entriesToSort * 4;
            if (sortedEntries == null || sortedEntries.length < arrayLen) {
                sortedEntries = new long[(int) (arrayLen * 2)];
            }

            long startTime = MathUtils.nowInNano();

            sortedEntriesIdx = 0;
            index.forEach((ledgerId, entryId, offset, length) -> {
                if (deletedLedgers.contains(ledgerId)) {
                    // Ignore deleted ledgers
                    return;
                }

                sortedEntries[sortedEntriesIdx] = ledgerId;
                sortedEntries[sortedEntriesIdx + 1] = entryId;
                sortedEntries[sortedEntriesIdx + 2] = offset;
                sortedEntries[sortedEntriesIdx + 3] = length;
                sortedEntriesIdx += 4;
            });

            if (log.isDebugEnabled()) {
                log.debug("iteration took {} ms", MathUtils.elapsedNanos(startTime) / 1e6);
            }
            startTime = MathUtils.nowInNano();

            // Sort entries by (ledgerId, entryId) maintaining the 4 items groups
            ArrayGroupSort.sort(sortedEntries, 0, sortedEntriesIdx);
            if (log.isDebugEnabled()) {
                log.debug("sorting {} ms", (MathUtils.elapsedNanos(startTime) / 1e6));
            }
            startTime = MathUtils.nowInNano();

            ByteBuf[] entrySegments = new ByteBuf[segmentsCount];
            for (int i = 0; i < segmentsCount; i++) {
                entrySegments[i] = cacheSegments[i].slice(0, cacheSegments[i].capacity());
            }

            for (int i = 0; i < sortedEntriesIdx; i += 4) {
                long ledgerId = sortedEntries[i];
                long entryId = sortedEntries[i + 1];
                long offset = sortedEntries[i + 2];
                long length = sortedEntries[i + 3];

                int localOffset = (int) (offset & segmentOffsetMask);
                int segmentIdx = (int) (offset >>> segmentOffsetBits);
                ByteBuf entry = entrySegments[segmentIdx];
                entry.setIndex(localOffset, localOffset + (int) length);
                consumer.accept(ledgerId, entryId, entry);
            }

            if (log.isDebugEnabled()) {
                log.debug("entry log adding {} ms", MathUtils.elapsedNanos(startTime) / 1e6);
            }
        } finally {
            sortedEntriesLock.unlock();
        }
    }

    public long size() {
        return cacheSize.get();
    }

    public long count() {
        return cacheCount.sum();
    }

    public boolean isEmpty() {
        return cacheSize.get() == 0L;
    }

    private static final int ALIGN_64_MASK = ~(64 - 1);

    static int align64(int size) {
        return (size + 64 - 1) & ALIGN_64_MASK;
    }

    private static long alignToPowerOfTwo(long n) {
        return (long) Math.pow(2, 64 - Long.numberOfLeadingZeros(n - 1));
    }

    private final ReentrantLock sortedEntriesLock = new ReentrantLock();
    private long[] sortedEntries;
    private int sortedEntriesIdx;

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