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

import org.apache.kafka.common.message.SnapshotFooterRecord;
import org.apache.kafka.common.message.SnapshotHeaderRecord;
import org.apache.kafka.common.record.ControlRecordUtils;
import org.apache.kafka.common.record.Record;
import org.apache.kafka.common.record.RecordBatch;
import org.apache.kafka.common.utils.BufferSupplier;
import org.apache.kafka.common.utils.BufferSupplier.GrowableBufferSupplier;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.raft.Batch;
import org.apache.kafka.raft.ControlRecord;
import org.apache.kafka.raft.RaftClientTestContext;
import org.apache.kafka.raft.internals.StringSerde;
import org.apache.kafka.server.common.OffsetAndEpoch;

import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.Random;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public final class SnapshotWriterReaderTest {
    private final int localId = 0;
    private final Set<Integer> voters = Set.of(localId);

    @Test
    public void testSnapshotDelimiters() throws Exception {
        int recordsPerBatch = 1;
        int batches = 0;
        int delimiterCount = 2;
        long magicTimestamp = 0xDEADBEEF;
        OffsetAndEpoch id = new OffsetAndEpoch(recordsPerBatch * batches, 3);

        RaftClientTestContext.Builder contextBuilder = new RaftClientTestContext.Builder(localId, voters);
        RaftClientTestContext context = contextBuilder.build();

        context.pollUntil(() -> context.currentLeader().equals(OptionalInt.of(localId)));
        context.advanceLocalLeaderHighWatermarkToLogEndOffset();

        // Create an empty snapshot and freeze it immediately
        try (SnapshotWriter<String> snapshot = context.client.createSnapshot(id, magicTimestamp).get()) {
            assertEquals(id, snapshot.snapshotId());
            snapshot.freeze();
        }

        // Verify that an empty snapshot has only the Header and Footer
        try (SnapshotReader<String> reader = readSnapshot(context, id, Integer.MAX_VALUE)) {
            assertEquals(magicTimestamp, reader.lastContainedLogTimestamp());

            RawSnapshotReader snapshot = context.log.readSnapshot(id).get();
            int recordCount = validateDelimiters(snapshot, magicTimestamp);
            assertEquals((recordsPerBatch * batches) + delimiterCount, recordCount);
        }
    }

    @Test
    public void testWritingSnapshot() throws Exception {
        int recordsPerBatch = 3;
        int batches = 3;
        int delimiterCount = 2;
        long magicTimestamp = 0xDEADBEEF;
        OffsetAndEpoch id = new OffsetAndEpoch(recordsPerBatch * batches, 3);
        List<List<String>> expected = buildRecords(recordsPerBatch, batches);

        RaftClientTestContext.Builder contextBuilder = new RaftClientTestContext.Builder(localId, voters);
        for (List<String> batch : expected) {
            contextBuilder.appendToLog(id.epoch(), batch);
        }
        RaftClientTestContext context = contextBuilder.build();

        context.pollUntil(() -> context.currentLeader().equals(OptionalInt.of(localId)));

        context.advanceLocalLeaderHighWatermarkToLogEndOffset();

        try (SnapshotWriter<String> snapshot = context.client.createSnapshot(id, magicTimestamp).get()) {
            assertEquals(id, snapshot.snapshotId());
            expected.forEach(batch -> assertDoesNotThrow(() -> snapshot.append(batch)));
            snapshot.freeze();
        }

        try (SnapshotReader<String> reader = readSnapshot(context, id, Integer.MAX_VALUE)) {
            RawSnapshotReader snapshot = context.log.readSnapshot(id).get();
            int recordCount = validateDelimiters(snapshot, magicTimestamp);
            assertEquals((recordsPerBatch * batches) + delimiterCount, recordCount);
            assertDataSnapshot(expected, reader);

            assertEquals(magicTimestamp, Snapshots.lastContainedLogTimestamp(snapshot, new LogContext()));
        }
    }

    @Test
    public void testAbortedSnapshot() throws Exception {
        int recordsPerBatch = 3;
        int batches = 3;
        OffsetAndEpoch id = new OffsetAndEpoch(recordsPerBatch * batches, 3);
        List<List<String>> expected = buildRecords(recordsPerBatch, batches);

        RaftClientTestContext.Builder contextBuilder = new RaftClientTestContext.Builder(localId, voters);
        for (List<String> batch : expected) {
            contextBuilder.appendToLog(id.epoch(), batch);
        }
        RaftClientTestContext context = contextBuilder.build();

        context.pollUntil(() -> context.currentLeader().equals(OptionalInt.of(localId)));

        context.advanceLocalLeaderHighWatermarkToLogEndOffset();

        try (SnapshotWriter<String> snapshot = context.client.createSnapshot(id, 0).get()) {
            assertEquals(id, snapshot.snapshotId());
            expected.forEach(batch ->
                assertDoesNotThrow(() -> snapshot.append(batch))
            );
        }

        assertEquals(Optional.empty(), context.log.readSnapshot(id));
    }

    @Test
    public void testAppendToFrozenSnapshot() throws Exception {
        int recordsPerBatch = 3;
        int batches = 3;
        OffsetAndEpoch id = new OffsetAndEpoch(recordsPerBatch * batches, 3);
        List<List<String>> expected = buildRecords(recordsPerBatch, batches);

        RaftClientTestContext.Builder contextBuilder = new RaftClientTestContext.Builder(localId, voters);
        for (List<String> batch : expected) {
            contextBuilder.appendToLog(id.epoch(), batch);
        }
        RaftClientTestContext context = contextBuilder.build();

        context.pollUntil(() -> context.currentLeader().equals(OptionalInt.of(localId)));

        context.advanceLocalLeaderHighWatermarkToLogEndOffset();

        try (SnapshotWriter<String> snapshot = context.client.createSnapshot(id, 0).get()) {
            assertEquals(id, snapshot.snapshotId());
            expected.forEach(batch ->
                assertDoesNotThrow(() -> snapshot.append(batch))
            );

            snapshot.freeze();

            assertThrows(RuntimeException.class, () -> snapshot.append(expected.get(0)));
        }
    }

    private List<List<String>> buildRecords(int recordsPerBatch, int batches) {
        Random random = new Random(0);
        List<List<String>> result = new ArrayList<>(batches);
        for (int i = 0; i < batches; i++) {
            List<String> batch = new ArrayList<>(recordsPerBatch);
            for (int j = 0; j < recordsPerBatch; j++) {
                batch.add(String.valueOf(random.nextInt()));
            }
            result.add(batch);
        }

        return result;
    }

    private SnapshotReader<String> readSnapshot(
        RaftClientTestContext context,
        OffsetAndEpoch snapshotId,
        int maxBatchSize
    ) {
        return RecordsSnapshotReader.of(
            context.log.readSnapshot(snapshotId).get(),
            context.serde,
            BufferSupplier.create(),
            maxBatchSize,
            true,
            new LogContext()
        );
    }

    private int validateDelimiters(
        RawSnapshotReader snapshot,
        long lastContainedLogTime
    ) {
        assertNotEquals(0, snapshot.sizeInBytes());

        int countRecords = 0;

        var recordBatches = snapshot.records().batchIterator();

        assertTrue(recordBatches.hasNext());
        RecordBatch batch = recordBatches.next();

        Iterator<Record> records = batch.streamingIterator(new GrowableBufferSupplier());

        // Verify existence of the header record
        assertTrue(batch.isControlBatch());
        assertTrue(records.hasNext());
        Record record = records.next();
        countRecords += 1;

        SnapshotHeaderRecord headerRecord = ControlRecordUtils.deserializeSnapshotHeaderRecord(record);
        assertEquals(ControlRecordUtils.SNAPSHOT_HEADER_CURRENT_VERSION, headerRecord.version());
        assertEquals(headerRecord.lastContainedLogTimestamp(), lastContainedLogTime);

        assertFalse(records.hasNext());

        // Loop over remaining records
        while (recordBatches.hasNext()) {
            batch = recordBatches.next();
            records = batch.streamingIterator(new GrowableBufferSupplier());

            while (records.hasNext()) {
                countRecords += 1;
                record = records.next();
            }
        }

        // Verify existence of the footer record in the end
        assertTrue(batch.isControlBatch());

        SnapshotFooterRecord footerRecord = ControlRecordUtils.deserializeSnapshotFooterRecord(record);
        assertEquals(ControlRecordUtils.SNAPSHOT_FOOTER_CURRENT_VERSION, footerRecord.version());

        return countRecords;
    }

    public static void assertDataSnapshot(List<List<String>> batches, RawSnapshotReader reader) {
        assertDataSnapshot(
            batches,
            RecordsSnapshotReader.of(
                reader,
                new StringSerde(),
                BufferSupplier.create(),
                Integer.MAX_VALUE,
                true,
                new LogContext()
            )
        );
    }

    public static void assertDataSnapshot(
        List<List<String>> batches,
        SnapshotReader<String> reader
    ) {
        List<String> expected = new ArrayList<>();
        batches.forEach(expected::addAll);

        List<String> actual = new ArrayList<>(expected.size());
        while (reader.hasNext()) {
            Batch<String> batch = reader.next();
            for (String value : batch) {
                actual.add(value);
            }
        }

        assertEquals(expected, actual);
    }

    public static void assertControlSnapshot(
        List<List<ControlRecord>> expectedBatches,
        SnapshotReader<?> reader
    ) {
        List<List<ControlRecord>> actualBatches = new ArrayList<>(expectedBatches.size());
        while (reader.hasNext()) {
            Batch<?> batch = reader.next();
            if (!batch.controlRecords().isEmpty()) {
                actualBatches.add(batch.controlRecords());
            }
        }

        assertEquals(expectedBatches, actualBatches);
    }
}
