/*
 * 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.fluss.client.table.scanner;

import org.apache.fluss.record.ChangeType;
import org.apache.fluss.record.LogRecord;
import org.apache.fluss.row.BinaryString;
import org.apache.fluss.row.GenericRow;
import org.apache.fluss.row.InternalRow;
import org.apache.fluss.row.KeyValueRow;
import org.apache.fluss.row.ProjectedRow;
import org.apache.fluss.types.IntType;
import org.apache.fluss.types.RowType;
import org.apache.fluss.types.StringType;
import org.apache.fluss.utils.CloseableIterator;

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.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

import static org.apache.fluss.testutils.DataTestUtils.row;
import static org.assertj.core.api.Assertions.assertThat;

/** Test for {@link SortMergeReader}. */
class SortMergeReaderTest {

    private static class FlussRowComparator implements Comparator<InternalRow> {

        private final int keyIndex;

        public FlussRowComparator(int keyIndex) {
            this.keyIndex = keyIndex;
        }

        @Override
        public int compare(InternalRow o1, InternalRow o2) {
            return o1.getInt(keyIndex) - o2.getInt(keyIndex);
        }
    }

    @ParameterizedTest
    @ValueSource(booleans = {false, true})
    void testReadBatch(boolean isProjected) {
        int keyIndex = 0;
        int[] pkIndexes = new int[] {keyIndex};
        int[] projectedFields = isProjected ? new int[] {keyIndex, 1} : null;
        List<LogRecord> logRecords1 = createRecords(0, 10, false);
        List<LogRecord> logRecords2 = createRecords(10, 10, false);
        List<LogRecord> logRecords3 = createRecords(20, 10, false);
        List<KeyValueRow> logRecords4 =
                createRecords(5, 20, true).stream()
                        .map(logRecord -> new KeyValueRow(pkIndexes, logRecord.getRow(), false))
                        .collect(Collectors.toList());

        SortMergeReader sortMergeReader =
                new SortMergeReader(
                        projectedFields,
                        new int[] {keyIndex},
                        Arrays.asList(
                                CloseableIterator.wrap(logRecords2.iterator()),
                                CloseableIterator.wrap(logRecords3.iterator()),
                                CloseableIterator.wrap(logRecords1.iterator())),
                        new FlussRowComparator(keyIndex),
                        CloseableIterator.wrap(logRecords4.iterator()));

        List<InternalRow> actualRows = new ArrayList<>();
        InternalRow.FieldGetter[] fieldGetters =
                InternalRow.createFieldGetters(
                        isProjected
                                ? RowType.of(new IntType(), new StringType())
                                : RowType.of(new IntType(), new StringType(), new StringType()));
        try (CloseableIterator<InternalRow> iterator = sortMergeReader.readBatch()) {
            actualRows.addAll(materializeRows(iterator, fieldGetters));
        }
        assertThat(actualRows).hasSize(30);

        List<LogRecord> expectedLogRecords = createRecords(0, 5, false);
        expectedLogRecords.addAll(createRecords(5, 20, true));
        expectedLogRecords.addAll(createRecords(25, 5, false));
        ProjectedRow projectedRow = isProjected ? ProjectedRow.from(projectedFields) : null;
        CloseableIterator<InternalRow> expected =
                isProjected
                        ? projected(
                                CloseableIterator.wrap(expectedLogRecords.iterator()), projectedRow)
                        : CloseableIterator.wrap(
                                expectedLogRecords.stream().map(LogRecord::getRow).iterator());
        assertThat(actualRows).isEqualTo(materializeRows(expected, fieldGetters));
    }

    @Test
    void testReadBatchSkipsDeletedSnapshotRows() {
        int keyIndex = 0;
        int[] pkIndexes = new int[] {keyIndex};

        List<LogRecord> snapshotRecords = createRecords(0, 10, false);
        List<KeyValueRow> changeLogRecords = createDeleteRecords(pkIndexes, Arrays.asList(2, 5, 8));

        SortMergeReader sortMergeReader =
                new SortMergeReader(
                        null,
                        pkIndexes,
                        CloseableIterator.wrap(snapshotRecords.iterator()),
                        new FlussRowComparator(keyIndex),
                        CloseableIterator.wrap(changeLogRecords.iterator()));

        InternalRow.FieldGetter[] fieldGetters =
                InternalRow.createFieldGetters(
                        RowType.of(new IntType(), new StringType(), new StringType()));
        List<InternalRow> actualRows;
        try (CloseableIterator<InternalRow> iterator = sortMergeReader.readBatch()) {
            actualRows = materializeRows(iterator, fieldGetters);
        }

        assertThat(actualRows).hasSize(7);
        assertThat(actualRows.stream().map(row -> row.getInt(0)).collect(Collectors.toList()))
                .containsExactly(0, 1, 3, 4, 6, 7, 9);
    }

    private CloseableIterator<InternalRow> projected(
            CloseableIterator<LogRecord> originElementIterator, final ProjectedRow projectedRow) {
        return new CloseableIterator<InternalRow>() {
            private final CloseableIterator<LogRecord> inner = originElementIterator;

            @Override
            public void close() {
                inner.close();
            }

            @Override
            public boolean hasNext() {
                return inner.hasNext();
            }

            @Override
            public InternalRow next() {
                LogRecord element = inner.next();
                return projectedRow.replaceRow(element.getRow());
            }
        };
    }

    private List<InternalRow> materializeRows(
            CloseableIterator<InternalRow> iterator, InternalRow.FieldGetter[] fieldGetters) {
        List<InternalRow> actualRows = new ArrayList<>();
        while (iterator != null && iterator.hasNext()) {
            InternalRow row = iterator.next();
            GenericRow genericRow = new GenericRow(fieldGetters.length);
            for (int i = 0; i < fieldGetters.length; i++) {
                genericRow.setField(i, fieldGetters[i].getFieldOrNull(row));
            }
            actualRows.add(genericRow);
        }
        return actualRows;
    }

    private List<LogRecord> createRecords(int startId, int count, boolean isLog) {
        List<LogRecord> logRecords = new ArrayList<>();

        for (int i = 0; i < count; i++) {
            GenericRow row =
                    row(
                            startId + i,
                            BinaryString.fromString(isLog ? "a" + "_updated" : "a"),
                            BinaryString.fromString(isLog ? "A" + "_updated" : "A"));
            logRecords.add(new ScanRecord(i, System.currentTimeMillis(), ChangeType.INSERT, row));
        }
        return logRecords;
    }

    private List<KeyValueRow> createDeleteRecords(int[] pkIndexes, List<Integer> keys) {
        List<KeyValueRow> deleteRecords = new ArrayList<>();
        for (Integer key : keys) {
            deleteRecords.add(
                    new KeyValueRow(
                            pkIndexes,
                            row(key, BinaryString.fromString("a"), BinaryString.fromString("A")),
                            true));
        }
        return deleteRecords;
    }
}
