/*
 * 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.connector.elasticsearch.sink;

import org.apache.flink.api.connector.sink2.SinkWriter;
import org.apache.flink.api.connector.sink2.SinkWriter.Context;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.connector.elasticsearch.ElasticsearchUtil;
import org.apache.flink.connector.elasticsearch.NetworkClientConfig;
import org.apache.flink.connector.elasticsearch.sink.ElasticsearchWriter.DefaultBulkResponseInspector;
import org.apache.flink.connector.elasticsearch.test.DockerImageVersions;
import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.Gauge;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.metrics.groups.OperatorIOMetricGroup;
import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
import org.apache.flink.metrics.testutils.MetricListener;
import org.apache.flink.runtime.mailbox.SyncMailboxExecutor;
import org.apache.flink.runtime.metrics.MetricNames;
import org.apache.flink.runtime.metrics.groups.InternalSinkWriterMetricGroup;
import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration;
import org.apache.flink.test.junit5.MiniClusterExtension;
import org.apache.flink.util.TestLoggerExtension;

import org.apache.http.HttpHost;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.bulk.BackoffPolicy;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.unit.TimeValue;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.elasticsearch.ElasticsearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;

import static org.apache.flink.connector.elasticsearch.sink.TestClientBase.DOCUMENT_TYPE;
import static org.apache.flink.connector.elasticsearch.sink.TestClientBase.buildMessage;
import static org.assertj.core.api.Assertions.assertThat;

/** Tests for {@link ElasticsearchWriter}. */
@Testcontainers
@ExtendWith(TestLoggerExtension.class)
class ElasticsearchWriterITCase {

    private static final Logger LOG = LoggerFactory.getLogger(ElasticsearchWriterITCase.class);

    @RegisterExtension
    private static final MiniClusterExtension MINI_CLUSTER_RESOURCE =
            new MiniClusterExtension(
                    new MiniClusterResourceConfiguration.Builder()
                            .setNumberTaskManagers(1)
                            .setNumberSlotsPerTaskManager(3)
                            .build());

    @Container
    private static final ElasticsearchContainer ES_CONTAINER =
            ElasticsearchUtil.createElasticsearchContainer(
                    DockerImageVersions.ELASTICSEARCH_7, LOG);

    private RestHighLevelClient client;
    private TestClientBase context;
    private MetricListener metricListener;

    @BeforeEach
    void setUp() {
        metricListener = new MetricListener();
        client =
                new RestHighLevelClient(
                        RestClient.builder(HttpHost.create(ES_CONTAINER.getHttpHostAddress())));
        context = new TestClient(client);
    }

    @AfterEach
    void tearDown() throws IOException {
        if (client != null) {
            client.close();
        }
    }

    @Test
    void testWriteOnBulkFlush() throws Exception {
        final String index = "test-bulk-flush-without-checkpoint";
        final int flushAfterNActions = 5;
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(flushAfterNActions, -1, -1, FlushBackoffType.NONE, 0, 0);

        SinkWriterMetricGroup metricGroup =
                InternalSinkWriterMetricGroup.wrap(
                        new TestingSinkWriterMetricGroup.Builder()
                                .setIoMetricGroupSupplier(
                                        UnregisteredMetricsGroup::createOperatorIOMetricGroup)
                                .setParentMetricGroup(
                                        UnregisteredMetricsGroup.createOperatorMetricGroup())
                                .build());
        final UpdatingEmitter updatingEmitter =
                new UpdatingEmitter(index, context.getDataFieldName()) {
                    @Override
                    public void emit(
                            Tuple2<Integer, String> element,
                            Context context,
                            RequestIndexer indexer) {
                        super.emit(element, context, indexer);
                        if (element.f0 == 8) {
                            indexer.flush();
                        }
                    }
                };

        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(false, bulkProcessorConfig, metricGroup, updatingEmitter)) {
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);
            writer.write(Tuple2.of(3, buildMessage(3)), null);
            writer.write(Tuple2.of(4, buildMessage(4)), null);

            // Ignore flush on checkpoint
            writer.flush(false);

            context.assertThatIdsAreNotWritten(index, 1, 2, 3, 4);

            // Trigger flush
            writer.write(Tuple2.of(5, "test-5"), null);
            context.assertThatIdsAreWritten(index, 1, 2, 3, 4, 5);

            writer.write(Tuple2.of(6, "test-6"), null);
            context.assertThatIdsAreNotWritten(index, 6);

            // Force flush
            writer.blockingFlushAllActions();
            context.assertThatIdsAreWritten(index, 1, 2, 3, 4, 5, 6);

            writer.write(Tuple2.of(7, "test-7"), null);
            context.assertThatIdsAreNotWritten(index, 7);
            writer.write(Tuple2.of(8, "test-8"), null);
            context.assertThatIdsAreWritten(index, 1, 2, 3, 4, 5, 6, 7, 8);
        }
    }

    @Test
    void testWriteOnBulkIntervalFlush() throws Exception {
        final String index = "test-bulk-flush-with-interval";

        // Configure bulk processor to flush every 1s;
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(-1, -1, 1000, FlushBackoffType.NONE, 0, 0);

        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(index, false, bulkProcessorConfig)) {
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);
            writer.write(Tuple2.of(3, buildMessage(3)), null);
            writer.write(Tuple2.of(4, buildMessage(4)), null);
            writer.blockingFlushAllActions();
        }

        context.assertThatIdsAreWritten(index, 1, 2, 3, 4);
    }

    @Test
    void testWriteOnCheckpoint() throws Exception {
        final String index = "test-bulk-flush-with-checkpoint";
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(-1, -1, -1, FlushBackoffType.NONE, 0, 0);

        // Enable flush on checkpoint
        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(index, true, bulkProcessorConfig)) {
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);
            writer.write(Tuple2.of(3, buildMessage(3)), null);

            context.assertThatIdsAreNotWritten(index, 1, 2, 3);

            // Trigger flush
            writer.flush(false);

            context.assertThatIdsAreWritten(index, 1, 2, 3);
        }
    }

    @Test
    void testIncrementByteOutMetric() throws Exception {
        final String index = "test-inc-byte-out";
        final OperatorIOMetricGroup operatorIOMetricGroup =
                UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup().getIOMetricGroup();
        final int flushAfterNActions = 2;
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(flushAfterNActions, -1, -1, FlushBackoffType.NONE, 0, 0);

        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(
                        index,
                        false,
                        bulkProcessorConfig,
                        getSinkWriterMetricGroup(operatorIOMetricGroup))) {
            final Counter numBytesOut = operatorIOMetricGroup.getNumBytesOutCounter();
            assertThat(numBytesOut.getCount()).isZero();
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);

            writer.blockingFlushAllActions();
            long first = numBytesOut.getCount();
            assertThat(first).isGreaterThan(0);

            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);

            writer.blockingFlushAllActions();
            assertThat(numBytesOut.getCount()).isGreaterThan(first);
        }
    }

    @Test
    void testIncrementRecordsSendMetric() throws Exception {
        final String index = "test-inc-records-send";
        final int flushAfterNActions = 2;
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(flushAfterNActions, -1, -1, FlushBackoffType.NONE, 0, 0);

        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(index, false, bulkProcessorConfig)) {
            final Optional<Counter> recordsSend =
                    metricListener.getCounter(MetricNames.NUM_RECORDS_SEND);
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            // Update existing index
            writer.write(Tuple2.of(1, "u" + buildMessage(2)), null);
            // Delete index
            writer.write(Tuple2.of(1, "d" + buildMessage(3)), null);

            writer.blockingFlushAllActions();

            assertThat(recordsSend).isPresent();
            assertThat(recordsSend.get().getCount()).isEqualTo(3L);
        }
    }

    @Test
    void testCurrentSendTime() throws Exception {
        final String index = "test-current-send-time";
        final int flushAfterNActions = 2;
        final BulkProcessorConfig bulkProcessorConfig =
                new BulkProcessorConfig(flushAfterNActions, -1, -1, FlushBackoffType.NONE, 0, 0);

        try (final ElasticsearchWriter<Tuple2<Integer, String>> writer =
                createWriter(index, false, bulkProcessorConfig)) {
            final Optional<Gauge<Long>> currentSendTime =
                    metricListener.getGauge("currentSendTime");
            writer.write(Tuple2.of(1, buildMessage(1)), null);
            writer.write(Tuple2.of(2, buildMessage(2)), null);

            writer.blockingFlushAllActions();

            assertThat(currentSendTime).isPresent();
            assertThat(currentSendTime.get().getValue()).isGreaterThan(0L);
        }
    }

    private ElasticsearchWriter<Tuple2<Integer, String>> createWriter(
            String index, boolean flushOnCheckpoint, BulkProcessorConfig bulkProcessorConfig) {
        return createWriter(
                index, flushOnCheckpoint, bulkProcessorConfig, getSinkWriterMetricGroup());
    }

    private ElasticsearchWriter<Tuple2<Integer, String>> createWriter(
            String index,
            boolean flushOnCheckpoint,
            BulkProcessorConfig bulkProcessorConfig,
            SinkWriterMetricGroup metricGroup) {
        return createWriter(
                flushOnCheckpoint,
                bulkProcessorConfig,
                metricGroup,
                new UpdatingEmitter(index, context.getDataFieldName()));
    }

    private static ElasticsearchWriter<Tuple2<Integer, String>> createWriter(
            boolean flushOnCheckpoint,
            BulkProcessorConfig bulkProcessorConfig,
            SinkWriterMetricGroup metricGroup,
            UpdatingEmitter updatingEmitter) {
        return new ElasticsearchWriter<>(
                Collections.singletonList(HttpHost.create(ES_CONTAINER.getHttpHostAddress())),
                updatingEmitter,
                flushOnCheckpoint,
                bulkProcessorConfig,
                new TestBulkProcessorBuilderFactory(),
                new DefaultBulkResponseInspector(),
                new NetworkClientConfig(null, null, null, null, null, null, null, null),
                metricGroup,
                new SyncMailboxExecutor());
    }

    private TestingSinkWriterMetricGroup getSinkWriterMetricGroup() {
        final OperatorIOMetricGroup operatorIOMetricGroup =
                UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup().getIOMetricGroup();
        return getSinkWriterMetricGroup(operatorIOMetricGroup);
    }

    private TestingSinkWriterMetricGroup getSinkWriterMetricGroup(
            OperatorIOMetricGroup operatorIOMetricGroup) {
        MetricGroup parentMetricGroup = metricListener.getMetricGroup();
        Counter numRecordsOutErrors = parentMetricGroup.counter(MetricNames.NUM_RECORDS_OUT_ERRORS);
        Counter numRecordsSendErrors =
                parentMetricGroup.counter(MetricNames.NUM_RECORDS_SEND_ERRORS, numRecordsOutErrors);
        Counter numRecordsWritten =
                parentMetricGroup.counter(
                        MetricNames.NUM_RECORDS_SEND,
                        operatorIOMetricGroup.getNumRecordsOutCounter());
        Counter numBytesWritten =
                parentMetricGroup.counter(
                        MetricNames.NUM_BYTES_SEND, operatorIOMetricGroup.getNumBytesOutCounter());
        Consumer<Gauge<Long>> currentSendTimeGaugeConsumer =
                currentSendTimeGauge ->
                        parentMetricGroup.gauge(
                                MetricNames.CURRENT_SEND_TIME, currentSendTimeGauge);
        return new TestingSinkWriterMetricGroup.Builder()
                .setParentMetricGroup(parentMetricGroup)
                .setIoMetricGroupSupplier(() -> operatorIOMetricGroup)
                .setNumRecordsOutErrorsCounterSupplier(() -> numRecordsOutErrors)
                .setNumRecordsSendErrorsCounterSupplier(() -> numRecordsSendErrors)
                .setNumRecordsSendCounterSupplier(() -> numRecordsWritten)
                .setNumBytesSendCounterSupplier(() -> numBytesWritten)
                .setCurrentSendTimeGaugeConsumer(currentSendTimeGaugeConsumer)
                .build();
    }

    private static class TestBulkProcessorBuilderFactory implements BulkProcessorBuilderFactory {
        @Override
        public BulkProcessor.Builder apply(
                RestHighLevelClient client,
                BulkProcessorConfig bulkProcessorConfig,
                BulkProcessor.Listener listener) {
            BulkProcessor.Builder builder =
                    BulkProcessor.builder(
                            new BulkRequestConsumerFactory() { // This cannot be inlined as a lambda
                                // because then deserialization fails
                                @Override
                                public void accept(
                                        BulkRequest bulkRequest,
                                        ActionListener<BulkResponse> bulkResponseActionListener) {
                                    client.bulkAsync(
                                            bulkRequest,
                                            RequestOptions.DEFAULT,
                                            bulkResponseActionListener);
                                }
                            },
                            listener);

            if (bulkProcessorConfig.getBulkFlushMaxActions() != -1) {
                builder.setBulkActions(bulkProcessorConfig.getBulkFlushMaxActions());
            }

            if (bulkProcessorConfig.getBulkFlushMaxMb() != -1) {
                builder.setBulkSize(
                        new ByteSizeValue(
                                bulkProcessorConfig.getBulkFlushMaxMb(), ByteSizeUnit.MB));
            }

            if (bulkProcessorConfig.getBulkFlushInterval() != -1) {
                builder.setFlushInterval(new TimeValue(bulkProcessorConfig.getBulkFlushInterval()));
            }

            BackoffPolicy backoffPolicy;
            final TimeValue backoffDelay =
                    new TimeValue(bulkProcessorConfig.getBulkFlushBackOffDelay());
            final int maxRetryCount = bulkProcessorConfig.getBulkFlushBackoffRetries();
            switch (bulkProcessorConfig.getFlushBackoffType()) {
                case CONSTANT:
                    backoffPolicy = BackoffPolicy.constantBackoff(backoffDelay, maxRetryCount);
                    break;
                case EXPONENTIAL:
                    backoffPolicy = BackoffPolicy.exponentialBackoff(backoffDelay, maxRetryCount);
                    break;
                case NONE:
                    backoffPolicy = BackoffPolicy.noBackoff();
                    break;
                default:
                    throw new IllegalArgumentException(
                            "Received unknown backoff policy type "
                                    + bulkProcessorConfig.getFlushBackoffType());
            }
            builder.setBackoffPolicy(backoffPolicy);
            return builder;
        }
    }

    private static class UpdatingEmitter implements ElasticsearchEmitter<Tuple2<Integer, String>> {

        private final String dataFieldName;
        private final String index;

        UpdatingEmitter(String index, String dataFieldName) {
            this.index = index;
            this.dataFieldName = dataFieldName;
        }

        @Override
        public void emit(
                Tuple2<Integer, String> element,
                SinkWriter.Context context,
                RequestIndexer indexer) {

            Map<String, Object> document = new HashMap<>();
            document.put(dataFieldName, element.f1);

            final char action = element.f1.charAt(0);
            final String id = element.f0.toString();
            switch (action) {
                case 'd':
                    {
                        indexer.add(new DeleteRequest(index).id(id));
                        break;
                    }
                case 'u':
                    {
                        indexer.add(new UpdateRequest().index(index).id(id).doc(document));
                        break;
                    }
                default:
                    {
                        indexer.add(
                                new IndexRequest(index)
                                        .id(id)
                                        .type(DOCUMENT_TYPE)
                                        .source(document));
                    }
            }
        }
    }

    private static class TestClient extends TestClientBase {

        TestClient(RestHighLevelClient client) {
            super(client);
        }

        @Override
        GetResponse getResponse(String index, int id) throws IOException {
            return client.get(new GetRequest(index, Integer.toString(id)), RequestOptions.DEFAULT);
        }
    }
}
