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

import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.metadata.AbortTransactionRecord;
import org.apache.kafka.common.metadata.BeginTransactionRecord;
import org.apache.kafka.common.metadata.EndTransactionRecord;
import org.apache.kafka.common.metadata.NoOpRecord;
import org.apache.kafka.common.metadata.PartitionRecord;
import org.apache.kafka.common.metadata.TopicRecord;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.image.MetadataDelta;
import org.apache.kafka.image.MetadataImage;
import org.apache.kafka.raft.Batch;
import org.apache.kafka.raft.LeaderAndEpoch;
import org.apache.kafka.server.common.ApiMessageAndVersion;
import org.apache.kafka.server.fault.MockFaultHandler;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;

import java.util.ArrayList;
import java.util.List;
import java.util.OptionalInt;
import java.util.stream.IntStream;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class MetadataBatchLoaderTest {

    static final Uuid TOPIC_FOO = Uuid.fromString("c6uHMgPkRp2Urjlh-RxMNQ");
    static final Uuid TOPIC_BAR = Uuid.fromString("tUWOOPvzQhmZZ_eXmTCcig");
    static final List<ApiMessageAndVersion> TOPIC_TXN_BATCH_1;
    static final List<ApiMessageAndVersion> TOPIC_TXN_BATCH_2;
    static final List<ApiMessageAndVersion> TOPIC_NO_TXN_BATCH;
    static final List<ApiMessageAndVersion> TXN_BEGIN_SINGLETON;
    static final List<ApiMessageAndVersion> TXN_END_SINGLETON;
    static final List<ApiMessageAndVersion> TXN_ABORT_SINGLETON;
    static final LeaderAndEpoch LEADER_AND_EPOCH = new LeaderAndEpoch(OptionalInt.of(1), 42);

    static {
        {
            TOPIC_TXN_BATCH_1 = List.of(
                new ApiMessageAndVersion(new BeginTransactionRecord().setName("txn-1"), (short) 0),
                new ApiMessageAndVersion(new TopicRecord()
                    .setName("foo")
                    .setTopicId(TOPIC_FOO), (short) 0),
                new ApiMessageAndVersion(new PartitionRecord()
                    .setPartitionId(0)
                    .setTopicId(TOPIC_FOO), (short) 0)
            );

            TOPIC_TXN_BATCH_2 = List.of(
                new ApiMessageAndVersion(new PartitionRecord()
                    .setPartitionId(1)
                    .setTopicId(TOPIC_FOO), (short) 0),
                new ApiMessageAndVersion(new PartitionRecord()
                    .setPartitionId(2)
                    .setTopicId(TOPIC_FOO), (short) 0),
                new ApiMessageAndVersion(new EndTransactionRecord(), (short) 0)
            );

            TOPIC_NO_TXN_BATCH = List.of(
                new ApiMessageAndVersion(new TopicRecord()
                    .setName("bar")
                    .setTopicId(TOPIC_BAR), (short) 0),
                new ApiMessageAndVersion(new PartitionRecord()
                    .setPartitionId(0)
                    .setTopicId(TOPIC_BAR), (short) 0),
                new ApiMessageAndVersion(new PartitionRecord()
                    .setPartitionId(1)
                    .setTopicId(TOPIC_BAR), (short) 0)
            );

            TXN_BEGIN_SINGLETON = List.of(
                new ApiMessageAndVersion(new BeginTransactionRecord().setName("txn-1"), (short) 0));

            TXN_END_SINGLETON = List.of(
                new ApiMessageAndVersion(new EndTransactionRecord(), (short) 0));

            TXN_ABORT_SINGLETON = List.of(
                new ApiMessageAndVersion(new AbortTransactionRecord(), (short) 0));
        }
    }

    static List<ApiMessageAndVersion> noOpRecords(int n) {
        return IntStream.range(0, n)
                .mapToObj(__ -> new ApiMessageAndVersion(new NoOpRecord(), (short) 0))
                .toList();
    }


    static class MockMetadataUpdater implements MetadataBatchLoader.MetadataUpdater {
        MetadataImage latestImage = null;
        MetadataDelta latestDelta = null;
        LogDeltaManifest latestManifest = null;
        int updates = 0;

        @Override
        public void update(MetadataDelta delta, MetadataImage image, LogDeltaManifest manifest) {
            latestDelta = delta;
            latestImage = image;
            latestManifest = manifest;
            updates++;
        }

        public void reset() {
            latestImage = null;
            latestDelta = null;
            latestManifest = null;
            updates = 0;
        }
    }

    @Test
    public void testAlignedTransactionBatches() {
        Batch<ApiMessageAndVersion> batch1 = Batch.data(
            10, 1, 0, 10, TOPIC_TXN_BATCH_1);
        Batch<ApiMessageAndVersion> batch2 = Batch.data(
            13, 2, 0, 10, noOpRecords(3));
        Batch<ApiMessageAndVersion> batch3 = Batch.data(
            16, 2, 0, 30, TOPIC_TXN_BATCH_2);

        MockMetadataUpdater updater = new MockMetadataUpdater();
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            new MockFaultHandler("testAlignedTransactionBatches"),
            updater
        );

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        batchLoader.loadBatch(batch3, LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);
        assertNotNull(updater.latestImage.topics().getTopic("foo"));
        assertEquals(18, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(2, updater.latestImage.provenance().lastContainedEpoch());
        assertTrue(updater.latestImage.provenance().isOffsetBatchAligned());
    }

    @Test
    public void testSingletonBeginAndEnd() {
        Batch<ApiMessageAndVersion> batch1 = Batch.data(
            13, 1, 0, 30, noOpRecords(3));

        Batch<ApiMessageAndVersion> batch2 = Batch.data(
            16, 2, 0, 30, TXN_BEGIN_SINGLETON);

        Batch<ApiMessageAndVersion> batch3 = Batch.data(
            17, 3, 0, 10, TOPIC_NO_TXN_BATCH);

        Batch<ApiMessageAndVersion> batch4 = Batch.data(
            20, 4, 0, 10, TXN_END_SINGLETON);
        MockMetadataUpdater updater = new MockMetadataUpdater();
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            new MockFaultHandler("testSingletonBeginAndEnd"),
            updater
        );

        // All in one commit
        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        // batch1 is flushed in this loadBatch call
        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        assertEquals(1, updater.updates);
        assertTrue(updater.latestImage.provenance().isOffsetBatchAligned());
        assertNull(updater.latestImage.topics().getTopic("bar"));
        batchLoader.loadBatch(batch3, LEADER_AND_EPOCH);
        assertEquals(1, updater.updates);
        batchLoader.loadBatch(batch4, LEADER_AND_EPOCH);
        assertEquals(1, updater.updates);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertNotNull(updater.latestImage.topics().getTopic("bar"));
        assertEquals(20, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(4, updater.latestImage.provenance().lastContainedEpoch());

        // Each batch in a separate commit
        updater.reset();
        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);

        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);

        batchLoader.loadBatch(batch3, LEADER_AND_EPOCH);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);

        batchLoader.loadBatch(batch4, LEADER_AND_EPOCH);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(2, updater.updates);
    }

    @Test
    public void testUnexpectedBeginTransaction() {
        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testUnexpectedBeginTransaction");
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            faultHandler,
            updater
        );

        Batch<ApiMessageAndVersion> batch1 = Batch.data(
            10, 2, 0, 30, TOPIC_TXN_BATCH_1);

        Batch<ApiMessageAndVersion> batch2 = Batch.data(
            13, 2, 0, 30, TXN_BEGIN_SINGLETON);

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        assertNull(faultHandler.firstException());
        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        assertEquals(RuntimeException.class, faultHandler.firstException().getCause().getClass());
        assertEquals(
            "Encountered BeginTransactionRecord while already in a transaction",
            faultHandler.firstException().getCause().getMessage()
        );
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(0, updater.updates);
    }

    @Test
    public void testUnexpectedEndTransaction() {
        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testUnexpectedAbortTransaction");
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
                new LogContext(),
                new MockTime(),
                faultHandler,
                updater
        );

        // First batch gets loaded fine
        Batch<ApiMessageAndVersion> batch1 = Batch.data(
            10, 2, 0, 30, TOPIC_NO_TXN_BATCH);

        // Second batch throws an error, but shouldn't interfere with prior batches
        Batch<ApiMessageAndVersion> batch2 = Batch.data(
            13, 2, 0, 30, TXN_END_SINGLETON);

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        assertNull(faultHandler.firstException());
        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        assertEquals(RuntimeException.class, faultHandler.firstException().getCause().getClass());
        assertEquals(
            "Encountered EndTransactionRecord without having seen a BeginTransactionRecord",
            faultHandler.firstException().getCause().getMessage()
        );
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);
        assertNotNull(updater.latestImage.topics().getTopic("bar"));
    }

    @Test
    public void testUnexpectedAbortTransaction() {
        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testUnexpectedAbortTransaction");
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            faultHandler,
            updater
        );

        // First batch gets loaded fine
        Batch<ApiMessageAndVersion> batch1 = Batch.data(
            10, 2, 0, 30, TOPIC_NO_TXN_BATCH);

        // Second batch throws an error, but shouldn't interfere with prior batches
        Batch<ApiMessageAndVersion> batch2 = Batch.data(
            13, 2, 0, 30, TXN_ABORT_SINGLETON);

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch1, LEADER_AND_EPOCH);
        assertNull(faultHandler.firstException());
        batchLoader.loadBatch(batch2, LEADER_AND_EPOCH);
        assertEquals(RuntimeException.class, faultHandler.firstException().getCause().getClass());
        assertEquals(
            "Encountered AbortTransactionRecord without having seen a BeginTransactionRecord",
            faultHandler.firstException().getCause().getMessage()
        );
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(1, updater.updates);
        assertNotNull(updater.latestImage.topics().getTopic("bar"));
    }

    private MetadataBatchLoader loadSingleBatch(
        MockMetadataUpdater updater,
        MockFaultHandler faultHandler,
        List<ApiMessageAndVersion> batchRecords
    ) {
        Batch<ApiMessageAndVersion> batch = Batch.data(
            10, 42, 0, 100, batchRecords);

        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            faultHandler,
            updater
        );

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(batch, LEADER_AND_EPOCH);
        return batchLoader;
    }

    @Test
    public void testMultipleTransactionsInOneBatch() {
        List<ApiMessageAndVersion> batchRecords = new ArrayList<>();
        batchRecords.addAll(TOPIC_TXN_BATCH_1);
        batchRecords.addAll(TOPIC_TXN_BATCH_2);
        batchRecords.addAll(TXN_BEGIN_SINGLETON);
        batchRecords.addAll(TOPIC_NO_TXN_BATCH);
        batchRecords.addAll(TXN_END_SINGLETON);

        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testMultipleTransactionsInOneBatch");
        MetadataBatchLoader batchLoader = loadSingleBatch(updater, faultHandler, batchRecords);

        assertEquals(1, updater.updates);
        assertEquals(0, updater.latestManifest.numBytes());
        assertEquals(15, updater.latestImage.provenance().lastContainedOffset());
        // The first transaction is flushed in the middle of the batch, the offset flushed is not batch-aligned
        assertFalse(updater.latestImage.provenance().isOffsetBatchAligned());
        assertEquals(42, updater.latestImage.provenance().lastContainedEpoch());

        assertNotNull(updater.latestImage.topics().getTopic("foo"));
        assertNull(updater.latestImage.topics().getTopic("bar"));
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(2, updater.updates);
        assertEquals(100, updater.latestManifest.numBytes());
        assertEquals(20, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(42, updater.latestImage.provenance().lastContainedEpoch());
        assertTrue(updater.latestImage.provenance().isOffsetBatchAligned());
        assertNotNull(updater.latestImage.topics().getTopic("foo"));
        assertNotNull(updater.latestImage.topics().getTopic("bar"));
    }

    @Test
    public void testMultipleTransactionsInOneBatchesWithNoOp() {
        List<ApiMessageAndVersion> batchRecords = new ArrayList<>();
        batchRecords.addAll(noOpRecords(1));
        batchRecords.addAll(TOPIC_TXN_BATCH_1);
        batchRecords.addAll(noOpRecords(1));
        batchRecords.addAll(TOPIC_TXN_BATCH_2);
        // A batch with non-transactional records between two transactions causes a delta to get published
        batchRecords.addAll(noOpRecords(1));
        batchRecords.addAll(TXN_BEGIN_SINGLETON);
        batchRecords.addAll(noOpRecords(1));
        batchRecords.addAll(TOPIC_NO_TXN_BATCH);
        batchRecords.addAll(noOpRecords(1));
        batchRecords.addAll(TXN_END_SINGLETON);
        batchRecords.addAll(noOpRecords(1));

        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testMultipleTransactionsInOneBatches");
        MetadataBatchLoader batchLoader = loadSingleBatch(updater, faultHandler, batchRecords);

        assertEquals(2, updater.updates);
        assertEquals(0, updater.latestManifest.numBytes());
        assertEquals(18, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(42, updater.latestImage.provenance().lastContainedEpoch());
        assertFalse(updater.latestImage.provenance().isOffsetBatchAligned());
        assertNotNull(updater.latestImage.topics().getTopic("foo"));
        assertNull(updater.latestImage.topics().getTopic("bar"));
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);
        assertEquals(3, updater.updates);
        assertEquals(100, updater.latestManifest.numBytes());
        assertEquals(26, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(42, updater.latestImage.provenance().lastContainedEpoch());
        assertNotNull(updater.latestImage.topics().getTopic("foo"));
        assertNotNull(updater.latestImage.topics().getTopic("bar"));
    }

    @ParameterizedTest
    @ValueSource(booleans = {true, false})
    public void testOneTransactionInMultipleBatches(boolean abortTxn) {
        MockMetadataUpdater updater = new MockMetadataUpdater();
        MetadataBatchLoader batchLoader = new MetadataBatchLoader(
            new LogContext(),
            new MockTime(),
            new MockFaultHandler("testOneTransactionInMultipleBatches"),
            updater
        );

        batchLoader.resetToImage(MetadataImage.EMPTY);
        batchLoader.loadBatch(Batch.data(
            16, 2, 0, 10, TXN_BEGIN_SINGLETON), LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        batchLoader.loadBatch(Batch.data(
            17, 3, 0, 30, TOPIC_NO_TXN_BATCH), LEADER_AND_EPOCH);
        assertEquals(0, updater.updates);
        if (abortTxn) {
            batchLoader.loadBatch(Batch.data(
                20, 4, 0, 10, TXN_ABORT_SINGLETON), LEADER_AND_EPOCH);
        } else {
            batchLoader.loadBatch(Batch.data(
                20, 4, 0, 10, TXN_END_SINGLETON), LEADER_AND_EPOCH);
        }
        assertEquals(0, updater.updates);
        batchLoader.maybeFlushBatches(LEADER_AND_EPOCH, true);

        // Regardless of end/abort, we should publish an updated MetadataProvenance and manifest
        assertEquals(50, updater.latestManifest.numBytes());
        assertEquals(3, updater.latestManifest.numBatches());
        assertEquals(20, updater.latestImage.provenance().lastContainedOffset());
        assertEquals(4, updater.latestImage.provenance().lastContainedEpoch());
        if (abortTxn) {
            assertNull(updater.latestImage.topics().getTopic("bar"));
        } else {
            assertNotNull(updater.latestImage.topics().getTopic("bar"));
        }
    }

    @Test
    public void testTransactionAlignmentOnBatchBoundary() {
        List<ApiMessageAndVersion> batchRecords = new ArrayList<>();
        batchRecords.addAll(noOpRecords(3));
        batchRecords.addAll(TOPIC_TXN_BATCH_1);
        batchRecords.addAll(TOPIC_TXN_BATCH_2);
        batchRecords.addAll(noOpRecords(3));

        MockMetadataUpdater updater = new MockMetadataUpdater();
        MockFaultHandler faultHandler = new MockFaultHandler("testMultipleTransactionsInOneBatch");
        MetadataBatchLoader batchLoader = loadSingleBatch(updater, faultHandler, batchRecords);

        assertEquals(1, updater.updates);
        assertEquals(0, updater.latestManifest.numBytes());
        assertEquals(12, updater.latestImage.provenance().lastContainedOffset());
        assertFalse(updater.latestImage.provenance().isOffsetBatchAligned());

        batchLoader.loadBatch(Batch.data(
                22, 42, 0, 10, TXN_BEGIN_SINGLETON), LEADER_AND_EPOCH);
        assertEquals(2, updater.updates);
        assertEquals(100, updater.latestManifest.numBytes());
        assertEquals(21, updater.latestImage.provenance().lastContainedOffset());
        assertTrue(updater.latestImage.provenance().isOffsetBatchAligned());
    }
}
