/*
 * 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.paimon.globalindex.sorted;

import org.apache.paimon.CoreOptions;
import org.apache.paimon.Snapshot;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.BlobData;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.globalindex.GlobalIndexSingleColumnWriter;
import org.apache.paimon.globalindex.KeySerializer;
import org.apache.paimon.globalindex.ResultEntry;
import org.apache.paimon.globalindex.btree.BTreeIndexOptions;
import org.apache.paimon.index.GlobalIndexMeta;
import org.apache.paimon.index.IndexFileHandler;
import org.apache.paimon.index.IndexFileMeta;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.manifest.IndexManifestEntry;
import org.apache.paimon.manifest.ManifestEntry;
import org.apache.paimon.memory.MemorySlice;
import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.schema.Schema;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.table.TableTestBase;
import org.apache.paimon.table.sink.BatchTableCommit;
import org.apache.paimon.table.sink.BatchTableWrite;
import org.apache.paimon.table.sink.BatchWriteBuilder;
import org.apache.paimon.table.sink.CommitMessage;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.Pair;

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

import java.io.Closeable;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/** Test class for {@link SortedGlobalIndexBuilder}. */
public class SortedGlobalIndexBuilderTest extends TableTestBase {

    private static final long PART_ROW_NUM = 1000L;
    private static final KeySerializer KEY_SERIALIZER = KeySerializer.create(DataTypes.INT());
    private static final Comparator<Object> COMPARATOR = KEY_SERIALIZER.createComparator();

    @Test
    public void testSingleColumnWriterRotationPreservesResultGroups() throws Exception {
        GlobalIndexSingleColumnWriter first = mock(GlobalIndexSingleColumnWriter.class);
        GlobalIndexSingleColumnWriter second = mock(GlobalIndexSingleColumnWriter.class);
        when(first.finish())
                .thenReturn(Collections.singletonList(new ResultEntry("index-1", 2, null)));
        when(second.finish())
                .thenReturn(Collections.singletonList(new ResultEntry("index-2", 1, null)));
        Queue<GlobalIndexSingleColumnWriter> writers =
                new ArrayDeque<>(Arrays.asList(first, second));
        SortedSingleColumnIndexWriter rotatingWriter =
                new SortedSingleColumnIndexWriter(2, writers::remove);

        rotatingWriter.write(10, 0);
        rotatingWriter.write(20, 1);
        rotatingWriter.write(30, 2);
        List<List<ResultEntry>> results = rotatingWriter.finish();

        verify(first).write(10, 0);
        verify(first).write(20, 1);
        verify(second).write(30, 2);
        assertThat(results).hasSize(2);
        assertThat(results.get(0)).extracting(ResultEntry::fileName).containsExactly("index-1");
        assertThat(results.get(1)).extracting(ResultEntry::fileName).containsExactly("index-2");
    }

    @Test
    public void testSingleColumnWriterClosesActiveWriter() throws Exception {
        GlobalIndexSingleColumnWriter activeWriter =
                mock(
                        GlobalIndexSingleColumnWriter.class,
                        org.mockito.Mockito.withSettings().extraInterfaces(Closeable.class));
        SortedSingleColumnIndexWriter rotatingWriter =
                new SortedSingleColumnIndexWriter(2, () -> activeWriter);
        rotatingWriter.write(10, 0);

        assertThat(rotatingWriter).isInstanceOf(AutoCloseable.class);
        ((AutoCloseable) rotatingWriter).close();

        verify((Closeable) activeWriter).close();
    }

    @Test
    public void testBuildForSinglePartitionClosesWriterAfterFailure() throws Exception {
        createTableDefault();
        GlobalIndexSingleColumnWriter activeWriter =
                mock(
                        GlobalIndexSingleColumnWriter.class,
                        org.mockito.Mockito.withSettings().extraInterfaces(Closeable.class));
        org.mockito.Mockito.doThrow(new RuntimeException("write failed"))
                .when(activeWriter)
                .write(10, 0);
        SortedGlobalIndexBuilder builder =
                new SortedGlobalIndexBuilder(getTableDefault(), "btree") {
                    @Override
                    public GlobalIndexSingleColumnWriter createWriter() {
                        return activeWriter;
                    }
                };
        builder.withIndexField("f0");

        org.assertj.core.api.Assertions.assertThatThrownBy(
                        () ->
                                builder.buildForSinglePartition(
                                        new org.apache.paimon.utils.Range(0, 0),
                                        null,
                                        Collections.<InternalRow>singletonList(
                                                        GenericRow.of(10, 0L))
                                                .iterator()))
                .isInstanceOf(RuntimeException.class)
                .hasMessage("write failed");

        verify((Closeable) activeWriter).close();
    }

    @Override
    public Schema schemaDefault() {
        Schema.Builder schemaBuilder = Schema.newBuilder();
        schemaBuilder.column("dt", DataTypes.STRING());
        schemaBuilder.column("f0", DataTypes.INT());
        schemaBuilder.column("f1", DataTypes.STRING());
        schemaBuilder.option(CoreOptions.ROW_TRACKING_ENABLED.key(), "true");
        schemaBuilder.option(CoreOptions.DATA_EVOLUTION_ENABLED.key(), "true");
        schemaBuilder.option(BTreeIndexOptions.BTREE_INDEX_RECORDS_PER_RANGE.key(), "100");
        schemaBuilder.partitionKeys(Collections.singletonList("dt"));
        return schemaBuilder.build();
    }

    private void write() throws Exception {
        createTableDefault();

        BatchWriteBuilder builder = getTableDefault().newBatchWriteBuilder();
        try (BatchTableWrite write0 = builder.newWrite()) {
            for (int i = 0; i < PART_ROW_NUM; i++) {
                write0.write(
                        GenericRow.of(
                                BinaryString.fromString("p0"),
                                i,
                                BinaryString.fromString("f1_" + i)));
            }
            BatchTableCommit commit = builder.newCommit();
            commit.commit(write0.prepareCommit());
        }

        try (BatchTableWrite write1 = builder.newWrite()) {
            for (int i = 0; i < PART_ROW_NUM; i++) {
                write1.write(
                        GenericRow.of(
                                BinaryString.fromString("p1"),
                                i,
                                BinaryString.fromString("f1_" + i)));
            }
            BatchTableCommit commit = builder.newCommit();
            commit.commit(write1.prepareCommit());
        }
    }

    private void createIndex(PartitionPredicate partitionPredicate) throws Exception {
        FileStoreTable table = getTableDefault();

        SortedGlobalIndexBuilder builder = new SortedGlobalIndexBuilder(table, "btree");
        builder.withIndexField("f0");
        builder.withPartitionPredicate(partitionPredicate);
        List<DataSplit> dataSplits =
                builder.scan()
                        .map(Pair::getRight)
                        .orElseThrow(
                                () ->
                                        new IllegalStateException(
                                                "Expected scan result when building index."));
        List<CommitMessage> commitMessages = new ArrayList<>();
        for (DataSplit dataSplit : dataSplits) {
            commitMessages.addAll(builder.build(dataSplit, ioManager));
        }

        try (BatchTableCommit commit = table.newBatchWriteBuilder().newCommit()) {
            commit.commit(commitMessages);
        }
    }

    @Test
    public void testCreateIndexForSinglePartition() throws Exception {
        write();

        FileStoreTable table = getTableDefault();
        RowType partType = table.rowType().project("dt");
        Predicate predicate =
                PartitionPredicate.createPartitionPredicate(
                        partType, Collections.singletonMap("dt", BinaryString.fromString("p0")));

        createIndex(PartitionPredicate.fromPredicate(partType, predicate));

        Map<BinaryRow, List<Pair<String, FileStats>>> metasByParts = gatherIndexMetas(table);

        Assertions.assertEquals(1, metasByParts.size());

        metasByParts.forEach(this::assertFilesNonOverlapping);
    }

    @Test
    public void testCreateIndexForMultiplePartitions() throws Exception {
        write();

        createIndex(null);

        FileStoreTable table = getTableDefault();

        Map<BinaryRow, List<Pair<String, FileStats>>> metasByParts = gatherIndexMetas(table);

        Assertions.assertEquals(2, metasByParts.size());

        metasByParts.forEach(this::assertFilesNonOverlapping);
    }

    @Test
    public void testIncrementalScanNoData() throws Exception {
        createTableDefault();

        FileStoreTable table = getTableDefault();
        SortedGlobalIndexBuilder builder = new SortedGlobalIndexBuilder(table, "btree");
        builder.withIndexField("f0");

        Assertions.assertFalse(
                builder.incrementalScan().isPresent(),
                "incrementalScan on empty table should return empty");
    }

    @Test
    public void testIncrementalScanAllIndexed() throws Exception {
        write();
        createIndex(null);

        FileStoreTable table = getTableDefault();
        SortedGlobalIndexBuilder builder = new SortedGlobalIndexBuilder(table, "btree");
        builder.withIndexField("f0");

        Assertions.assertFalse(
                builder.incrementalScan().isPresent(),
                "incrementalScan should return empty when all data is already indexed");
    }

    @Test
    public void testIncrementalScanWithNewData() throws Exception {
        write();
        createIndex(null);

        // Write additional data to partition p0
        FileStoreTable table = getTableDefault();
        BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
        try (BatchTableWrite batchWrite = writeBuilder.newWrite()) {
            for (int i = (int) PART_ROW_NUM; i < PART_ROW_NUM + 500; i++) {
                batchWrite.write(
                        GenericRow.of(
                                BinaryString.fromString("p0"),
                                i,
                                BinaryString.fromString("f1_" + i)));
            }
            try (BatchTableCommit commit = writeBuilder.newCommit()) {
                commit.commit(batchWrite.prepareCommit());
            }
        }

        table = getTableDefault();
        SortedGlobalIndexBuilder builder = new SortedGlobalIndexBuilder(table, "btree");
        builder.withIndexField("f0");

        Optional<Pair<org.apache.paimon.utils.RowRangeIndex, List<DataSplit>>> incrementalScan =
                builder.incrementalScan();
        Assertions.assertTrue(
                incrementalScan.isPresent(),
                "incrementalScan should return non-empty splits for newly written data");
        List<DataSplit> splits = incrementalScan.get().getRight();
        Assertions.assertFalse(
                splits.isEmpty(),
                "incrementalScan should return non-empty splits for newly written data");

        long totalRowCount = 0;
        for (DataSplit split : splits) {
            totalRowCount += split.rowCount();
        }
        Assertions.assertEquals(
                500, totalRowCount, "incrementalScan should only return the newly written rows");
    }

    @Test
    public void testIncrementalScanWithPartitionPredicate() throws Exception {
        write();

        // Only index partition p0
        FileStoreTable table = getTableDefault();
        RowType partType = table.rowType().project("dt");
        Predicate predicate =
                PartitionPredicate.createPartitionPredicate(
                        partType, Collections.singletonMap("dt", BinaryString.fromString("p0")));
        createIndex(PartitionPredicate.fromPredicate(partType, predicate));

        // Write new data to both partitions
        BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
        try (BatchTableWrite batchWrite = writeBuilder.newWrite()) {
            for (int i = (int) PART_ROW_NUM; i < PART_ROW_NUM + 200; i++) {
                batchWrite.write(
                        GenericRow.of(
                                BinaryString.fromString("p0"),
                                i,
                                BinaryString.fromString("f1_" + i)));
                batchWrite.write(
                        GenericRow.of(
                                BinaryString.fromString("p1"),
                                i,
                                BinaryString.fromString("f1_" + i)));
            }
            try (BatchTableCommit commit = writeBuilder.newCommit()) {
                commit.commit(batchWrite.prepareCommit());
            }
        }

        // incrementalScan with p0 filter: should return only the new p0 data
        table = getTableDefault();
        predicate =
                PartitionPredicate.createPartitionPredicate(
                        partType, Collections.singletonMap("dt", BinaryString.fromString("p0")));
        SortedGlobalIndexBuilder builder = new SortedGlobalIndexBuilder(table, "btree");
        builder.withIndexField("f0");
        builder.withPartitionPredicate(PartitionPredicate.fromPredicate(partType, predicate));

        List<DataSplit> splits =
                builder.incrementalScan()
                        .map(Pair::getRight)
                        .orElseThrow(
                                () ->
                                        new IllegalStateException(
                                                "Expected incremental scan result for new data."));
        Assertions.assertFalse(splits.isEmpty());

        for (DataSplit split : splits) {
            Assertions.assertEquals(
                    BinaryString.fromString("p0"),
                    split.partition().getString(0),
                    "incrementalScan with partition predicate should only return matching partition");
        }

        long totalRowCount = 0;
        for (DataSplit split : splits) {
            totalRowCount += split.rowCount();
        }
        Assertions.assertEquals(
                200,
                totalRowCount,
                "incrementalScan should only return the new rows in partition p0");
    }

    @Test
    public void testScanFiltersBlobFilesByManifestEntryFilter() throws Exception {
        Schema.Builder schemaBuilder = Schema.newBuilder();
        schemaBuilder.column("dt", DataTypes.STRING());
        schemaBuilder.column("f0", DataTypes.INT());
        schemaBuilder.column("f1", DataTypes.BLOB());
        schemaBuilder.option(CoreOptions.ROW_TRACKING_ENABLED.key(), "true");
        schemaBuilder.option(CoreOptions.DATA_EVOLUTION_ENABLED.key(), "true");
        schemaBuilder.option(CoreOptions.BLOB_TARGET_FILE_SIZE.key(), "1 b");
        schemaBuilder.partitionKeys(Collections.singletonList("dt"));

        catalog.createTable(identifier("BlobTable"), schemaBuilder.build(), false);
        FileStoreTable table = getTable(identifier("BlobTable"));

        byte[] blobBytes = new byte[] {1};
        BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
        try (BatchTableWrite write = writeBuilder.newWrite()) {
            for (int i = 0; i < 10; i++) {
                write.write(
                        GenericRow.of(BinaryString.fromString("p0"), i, new BlobData(blobBytes)));
            }
            try (BatchTableCommit commit = writeBuilder.newCommit()) {
                commit.commit(write.prepareCommit());
            }
        }

        Assertions.assertTrue(
                containsBlobFile(table.store().newScan().plan().files()),
                "Test table should contain blob manifest entries.");

        SortedGlobalIndexBuilder builder =
                new SortedGlobalIndexBuilder(table, "btree").withIndexField("f0");
        assertNoBlobFiles(
                builder.scan()
                        .map(Pair::getRight)
                        .orElseThrow(
                                () ->
                                        new IllegalStateException(
                                                "Expected scan result for blob table.")));
        assertNoBlobFiles(
                builder.incrementalScan()
                        .map(Pair::getRight)
                        .orElseThrow(
                                () ->
                                        new IllegalStateException(
                                                "Expected incremental scan result for blob table.")));
    }

    private Map<BinaryRow, List<Pair<String, FileStats>>> gatherIndexMetas(FileStoreTable table) {
        IndexFileHandler handler = table.store().newIndexFileHandler();

        Snapshot snapshot = table.latestSnapshot().get();
        List<IndexManifestEntry> entries = handler.scan(snapshot, "btree");

        Map<BinaryRow, List<Pair<String, FileStats>>> metasByParts = new HashMap<>();
        for (IndexManifestEntry entry : entries) {
            IndexFileMeta indexFileMeta = entry.indexFile();
            Assertions.assertNotNull(
                    indexFileMeta.globalIndexMeta(), "Global index meta should not be null");

            metasByParts
                    .computeIfAbsent(entry.partition(), part -> new ArrayList<>())
                    .add(
                            Pair.of(
                                    indexFileMeta.fileName(),
                                    FileStats.fromIndexFileMeta(indexFileMeta)));
        }

        return metasByParts;
    }

    private boolean containsBlobFile(List<ManifestEntry> entries) {
        for (ManifestEntry entry : entries) {
            if ("blob".equals(entry.file().fileFormat())) {
                return true;
            }
        }
        return false;
    }

    private void assertNoBlobFiles(List<DataSplit> splits) {
        for (DataSplit split : splits) {
            for (DataFileMeta file : split.dataFiles()) {
                Assertions.assertNotEquals(
                        "blob",
                        file.fileFormat(),
                        "Sorted global index scan should not include blob files.");
            }
        }
    }

    private void assertFilesNonOverlapping(
            BinaryRow partition, List<Pair<String, FileStats>> metas) {
        if (metas.isEmpty()) {
            return;
        }

        metas.sort((m1, m2) -> COMPARATOR.compare(m1.getValue().firstKey, m2.getValue().firstKey));
        String lastFileName = metas.get(0).getKey();
        FileStats lastMeta = metas.get(0).getValue();
        long rowCount = lastMeta.rowCount;
        for (int i = 1; i < metas.size(); i++) {
            String fileName = metas.get(i).getKey();
            FileStats fileMeta = metas.get(i).getValue();
            rowCount += fileMeta.rowCount;

            Assertions.assertTrue(
                    COMPARATOR.compare(lastMeta.lastKey, fileMeta.firstKey) <= 0,
                    String.format(
                            "In partition %s, key range [%s:%s] of file %s overlaps with adjacent file %s [%s:%s]",
                            partition.getString(0),
                            lastMeta.firstKey,
                            lastMeta.lastKey,
                            lastFileName,
                            fileName,
                            fileMeta.firstKey,
                            fileMeta.lastKey));

            lastFileName = fileName;
            lastMeta = fileMeta;
        }

        Assertions.assertEquals(
                PART_ROW_NUM,
                rowCount,
                String.format(
                        "In partition %s, total row count of all btree index files not equals to original data row count.",
                        partition.getString(0)));
    }

    static Object deserialize(byte[] bytes) {
        return KEY_SERIALIZER.deserialize(MemorySlice.wrap(bytes));
    }

    /** File Stats class for assertion. */
    private static class FileStats {

        final long rowCount;
        final Object firstKey;
        final Object lastKey;

        FileStats(long rowCount, Object firstKey, Object lastKey) {
            this.rowCount = rowCount;
            this.firstKey = firstKey;
            this.lastKey = lastKey;
        }

        static FileStats fromIndexFileMeta(org.apache.paimon.index.IndexFileMeta meta) {
            Assertions.assertNotNull(meta.globalIndexMeta());
            GlobalIndexMeta globalIndexMeta = meta.globalIndexMeta();
            org.apache.paimon.globalindex.SortedIndexFileMeta indexMeta =
                    org.apache.paimon.globalindex.SortedIndexFileMeta.deserialize(
                            globalIndexMeta.indexMeta());

            return new FileStats(
                    meta.rowCount(),
                    deserialize(indexMeta.getFirstKey()),
                    deserialize(indexMeta.getLastKey()));
        }
    }
}
