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

import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.compress.Compression;
import org.apache.kafka.common.errors.CorruptRecordException;
import org.apache.kafka.common.errors.OffsetOutOfRangeException;
import org.apache.kafka.common.record.MemoryRecords;
import org.apache.kafka.common.record.MemoryRecordsBuilder;
import org.apache.kafka.common.record.Record;
import org.apache.kafka.common.record.RecordBatch;
import org.apache.kafka.common.record.Records;
import org.apache.kafka.common.record.SimpleRecord;
import org.apache.kafka.common.record.TimestampType;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.server.common.OffsetAndEpoch;
import org.apache.kafka.snapshot.MockRawSnapshotReader;
import org.apache.kafka.snapshot.MockRawSnapshotWriter;
import org.apache.kafka.snapshot.RawSnapshotReader;
import org.apache.kafka.snapshot.RawSnapshotWriter;

import org.slf4j.Logger;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.OptionalLong;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class MockLog implements ReplicatedLog {
    private static final AtomicLong ID_GENERATOR = new AtomicLong();

    private final List<EpochStartOffset> epochStartOffsets = new ArrayList<>();
    private final List<LogBatch> batches = new ArrayList<>();
    private final NavigableMap<OffsetAndEpoch, MockRawSnapshotReader> snapshots = new TreeMap<>();
    private final TopicPartition topicPartition;
    private final Uuid topicId;
    private final Logger logger;

    private long nextId = ID_GENERATOR.getAndIncrement();
    private LogOffsetMetadata highWatermark = new LogOffsetMetadata(0, Optional.empty());
    private long firstUnflushedOffset = 0;
    private boolean flushedSinceLastChecked = false;

    public MockLog(
        TopicPartition topicPartition,
        Uuid topicId,
        LogContext logContext
    ) {
        this.topicPartition = topicPartition;
        this.topicId = topicId;
        this.logger = logContext.logger(MockLog.class);
    }

    @Override
    public void truncateTo(long offset) {
        if (offset < highWatermark.offset()) {
            throw new IllegalArgumentException("Illegal attempt to truncate to offset " + offset +
                " which is below the current high watermark " + highWatermark);
        }

        logger.debug("Truncating log to end offset {}", offset);
        batches.removeIf(entry -> entry.lastOffset() >= offset);
        epochStartOffsets.removeIf(epochStartOffset -> epochStartOffset.startOffset >= offset);
        firstUnflushedOffset = Math.min(firstUnflushedOffset, endOffset().offset());
    }

    @Override
    public boolean truncateToLatestSnapshot() {
        AtomicBoolean truncated = new AtomicBoolean(false);
        latestSnapshotId().ifPresent(snapshotId -> {
            if (snapshotId.epoch() > logLastFetchedEpoch().orElse(0) ||
                (snapshotId.epoch() == logLastFetchedEpoch().orElse(0) &&
                 snapshotId.offset() > endOffset().offset())) {

                logger.debug("Truncating to the latest snapshot at {}", snapshotId);

                batches.clear();
                epochStartOffsets.clear();
                snapshots.headMap(snapshotId, false).clear();
                updateHighWatermark(new LogOffsetMetadata(snapshotId.offset()));
                flush(false);

                truncated.set(true);
            }
        });

        return truncated.get();
    }

    @Override
    public void updateHighWatermark(LogOffsetMetadata offsetMetadata) {
        if (this.highWatermark.offset() > offsetMetadata.offset()) {
            throw new IllegalArgumentException("Non-monotonic update of current high watermark " +
                highWatermark + " to new value " + offsetMetadata);
        } else if (offsetMetadata.offset() > endOffset().offset()) {
            throw new IllegalArgumentException("Attempt to update high watermark to " + offsetMetadata +
                " which is larger than the current end offset " + endOffset());
        } else if (offsetMetadata.offset() < startOffset()) {
            throw new IllegalArgumentException("Attempt to update high watermark to " + offsetMetadata +
                " which is smaller than the current start offset " + startOffset());
        }

        assertValidHighWatermarkMetadata(offsetMetadata);
        this.highWatermark = offsetMetadata;
    }

    @Override
    public LogOffsetMetadata highWatermark() {
        return highWatermark;
    }

    @Override
    public TopicPartition topicPartition() {
        return topicPartition;
    }

    @Override
    public Uuid topicId() {
        return topicId;
    }

    private Optional<OffsetMetadata> metadataForOffset(long offset) {
        if (offset == endOffset().offset()) {
            return endOffset().metadata();
        }

        for (LogBatch batch : batches) {
            if (batch.lastOffset() < offset)
                continue;

            for (LogEntry entry : batch.entries) {
                if (entry.offset == offset) {
                    return Optional.of(entry.metadata);
                }
            }
        }

        return Optional.empty();
    }

    private void assertValidHighWatermarkMetadata(LogOffsetMetadata offsetMetadata) {
        if (offsetMetadata.metadata().isEmpty()) {
            return;
        }

        long id = ((MockOffsetMetadata) offsetMetadata.metadata().get()).id;
        long offset = offsetMetadata.offset();

        metadataForOffset(offset).ifPresent(metadata -> {
            long entryId = ((MockOffsetMetadata) metadata).id;
            if (entryId != id) {
                throw new IllegalArgumentException("High watermark " + offset +
                    " metadata uuid " + id + " does not match the " +
                    " log's record entry maintained uuid " + entryId);
            }
        });
    }

    private OptionalInt logLastFetchedEpoch() {
        if (epochStartOffsets.isEmpty()) {
            return OptionalInt.empty();
        } else {
            return OptionalInt.of(epochStartOffsets.get(epochStartOffsets.size() - 1).epoch);
        }
    }

    @Override
    public int lastFetchedEpoch() {
        return logLastFetchedEpoch().orElseGet(() -> latestSnapshotId().map(OffsetAndEpoch::epoch).orElse(0));
    }

    @Override
    public OffsetAndEpoch endOffsetForEpoch(int epoch) {
        return lastOffsetAndEpochFiltered(epochStartOffset -> epochStartOffset.epoch <= epoch);
    }

    private OffsetAndEpoch lastOffsetAndEpochFiltered(Predicate<EpochStartOffset> predicate) {
        int epochLowerBound = earliestSnapshotId().map(OffsetAndEpoch::epoch).orElse(0);
        for (EpochStartOffset epochStartOffset : epochStartOffsets) {
            if (!predicate.test(epochStartOffset)) {
                return new OffsetAndEpoch(epochStartOffset.startOffset, epochLowerBound);
            }
            epochLowerBound = epochStartOffset.epoch;
        }

        return new OffsetAndEpoch(endOffset().offset(), lastFetchedEpoch());
    }

    private Optional<LogEntry> lastEntry() {
        if (batches.isEmpty())
            return Optional.empty();
        return Optional.of(batches.get(batches.size() - 1).last());
    }

    private Optional<LogEntry> firstEntry() {
        if (batches.isEmpty())
            return Optional.empty();
        return Optional.of(batches.get(0).first());
    }

    @Override
    public LogOffsetMetadata endOffset() {
        long nextOffset = lastEntry()
            .map(entry -> entry.offset + 1)
            .orElse(
                latestSnapshotId()
                    .map(OffsetAndEpoch::offset)
                    .orElse(0L)
            );
        return new LogOffsetMetadata(nextOffset, Optional.of(new MockOffsetMetadata(nextId)));
    }

    @Override
    public long startOffset() {
        return firstEntry()
            .map(entry -> entry.offset)
            .orElse(
                earliestSnapshotId()
                    .map(OffsetAndEpoch::offset)
                    .orElse(0L)
            );
    }

    private List<LogEntry> buildEntries(RecordBatch batch, Function<Record, Long> offsetSupplier) {
        List<LogEntry> entries = new ArrayList<>();
        for (Record record : batch) {
            long offset = offsetSupplier.apply(record);
            long timestamp = record.timestamp();
            ByteBuffer key = copy(record.key());
            ByteBuffer value = copy(record.value());
            entries.add(buildEntry(offset, new SimpleRecord(timestamp, key, value)));
        }
        return entries;
    }

    private ByteBuffer copy(ByteBuffer nullableByteBuffer) {
        if (nullableByteBuffer == null) {
            return null;
        } else {
            byte[] array = Utils.toArray(nullableByteBuffer, nullableByteBuffer.position(), nullableByteBuffer.limit());
            return ByteBuffer.wrap(array);
        }
    }

    private LogEntry buildEntry(Long offset, SimpleRecord record) {
        long id = nextId;
        nextId = ID_GENERATOR.getAndIncrement();
        return new LogEntry(new MockOffsetMetadata(id), offset, record);
    }


    @Override
    public LogAppendInfo appendAsLeader(Records records, int epoch) {
        return append(records, epoch, true);
    }

    private long appendBatch(LogBatch batch) {
        if (batch.epoch > lastFetchedEpoch()) {
            epochStartOffsets.add(new EpochStartOffset(batch.epoch, batch.firstOffset()));
        }

        batches.add(batch);
        return batch.firstOffset();
    }

    @Override
    public LogAppendInfo appendAsFollower(Records records, int epoch) {
        return append(records, epoch, false);
    }

    private LogAppendInfo append(Records records, int epoch, boolean isLeader) {
        if (records.sizeInBytes() == 0) {
            throw new IllegalArgumentException("Attempt to append an empty record set");
        }

        long baseOffset = endOffset().offset();
        long lastOffset = baseOffset;
        boolean hasBatches = false;
        for (RecordBatch batch : records.batches()) {
            if (batch.baseOffset() != endOffset().offset()) {
                /* KafkaMetadataLog throws an kafka.common.UnexpectedAppendOffsetException this is the
                 * best we can do from this module.
                 */
                throw new RuntimeException(
                    String.format(
                        "Illegal append at offset %s with current end offset of %s",
                        batch.baseOffset(),
                        endOffset().offset()
                    )
                );
            } else if (isLeader && epoch != batch.partitionLeaderEpoch()) {
                // the partition leader epoch is set and does not match the one set in the batch
                throw new RuntimeException(
                    String.format(
                        "Epoch %s doesn't match batch leader epoch %s",
                        epoch,
                        batch.partitionLeaderEpoch()
                    )
                );
            } else if (!isLeader && batch.partitionLeaderEpoch() > epoch) {
                /* To avoid inconsistent log replication, follower should only append record
                 * batches with an epoch less than or equal to the leader epoch. There is more
                 * details on this issue and scenario in KAFKA-18723.
                 */
                break;
            }

            hasBatches = true;
            LogBatch logBatch = new LogBatch(
                batch.partitionLeaderEpoch(),
                batch.isControlBatch(),
                buildEntries(batch, Record::offset)
            );

            if (logger.isDebugEnabled()) {
                logger.debug(
                    "{} appending to the log {}",
                    isLeader ? "Leader" : "Follower",
                    logBatch
                );
            }

            appendBatch(logBatch);
            lastOffset = logBatch.last().offset;
        }

        if (!hasBatches) {
            // This emulates the default handling when records doesn't have enough bytes for a batch
            throw new CorruptRecordException("Append failed unexpectedly");
        }

        return new LogAppendInfo(baseOffset, lastOffset);
    }

    @Override
    public void flush(boolean forceFlushActiveSegment) {
        flushedSinceLastChecked = true;
        firstUnflushedOffset = endOffset().offset();
    }

    @Override
    public boolean maybeClean() {
        return false;
    }

    public long firstUnflushedOffset() {
        return firstUnflushedOffset;
    }

    /**
     * Returns true if the log was flushed since the last time this method was called.
     */
    public boolean flushedSinceLastChecked() {
        boolean oldValue = flushedSinceLastChecked;
        flushedSinceLastChecked = false;
        return oldValue;
    }

    /**
     * Reopening the log causes all unflushed data to be lost.
     */
    public void reopen() {
        batches.removeIf(batch -> batch.firstOffset() >= firstUnflushedOffset);
        epochStartOffsets.removeIf(epochStartOffset -> epochStartOffset.startOffset >= firstUnflushedOffset);
        highWatermark = new LogOffsetMetadata(0L, Optional.empty());
    }

    public List<LogBatch> readBatches(long startOffset, OptionalLong maxOffsetOpt) {
        verifyOffsetInRange(startOffset);

        long maxOffset = maxOffsetOpt.orElse(endOffset().offset());
        if (startOffset == maxOffset) {
            return List.of();
        }

        return batches.stream()
            .filter(batch -> batch.lastOffset() >= startOffset && batch.lastOffset() < maxOffset)
            .collect(Collectors.toList());
    }

    private void verifyOffsetInRange(long offset) {
        if (offset > endOffset().offset()) {
            throw new OffsetOutOfRangeException("Requested offset " + offset + " is larger than " +
                "then log end offset " + endOffset().offset());
        }

        if (offset < this.startOffset()) {
            throw new OffsetOutOfRangeException("Requested offset " + offset + " is smaller than " +
                "then log start offset " + this.startOffset());
        }
    }

    @Override
    public LogFetchInfo read(long startOffset, Isolation isolation) {
        verifyOffsetInRange(startOffset);

        long maxOffset = isolation == Isolation.COMMITTED ? highWatermark.offset() : endOffset().offset();
        if (startOffset >= maxOffset) {
            return new LogFetchInfo(MemoryRecords.EMPTY, new LogOffsetMetadata(
                startOffset, metadataForOffset(startOffset)));
        }

        ByteBuffer buffer = ByteBuffer.allocate(512);
        int batchCount = 0;
        LogOffsetMetadata batchStartOffset = null;

        logger.debug(
            "Looking for a batch that starts at {} and ends at {} for isolation {}",
            startOffset,
            maxOffset,
            isolation
        );

        for (LogBatch batch : batches) {
            // Note that start offset is inclusive while max offset is exclusive. We only return
            // complete batches, so batches which end at an offset larger than the max offset are
            // filtered, which is effectively the same as having the consumer drop an incomplete
            // batch returned in a fetch response.
            if (batch.lastOffset() >= startOffset && batch.lastOffset() < maxOffset && !batch.entries.isEmpty()) {
                buffer = batch.writeTo(buffer);

                if (batchStartOffset == null) {
                    batchStartOffset = batch.entries.get(0).logOffsetMetadata();
                }

                // Read on the mock log should return at most 2 batches. This is a simple solution
                // for testing interesting partial read scenarios.
                batchCount += 1;
                if (batchCount >= 2) {
                    break;
                }
            }
        }

        buffer.flip();
        Records records = MemoryRecords.readableRecords(buffer);

        if (batchStartOffset == null) {
            throw new RuntimeException("Expected to find at least one entry starting from offset " +
                startOffset + " but found none");
        }

        return new LogFetchInfo(records, batchStartOffset);
    }

    @Override
    public void initializeLeaderEpoch(int epoch) {
        long startOffset = endOffset().offset();
        epochStartOffsets.removeIf(epochStartOffset ->
            epochStartOffset.startOffset >= startOffset || epochStartOffset.epoch >= epoch);
        epochStartOffsets.add(new EpochStartOffset(epoch, startOffset));
    }

    @Override
    public Optional<RawSnapshotWriter> createNewSnapshot(OffsetAndEpoch snapshotId) {
        if (snapshotId.offset() < startOffset()) {
            logger.info(
                "Cannot create a snapshot with an id ({}) less than the log start offset ({})",
                snapshotId,
                startOffset()
            );

            return Optional.empty();
        }

        long highWatermarkOffset = highWatermark().offset();
        if (snapshotId.offset() > highWatermarkOffset) {
            throw new IllegalArgumentException(
                String.format(
                    "Cannot create a snapshot with an id (%s) greater than the high-watermark (%s)",
                    snapshotId,
                    highWatermarkOffset
                )
            );
        }

        ValidOffsetAndEpoch validOffsetAndEpoch = validateOffsetAndEpoch(snapshotId.offset(), snapshotId.epoch());
        if (validOffsetAndEpoch.kind() != ValidOffsetAndEpoch.Kind.VALID) {
            throw new IllegalArgumentException(
                String.format(
                    "Snapshot id (%s) is not valid according to the log: %s",
                    snapshotId,
                    validOffsetAndEpoch
                )
            );
        }

        long baseOffset = read(snapshotId.offset(), Isolation.COMMITTED).startOffsetMetadata.offset();
        if (snapshotId.offset() != baseOffset) {
            throw new IllegalArgumentException(
                String.format(
                    "Cannot create snapshot at offset (%s) because it is not batch aligned. " +
                    "The batch containing the requested offset has a base offset of (%s)",
                    snapshotId.offset(),
                    baseOffset
                )
            );
        }

        return createNewSnapshotUnchecked(snapshotId);
    }

    @Override
    public Optional<RawSnapshotWriter> createNewSnapshotUnchecked(OffsetAndEpoch snapshotId) {
        if (snapshots.containsKey(snapshotId)) {
            return Optional.empty();
        } else {
            return Optional.of(
                new MockRawSnapshotWriter(snapshotId, buffer ->
                    snapshots.putIfAbsent(snapshotId, new MockRawSnapshotReader(snapshotId, buffer))
                )
            );
        }
    }

    @Override
    public Optional<RawSnapshotReader> readSnapshot(OffsetAndEpoch snapshotId) {
        return Optional.ofNullable(snapshots.get(snapshotId));
    }

    @Override
    public Optional<RawSnapshotReader> latestSnapshot() {
        return latestSnapshotId().flatMap(this::readSnapshot);
    }

    @Override
    public Optional<OffsetAndEpoch> latestSnapshotId() {
        return Optional.ofNullable(snapshots.lastEntry())
            .map(Map.Entry::getKey);
    }

    @Override
    public Optional<OffsetAndEpoch> earliestSnapshotId() {
        return Optional.ofNullable(snapshots.firstEntry())
            .map(Map.Entry::getKey);
    }

    @Override
    public void onSnapshotFrozen(OffsetAndEpoch snapshotId) {}

    @Override
    public boolean deleteBeforeSnapshot(OffsetAndEpoch snapshotId) {
        if (startOffset() > snapshotId.offset()) {
            throw new OffsetOutOfRangeException(
                String.format(
                    "New log start (%s) is less than the current log start offset (%s)",
                    snapshotId,
                    startOffset()
                )
            );
        }
        if (highWatermark.offset() < snapshotId.offset()) {
            throw new OffsetOutOfRangeException(
                String.format(
                    "New log start (%s) is greater than the high watermark (%s)",
                    snapshotId,
                    highWatermark.offset()
                )
            );
        }

        boolean updated = false;
        if (snapshots.containsKey(snapshotId)) {
            snapshots.headMap(snapshotId, false).clear();

            logger.debug("Deleting batches included in the snapshot {}", snapshotId);
            batches.removeIf(entry -> entry.lastOffset() < snapshotId.offset());

            AtomicReference<Optional<EpochStartOffset>> last = new AtomicReference<>(Optional.empty());
            epochStartOffsets.removeIf(epochStartOffset -> {
                if (epochStartOffset.startOffset <= snapshotId.offset()) {
                    last.set(Optional.of(epochStartOffset));
                    return true;
                }

                return false;
            });

            last.get().ifPresent(epochStartOffset ->
                epochStartOffsets.add(0, new EpochStartOffset(epochStartOffset.epoch, snapshotId.offset())
                )
            );

            updated = true;
        }

        return updated;
    }

    @Override
    public String toString() {
        return String.format(
            "MockLog(epochStartOffsets=%s, batches=%s, snapshots=%s, highWatermark=%s",
            epochStartOffsets,
            batches,
            snapshots,
            highWatermark
        );
    }

    record MockOffsetMetadata(long id) implements OffsetMetadata {
    }

    record LogEntry(MockOffsetMetadata metadata, long offset, SimpleRecord record) {

        LogOffsetMetadata logOffsetMetadata() {
            return new LogOffsetMetadata(offset, Optional.of(metadata));
        }
    }

    record LogBatch(int epoch, boolean isControlBatch, List<LogEntry> entries) {
        LogBatch {
            if (entries.isEmpty())
                throw new IllegalArgumentException("Empty batches are not supported");
        }

        long firstOffset() {
            return first().offset;
        }

        LogEntry first() {
            return entries.get(0);
        }

        long lastOffset() {
            return last().offset;
        }

        LogEntry last() {
            return entries.get(entries.size() - 1);
        }

        ByteBuffer writeTo(ByteBuffer buffer) {
            LogEntry first = first();

            MemoryRecordsBuilder builder = MemoryRecords.builder(
                buffer, RecordBatch.CURRENT_MAGIC_VALUE, Compression.NONE,
                TimestampType.CREATE_TIME, first.offset, first.record.timestamp(),
                RecordBatch.NO_PRODUCER_ID, RecordBatch.NO_PRODUCER_EPOCH,
                RecordBatch.NO_SEQUENCE, false,
                isControlBatch, epoch);

            for (LogEntry entry : entries) {
                if (isControlBatch) {
                    builder.appendControlRecordWithOffset(entry.offset, entry.record);
                } else {
                    builder.appendWithOffset(entry.offset, entry.record);
                }
            }

            builder.close();
            return builder.buffer();
        }
    }

    private record EpochStartOffset(int epoch, long startOffset) {
    }
}
