/*
 * 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.flink.streaming.connectors.kafka.table;

import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.operators.ProcessingTimeService;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.sink2.SinkWriter;
import org.apache.flink.connector.kafka.sink.TwoPhaseCommittingStatefulSink;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.catalog.Column;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.data.TimestampData;
import org.apache.flink.table.runtime.connector.sink.SinkRuntimeProviderContext;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
import org.apache.flink.table.runtime.typeutils.RowDataSerializer;

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

import java.io.IOException;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ScheduledFuture;
import java.util.stream.Stream;

import static org.apache.flink.types.RowKind.DELETE;
import static org.apache.flink.types.RowKind.INSERT;
import static org.apache.flink.types.RowKind.UPDATE_AFTER;
import static org.assertj.core.api.Assertions.assertThat;

/** Tests for {@link ReducingUpsertWriter}. */
class ReducingUpsertWriterTest {

    static Stream<Boolean> enableObjectReuse() {
        return Stream.of(true, false);
    }

    private static final ResolvedSchema SCHEMA =
            ResolvedSchema.of(
                    Column.physical("id", DataTypes.INT().notNull()),
                    Column.physical("title", DataTypes.STRING().notNull()),
                    Column.physical("author", DataTypes.STRING()),
                    Column.physical("price", DataTypes.DOUBLE()),
                    Column.physical("qty", DataTypes.INT()),
                    Column.physical("ts", DataTypes.TIMESTAMP_LTZ(3)));

    private static final int keyIndices = 0;
    private static final int TIMESTAMP_INDICES = 5;
    private static final SinkBufferFlushMode BUFFER_FLUSH_MODE =
            new SinkBufferFlushMode(4, Long.MAX_VALUE);

    public static final RowData[] TEST_DATA = {
        GenericRowData.ofKind(
                INSERT,
                1001,
                StringData.fromString("Java public for dummies"),
                StringData.fromString("Tan Ah Teck"),
                11.11,
                11,
                TimestampData.fromInstant(Instant.parse("2021-03-30T15:00:00Z"))),
        GenericRowData.ofKind(
                INSERT,
                1002,
                StringData.fromString("More Java for dummies"),
                StringData.fromString("Tan Ah Teck"),
                22.22,
                22,
                TimestampData.fromInstant(Instant.parse("2021-03-30T16:00:00Z"))),
        GenericRowData.ofKind(
                INSERT,
                1004,
                StringData.fromString("A Cup of Java"),
                StringData.fromString("Kumar"),
                44.44,
                44,
                TimestampData.fromInstant(Instant.parse("2021-03-30T17:00:00Z"))),
        GenericRowData.ofKind(
                UPDATE_AFTER,
                1004,
                StringData.fromString("A Teaspoon of Java"),
                StringData.fromString("Kevin Jones"),
                55.55,
                55,
                TimestampData.fromInstant(Instant.parse("2021-03-30T18:00:00Z"))),
        GenericRowData.ofKind(
                UPDATE_AFTER,
                1004,
                StringData.fromString("A Teaspoon of Java 1.4"),
                StringData.fromString("Kevin Jones"),
                66.66,
                66,
                TimestampData.fromInstant(Instant.parse("2021-03-30T19:00:00Z"))),
        GenericRowData.ofKind(
                UPDATE_AFTER,
                1004,
                StringData.fromString("A Teaspoon of Java 1.5"),
                StringData.fromString("Kevin Jones"),
                77.77,
                77,
                TimestampData.fromInstant(Instant.parse("2021-03-30T20:00:00Z"))),
        GenericRowData.ofKind(
                DELETE,
                1004,
                StringData.fromString("A Teaspoon of Java 1.8"),
                StringData.fromString("Kevin Jones"),
                null,
                1010,
                TimestampData.fromInstant(Instant.parse("2021-03-30T21:00:00Z"))),
        GenericRowData.ofKind(
                DELETE,
                1005,
                StringData.fromString("A Teaspoon of Java 1.8"),
                StringData.fromString("Kevin Jones"),
                null,
                1010,
                TimestampData.fromInstant(Instant.parse("2021-03-30T21:00:00Z")))
    };

    @ParameterizedTest(name = "object reuse = {0}")
    @MethodSource("enableObjectReuse")
    public void testWriteData(boolean enableObjectReuse) throws Exception {
        final MockedSinkWriter writer = new MockedSinkWriter();
        final ReducingUpsertWriter<?, ?> bufferedWriter =
                createBufferedWriter(writer, enableObjectReuse);

        // write 4 records which doesn't trigger batch size
        writeData(bufferedWriter, new ReusableIterator(0, 4, enableObjectReuse));
        assertThat(writer.rowDataCollectors).isEmpty();

        // write one more record, and should flush the buffer
        writeData(bufferedWriter, new ReusableIterator(7, 1, enableObjectReuse));

        HashMap<Integer, List<RowData>> expected = new HashMap<>();
        expected.put(
                1001,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1001,
                                StringData.fromString("Java public for dummies"),
                                StringData.fromString("Tan Ah Teck"),
                                11.11,
                                11,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T15:00:00Z")))));
        expected.put(
                1002,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1002,
                                StringData.fromString("More Java for dummies"),
                                StringData.fromString("Tan Ah Teck"),
                                22.22,
                                22,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T16:00:00Z")))));
        expected.put(
                1004,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1004,
                                StringData.fromString("A Teaspoon of Java"),
                                StringData.fromString("Kevin Jones"),
                                55.55,
                                55,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T18:00:00Z")))));

        expected.put(
                1005,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                DELETE,
                                1005,
                                StringData.fromString("A Teaspoon of Java 1.8"),
                                StringData.fromString("Kevin Jones"),
                                null,
                                1010,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T21:00:00Z")))));

        compareCompactedResult(expected, writer.rowDataCollectors);

        writer.rowDataCollectors.clear();
        // write remaining data, and they are still buffered
        writeData(bufferedWriter, new ReusableIterator(4, 3, enableObjectReuse));
        assertThat(writer.rowDataCollectors).isEmpty();
    }

    @ParameterizedTest(name = "object reuse = {0}")
    @MethodSource("enableObjectReuse")
    public void testFlushDataWhenCheckpointing(boolean enableObjectReuse) throws Exception {
        final MockedSinkWriter writer = new MockedSinkWriter();
        final ReducingUpsertWriter<?, ?> bufferedWriter =
                createBufferedWriter(writer, enableObjectReuse);
        // write all data, there should be 3 records are still buffered
        writeData(bufferedWriter, new ReusableIterator(0, 4, enableObjectReuse));
        // snapshot should flush the buffer
        bufferedWriter.flush(true);
        assertThat(writer.flushed).isTrue();

        HashMap<Integer, List<RowData>> expected = new HashMap<>();
        expected.put(
                1001,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1001,
                                StringData.fromString("Java public for dummies"),
                                StringData.fromString("Tan Ah Teck"),
                                11.11,
                                11,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T15:00:00Z")))));
        expected.put(
                1002,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1002,
                                StringData.fromString("More Java for dummies"),
                                StringData.fromString("Tan Ah Teck"),
                                22.22,
                                22,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T16:00:00Z")))));
        expected.put(
                1004,
                Arrays.asList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1004,
                                StringData.fromString("A Teaspoon of Java"),
                                StringData.fromString("Kevin Jones"),
                                55.55,
                                55,
                                TimestampData.fromInstant(Instant.parse("2021-03-30T18:00:00Z")))));

        compareCompactedResult(expected, writer.rowDataCollectors);
    }

    @Test
    void testWriteDataWithNullTimestamp() throws Exception {
        final MockedSinkWriter writer = new MockedSinkWriter();
        final ReducingUpsertWriter<?, ?> bufferedWriter = createBufferedWriter(writer, false);

        bufferedWriter.write(
                GenericRowData.ofKind(
                        INSERT,
                        1001,
                        StringData.fromString("Java public for dummies"),
                        StringData.fromString("Tan Ah Teck"),
                        11.11,
                        11,
                        null),
                new SinkWriter.Context() {
                    @Override
                    public long currentWatermark() {
                        throw new UnsupportedOperationException("Not implemented.");
                    }

                    @Override
                    public Long timestamp() {
                        return null;
                    }
                });

        bufferedWriter.flush(true);

        final HashMap<Integer, List<RowData>> expected = new HashMap<>();
        expected.put(
                1001,
                Collections.singletonList(
                        GenericRowData.ofKind(
                                UPDATE_AFTER,
                                1001,
                                StringData.fromString("Java public for dummies"),
                                StringData.fromString("Tan Ah Teck"),
                                11.11,
                                11,
                                null)));

        compareCompactedResult(expected, writer.rowDataCollectors);
    }

    private void compareCompactedResult(
            Map<Integer, List<RowData>> expected, List<RowData> actual) {
        Map<Integer, List<RowData>> actualMap = new HashMap<>();

        for (RowData rowData : actual) {
            Integer id = rowData.getInt(keyIndices);
            actualMap.computeIfAbsent(id, key -> new ArrayList<>()).add(rowData);
        }

        assertThat(actualMap).hasSameSizeAs(expected);
        for (Integer id : expected.keySet()) {
            assertThat(actualMap.get(id)).isEqualTo(expected.get(id));
        }
    }

    private void writeData(ReducingUpsertWriter<?, ?> writer, Iterator<RowData> iterator)
            throws Exception {
        while (iterator.hasNext()) {
            RowData next = iterator.next();
            long rowtime = next.getTimestamp(TIMESTAMP_INDICES, 3).getMillisecond();
            writer.write(
                    next,
                    new org.apache.flink.api.connector.sink2.SinkWriter.Context() {
                        @Override
                        public long currentWatermark() {
                            throw new UnsupportedOperationException("Not implemented.");
                        }

                        @Override
                        public Long timestamp() {
                            return rowtime;
                        }
                    });
        }
    }

    @SuppressWarnings("unchecked")
    private ReducingUpsertWriter<?, ?> createBufferedWriter(
            MockedSinkWriter sinkWriter, boolean enableObjectReuse) {
        TypeInformation<RowData> typeInformation =
                (TypeInformation<RowData>)
                        new SinkRuntimeProviderContext(false)
                                .createTypeInformation(SCHEMA.toPhysicalRowDataType());
        return new ReducingUpsertWriter<>(
                sinkWriter,
                SCHEMA.toPhysicalRowDataType(),
                new int[] {keyIndices},
                BUFFER_FLUSH_MODE,
                new ProcessingTimeService() {
                    @Override
                    public long getCurrentProcessingTime() {
                        return 0;
                    }

                    @Override
                    public ScheduledFuture<?> registerTimer(
                            long timestamp, ProcessingTimeCallback target) {
                        return null;
                    }
                },
                enableObjectReuse
                        ? typeInformation.createSerializer(
                                        new ExecutionConfig().getSerializerConfig())
                                ::copy
                        : r -> r);
    }

    private static class MockedSinkWriter
            implements TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter<
                    RowData, Void, Void> {

        boolean flushed = false;

        transient List<RowData> rowDataCollectors;

        MockedSinkWriter() {
            rowDataCollectors = new ArrayList<>();
        }

        public void write(RowData element, Context context)
                throws IOException, InterruptedException {
            // Allow comparison between null timestamps
            if (context.timestamp() == null) {
                assertThat(element.getTimestamp(TIMESTAMP_INDICES, 3)).isNull();
            } else {
                assertThat(Instant.ofEpochMilli(context.timestamp()))
                        .isEqualTo(element.getTimestamp(TIMESTAMP_INDICES, 3).toInstant());
            }
            rowDataCollectors.add(element);
        }

        public void flush(boolean endOfInput) throws IOException, InterruptedException {
            flushed = true;
        }

        @Override
        public void close() throws Exception {}

        @Override
        public List<Void> snapshotState(long checkpointId) throws IOException {
            return null;
        }

        @Override
        public Collection<Void> prepareCommit() throws IOException, InterruptedException {
            return null;
        }
    }

    private class ReusableIterator implements Iterator<RowData> {

        private final RowDataSerializer serializer =
                InternalTypeInfo.of(SCHEMA.toSinkRowDataType().getLogicalType()).toRowSerializer();
        private final RowData reusedRow = new GenericRowData(SCHEMA.getColumnCount());
        private final boolean enableObjectReuse;

        private int begin;
        private final int end;

        ReusableIterator(int begin, int size, boolean enableObjectReuse) {
            this.begin = begin;
            this.end = begin + size;
            this.enableObjectReuse = enableObjectReuse;
        }

        @Override
        public boolean hasNext() {
            return begin < end;
        }

        @Override
        public RowData next() {
            if (enableObjectReuse) {
                return serializer.copy(TEST_DATA[begin++], reusedRow);
            } else {
                return TEST_DATA[begin++];
            }
        }
    }
}
