/*
 * 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.kafka.storage.internals.epoch;

import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.server.util.Scheduler;
import org.apache.kafka.storage.internals.checkpoint.LeaderEpochCheckpointFile;
import org.apache.kafka.storage.internals.log.EpochEntry;

import org.slf4j.Logger;

import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.NavigableMap;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.TreeMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.Predicate;

import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH;
import static org.apache.kafka.common.requests.OffsetsForLeaderEpochResponse.UNDEFINED_EPOCH_OFFSET;

/**
 * Represents a cache of (LeaderEpoch => Offset) mappings for a particular replica.
 * <p>
 * Leader Epoch = epoch assigned to each leader by the controller.
 * Offset = offset of the first message in each epoch.
 * <p>
 * Note that {@link #truncateFromStartAsyncFlush},{@link #truncateFromEndAsyncFlush} flush the epoch-entry changes to checkpoint asynchronously.
 * Hence, it is instantiater's responsibility to ensure restoring the cache to the correct state after instantiating
 * this class from checkpoint (which might contain stale epoch entries right after instantiation).
 */
public final class LeaderEpochFileCache {
    private final TopicPartition topicPartition;
    private final LeaderEpochCheckpointFile checkpoint;
    private final Scheduler scheduler;
    private final Logger log;

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final TreeMap<Integer, EpochEntry> epochs = new TreeMap<>();


    /**
     * @param topicPartition the associated topic partition
     * @param checkpoint     the checkpoint file
     * @param scheduler      the scheduler to use for async I/O operations
     */
    public LeaderEpochFileCache(TopicPartition topicPartition, LeaderEpochCheckpointFile checkpoint, Scheduler scheduler) {
        this.checkpoint = checkpoint;
        this.topicPartition = topicPartition;
        this.scheduler = scheduler;
        LogContext logContext = new LogContext("[LeaderEpochCache " + topicPartition + "] ");
        log = logContext.logger(LeaderEpochFileCache.class);
        checkpoint.read().forEach(this::assign);
    }

    /**
     * Instantiate a new LeaderEpochFileCache with provided epoch entries instead of from the backing checkpoint file.
     * The provided epoch entries are expected to be no less fresh than the checkpoint file.
     * @param epochEntries the current epoch entries
     * @param topicPartition the associated topic partition
     * @param checkpoint the checkpoint file
     * @param scheduler the scheduler to use for async I/O operations
     */
    private LeaderEpochFileCache(List<EpochEntry> epochEntries,
                                 TopicPartition topicPartition,
                                 LeaderEpochCheckpointFile checkpoint,
                                 Scheduler scheduler) {
        this.checkpoint = checkpoint;
        this.topicPartition = topicPartition;
        this.scheduler = scheduler;
        LogContext logContext = new LogContext("[LeaderEpochCache " + topicPartition + "] ");
        log = logContext.logger(LeaderEpochFileCache.class);
        for (EpochEntry entry : epochEntries) {
            epochs.put(entry.epoch(), entry);
        }
    }

    /**
     * Assigns the supplied Leader Epoch to the supplied Offset
     * Once the epoch is assigned it cannot be reassigned
     */
    public void assign(int epoch, long startOffset) {
        EpochEntry entry = new EpochEntry(epoch, startOffset);
        if (assign(entry)) {
            log.debug("Appended new epoch entry {}. Cache now contains {} entries.", entry, epochs.size());
            writeToFile();
        }
    }

    public void assign(List<EpochEntry> entries) {
        entries.forEach(entry -> {
            if (assign(entry)) {
                log.debug("Appended new epoch entry {}. Cache now contains {} entries.", entry, epochs.size());
            }
        });
        if (!entries.isEmpty()) writeToFile();
    }

    private boolean isUpdateNeeded(EpochEntry entry) {
        return latestEntry().map(epochEntry -> entry.epoch() != epochEntry.epoch() || entry.startOffset() < epochEntry.startOffset()).orElse(true);
    }

    private boolean assign(EpochEntry entry) {
        if (entry.epoch() < 0 || entry.startOffset() < 0) {
            throw new IllegalArgumentException("Received invalid partition leader epoch entry " + entry);
        }

        // Check whether the append is needed before acquiring the write lock
        // in order to avoid contention with readers in the common case
        if (!isUpdateNeeded(entry)) return false;

        lock.writeLock().lock();
        try {
            if (isUpdateNeeded(entry)) {
                maybeTruncateNonMonotonicEntries(entry);
                epochs.put(entry.epoch(), entry);
                return true;
            } else {
                return false;
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * Remove any entries which violate monotonicity prior to appending a new entry
     */
    private void maybeTruncateNonMonotonicEntries(EpochEntry newEntry) {
        List<EpochEntry> removedEpochs = removeWhileMatching(
                epochs.descendingMap().entrySet().iterator(),
                entry -> entry.epoch() >= newEntry.epoch() || entry.startOffset() >= newEntry.startOffset());

        if (removedEpochs.size() > 1 || (!removedEpochs.isEmpty() && removedEpochs.get(0).startOffset() != newEntry.startOffset())) {

            // Only log a warning if there were non-trivial removals. If the start offset of the new entry
            // matches the start offset of the removed epoch, then no data has been written and the truncation
            // is expected.
            log.warn("New epoch entry {} caused truncation of conflicting entries {}. " + "Cache now contains {} entries.", newEntry, removedEpochs, epochs.size());
        }
    }

    private static List<EpochEntry> removeWhileMatching(Iterator<Map.Entry<Integer, EpochEntry>> iterator, Predicate<EpochEntry> predicate) {
        ArrayList<EpochEntry> removedEpochs = new ArrayList<>();

        while (iterator.hasNext()) {
            EpochEntry entry = iterator.next().getValue();
            if (predicate.test(entry)) {
                removedEpochs.add(entry);
                iterator.remove();
            } else {
                return removedEpochs;
            }
        }

        return removedEpochs;
    }

    public boolean nonEmpty() {
        lock.readLock().lock();
        try {
            return !epochs.isEmpty();
        } finally {
            lock.readLock().unlock();
        }
    }

    public Optional<EpochEntry> latestEntry() {
        lock.readLock().lock();
        try {
            return Optional.ofNullable(epochs.lastEntry()).map(Map.Entry::getValue);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Returns the current Leader Epoch if one exists. This is the latest epoch
     * which has messages assigned to it.
     */
    public Optional<Integer> latestEpoch() {
        return latestEntry().map(epochEntry -> epochEntry.epoch());
    }

    public OptionalInt previousEpoch() {
        lock.readLock().lock();
        try {
            return latestEntry().flatMap(entry -> Optional.ofNullable(epochs.lowerEntry(entry.epoch())))
                    .map(integerEpochEntryEntry -> OptionalInt.of(integerEpochEntryEntry.getKey())).orElseGet(OptionalInt::empty);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Get the earliest cached entry if one exists.
     */
    public Optional<EpochEntry> earliestEntry() {
        lock.readLock().lock();
        try {
            return Optional.ofNullable(epochs.firstEntry()).map(Map.Entry::getValue);
        } finally {
            lock.readLock().unlock();
        }
    }

    public OptionalInt previousEpoch(int epoch) {
        lock.readLock().lock();
        try {
            return toOptionalInt(epochs.lowerKey(epoch));
        } finally {
            lock.readLock().unlock();
        }
    }

    public Optional<EpochEntry> previousEntry(int epoch) {
        lock.readLock().lock();
        try {
            return Optional.ofNullable(epochs.lowerEntry(epoch)).map(Map.Entry::getValue);
        } finally {
            lock.readLock().unlock();
        }
    }

    public OptionalInt nextEpoch(int epoch) {
        lock.readLock().lock();
        try {
            return toOptionalInt(epochs.higherKey(epoch));
        } finally {
            lock.readLock().unlock();
        }
    }

    private static OptionalInt toOptionalInt(Integer value) {
        return (value != null) ? OptionalInt.of(value) : OptionalInt.empty();
    }

    public Optional<EpochEntry> epochEntry(int epoch) {
        lock.readLock().lock();
        try {
            return Optional.ofNullable(epochs.get(epoch));
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Returns the Leader Epoch and the End Offset for a requested Leader Epoch.
     * <p>
     * The Leader Epoch returned is the largest epoch less than or equal to the requested Leader
     * Epoch. The End Offset is the end offset of this epoch, which is defined as the start offset
     * of the first Leader Epoch larger than the Leader Epoch requested, or else the Log End
     * Offset if the latest epoch was requested.
     * <p>
     * During the upgrade phase, where there are existing messages may not have a leader epoch,
     * if requestedEpoch is < the first epoch cached, UNDEFINED_EPOCH_OFFSET will be returned
     * so that the follower falls back to High Water Mark.
     *
     * @param requestedEpoch requested leader epoch
     * @param logEndOffset   the existing Log End Offset
     * @return found leader epoch and end offset
     */
    public Map.Entry<Integer, Long> endOffsetFor(int requestedEpoch, long logEndOffset) {
        lock.readLock().lock();
        try {
            Map.Entry<Integer, Long> epochAndOffset;
            if (requestedEpoch == UNDEFINED_EPOCH) {
                // This may happen if a bootstrapping follower sends a request with undefined epoch or
                // a follower is on the older message format where leader epochs are not recorded
                epochAndOffset = new AbstractMap.SimpleImmutableEntry<>(UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
            } else if (latestEpoch().isPresent() && latestEpoch().get() == requestedEpoch) {
                // For the leader, the latest epoch is always the current leader epoch that is still being written to.
                // Followers should not have any reason to query for the end offset of the current epoch, but a consumer
                // might if it is verifying its committed offset following a group rebalance. In this case, we return
                // the current log end offset which makes the truncation check work as expected.
                epochAndOffset = new AbstractMap.SimpleImmutableEntry<>(requestedEpoch, logEndOffset);
            } else {
                Map.Entry<Integer, EpochEntry> higherEntry = epochs.higherEntry(requestedEpoch);
                if (higherEntry == null) {
                    // The requested epoch is larger than any known epoch. This case should never be hit because
                    // the latest cached epoch is always the largest.
                    epochAndOffset = new AbstractMap.SimpleImmutableEntry<>(UNDEFINED_EPOCH, UNDEFINED_EPOCH_OFFSET);
                } else {
                    Map.Entry<Integer, EpochEntry> floorEntry = epochs.floorEntry(requestedEpoch);
                    if (floorEntry == null) {
                        // The requested epoch is smaller than any known epoch, so we return the start offset of the first
                        // known epoch which is larger than it. This may be inaccurate as there could have been
                        // epochs in between, but the point is that the data has already been removed from the log
                        // and we want to ensure that the follower can replicate correctly beginning from the leader's
                        // start offset.
                        epochAndOffset = new AbstractMap.SimpleImmutableEntry<>(requestedEpoch, higherEntry.getValue().startOffset());
                    } else {
                        // We have at least one previous epoch and one subsequent epoch. The result is the first
                        // prior epoch and the starting offset of the first subsequent epoch.
                        epochAndOffset = new AbstractMap.SimpleImmutableEntry<>(floorEntry.getValue().epoch(), higherEntry.getValue().startOffset());
                    }
                }
            }

            if (log.isTraceEnabled())
                log.trace("Processed end offset request for epoch {} and returning epoch {} with end offset {} from epoch cache of size {}}", requestedEpoch, epochAndOffset.getKey(), epochAndOffset.getValue(), epochs.size());

            return epochAndOffset;
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Removes all epoch entries from the store with start offsets greater than or equal to the passed offset.
     * <p>
     * Checkpoint-flushing is done asynchronously.
     */
    public void truncateFromEndAsyncFlush(long endOffset) {
        lock.writeLock().lock();
        try {
            List<EpochEntry> removedEntries = truncateFromEnd(epochs, endOffset);
            if (!removedEntries.isEmpty()) {
                // We flush the change to the device in the background because:
                // - To avoid fsync latency
                //   * fsync latency could be huge on a disk glitch, which is not rare in spinning drives
                //   * This method is called by ReplicaFetcher threads, which could block replica fetching
                //     then causing ISR shrink or high produce response time degradation in remote scope on high fsync latency.
                // - We still flush the change in #assign synchronously, meaning that it's guaranteed that the checkpoint file always has no missing entries.
                //   * Even when stale epochs are restored from the checkpoint file after the unclean shutdown, it will be handled by
                //     another truncateFromEnd call on log loading procedure, so it won't be a problem
                scheduler.scheduleOnce("leader-epoch-cache-flush-" + topicPartition, this::writeIfDirExists);

                log.debug("Cleared entries {} from epoch cache after truncating to end offset {}, leaving {} entries in the cache.", removedEntries, endOffset, epochs.size());
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * Clears old epoch entries. This method searches for the oldest epoch < offset, updates the saved epoch offset to
     * be offset, then clears any previous epoch entries.
     * <p>
     * This method is exclusive: so truncateFromStart(6) will retain an entry at offset 6.
     * <p>
     * Checkpoint-flushing is done asynchronously.
     *
     * @param startOffset the offset to clear up to
     */
    public void truncateFromStartAsyncFlush(long startOffset) {
        lock.writeLock().lock();
        try {
            List<EpochEntry> removedEntries = truncateFromStart(epochs, startOffset);
            if (!removedEntries.isEmpty()) {
                // We flush the change to the device in the background because:
                // - To avoid fsync latency
                //   * fsync latency could be huge on a disk glitch, which is not rare in spinning drives
                //   * This method is called as part of deleteRecords with holding UnifiedLog#lock.
                //      - Meanwhile all produces against the partition will be blocked, which causes req-handlers to exhaust
                // - We still flush the change in #assign synchronously, meaning that it's guaranteed that the checkpoint file always has no missing entries.
                //   * Even when stale epochs are restored from the checkpoint file after the unclean shutdown, it will be handled by
                //     another truncateFromStart call on log loading procedure, so it won't be a problem
                scheduler.scheduleOnce("leader-epoch-cache-flush-" + topicPartition, this::writeIfDirExists);

                EpochEntry updatedFirstEntry = removedEntries.get(removedEntries.size() - 1);
                log.debug("Cleared entries {} and rewrote first entry {} after truncating to start offset {}, leaving {} in the cache.", removedEntries, updatedFirstEntry, startOffset, epochs.size());
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    private static List<EpochEntry> truncateFromStart(TreeMap<Integer, EpochEntry> epochs, long startOffset) {
        List<EpochEntry> removedEntries = removeWhileMatching(
                epochs.entrySet().iterator(), entry -> entry.startOffset() <= startOffset);

        if (!removedEntries.isEmpty()) {
            EpochEntry firstBeforeStartOffset = removedEntries.get(removedEntries.size() - 1);
            EpochEntry updatedFirstEntry = new EpochEntry(firstBeforeStartOffset.epoch(), startOffset);
            epochs.put(updatedFirstEntry.epoch(), updatedFirstEntry);
        }

        return removedEntries;
    }

    private static List<EpochEntry> truncateFromEnd(TreeMap<Integer, EpochEntry> epochs, long endOffset) {
        Optional<EpochEntry> epochEntry = Optional.ofNullable(epochs.lastEntry()).map(Entry::getValue);
        if (endOffset >= 0 && epochEntry.isPresent() && epochEntry.get().startOffset() >= endOffset) {
            return removeWhileMatching(epochs.descendingMap().entrySet().iterator(), x -> x.startOffset() >= endOffset);
        }
        return List.of();
    }

    public OptionalInt epochForOffset(long offset) {
        lock.readLock().lock();
        try {
            OptionalInt previousEpoch = OptionalInt.empty();
            for (EpochEntry epochEntry : epochs.values()) {
                int epoch = epochEntry.epoch();
                long startOffset = epochEntry.startOffset();

                // Found the exact offset, return the respective epoch.
                if (startOffset == offset) return OptionalInt.of(epoch);

                // Return the previous found epoch as this epoch's start-offset is more than the target offset.
                if (startOffset > offset) return previousEpoch;

                previousEpoch = OptionalInt.of(epoch);
            }

            return previousEpoch;
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Returns a new LeaderEpochFileCache which contains same
     * epoch entries with replacing backing checkpoint file.
     * @param leaderEpochCheckpoint the new checkpoint file
     * @return a new LeaderEpochFileCache instance
     */
    public LeaderEpochFileCache withCheckpoint(LeaderEpochCheckpointFile leaderEpochCheckpoint) {
        lock.readLock().lock();
        try {
            return new LeaderEpochFileCache(epochEntries(),
                                            topicPartition,
                                            leaderEpochCheckpoint,
                                            scheduler);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Returns the leader epoch entries within the range of the given start and end offset
     * @param startOffset The start offset of the epoch entries (inclusive).
     * @param endOffset   The end offset of the epoch entries (exclusive)
     * @return the leader epoch entries
     */
    public List<EpochEntry> epochEntriesInRange(long startOffset, long endOffset) {
        lock.readLock().lock();
        try {
            TreeMap<Integer, EpochEntry> epochsCopy = new TreeMap<>(this.epochs);
            if (startOffset >= 0) {
                truncateFromStart(epochsCopy, startOffset);
            }
            truncateFromEnd(epochsCopy, endOffset);
            return new ArrayList<>(epochsCopy.values());
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * Delete all entries.
     */
    public void clearAndFlush() {
        lock.writeLock().lock();
        try {
            epochs.clear();
            writeToFile();
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void clear() {
        lock.writeLock().lock();
        try {
            epochs.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }

    public List<EpochEntry> epochEntries() {
        lock.readLock().lock();
        try {
            return new ArrayList<>(epochs.values());
        } finally {
            lock.readLock().unlock();
        }
    }

    public NavigableMap<Integer, Long> epochWithOffsets() {
        lock.readLock().lock();
        try {
            NavigableMap<Integer, Long> epochWithOffsets = new TreeMap<>();
            for (EpochEntry epochEntry : epochs.values()) {
                epochWithOffsets.put(epochEntry.epoch(), epochEntry.startOffset());
            }
            return epochWithOffsets;
        } finally {
            lock.readLock().unlock();
        }
    }

    private void writeToFile() {
        lock.readLock().lock();
        try {
            checkpoint.write(epochs.values());
        } finally {
            lock.readLock().unlock();
        }
    }

    private void writeIfDirExists() {
        lock.readLock().lock();
        try {
            // If we take a snapshot of the epoch entries here and flush them to disk outside the read lock,
            // by the time of flushing, the leader epoch file may already be updated with newer epoch entries.
            // Those newer entries will then be overridden with the old snapshot.
            checkpoint.writeIfDirExists(epochs.values());
        } finally {
            lock.readLock().unlock();
        }
    }
}
