/*
 * 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.inlong.sort.kafka;

import org.apache.inlong.sort.base.dirty.DirtyOptions;
import org.apache.inlong.sort.base.dirty.sink.DirtySink;
import org.apache.inlong.sort.kafka.DynamicKafkaSerializationSchema.MetadataConverter;
import org.apache.inlong.sort.protocol.enums.SchemaChangePolicy;
import org.apache.inlong.sort.protocol.enums.SchemaChangeType;

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.streaming.connectors.kafka.table.BufferedUpsertSinkFunction;
import org.apache.flink.streaming.connectors.kafka.table.KafkaSinkSemantic;
import org.apache.flink.streaming.connectors.kafka.table.SinkBufferFlushMode;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.catalog.CatalogTable;
import org.apache.flink.table.connector.ChangelogMode;
import org.apache.flink.table.connector.format.EncodingFormat;
import org.apache.flink.table.connector.sink.DynamicTableSink;
import org.apache.flink.table.connector.sink.SinkFunctionProvider;
import org.apache.flink.table.connector.sink.abilities.SupportsWritingMetadata;
import org.apache.flink.table.data.ArrayData;
import org.apache.flink.table.data.MapData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.utils.DataTypeUtils;
import org.apache.kafka.common.header.Header;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.apache.inlong.sort.kafka.table.KafkaOptions.KAFKA_IGNORE_ALL_CHANGELOG;

/**
 * A version-agnostic Kafka {@link DynamicTableSink}.
 *
 * Add an option `inlong.metric` to support metrics.
 */
@Internal
public class KafkaDynamicSink implements DynamicTableSink, SupportsWritingMetadata {

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

    // --------------------------------------------------------------------------------------------
    // Mutable attributes
    // --------------------------------------------------------------------------------------------

    // --------------------------------------------------------------------------------------------
    // Format attributes
    // --------------------------------------------------------------------------------------------
    private static final String VALUE_METADATA_PREFIX = "value.";
    /**
     * Data type to configure the formats.
     */
    protected final DataType physicalDataType;
    /**
     * Optional format for encoding keys to Kafka.
     */
    protected final @Nullable EncodingFormat<SerializationSchema<RowData>> keyEncodingFormat;
    /**
     * Format for encoding values to Kafka.
     */
    protected final EncodingFormat<SerializationSchema<RowData>> valueEncodingFormat;
    /**
     * Indices that determine the key fields and the source position in the consumed row.
     */
    protected final int[] keyProjection;
    /**
     * Indices that determine the value fields and the source position in the consumed row.
     */
    protected final int[] valueProjection;
    /**
     * Prefix that needs to be removed from fields when constructing the physical data type.
     */
    protected final @Nullable String keyPrefix;
    /**
     * The Kafka topic to write to.
     */
    protected final String topic;
    protected final String topicPattern;
    /**
     * Properties for the Kafka producer.
     */
    protected final Properties properties;
    /**
     * Partitioner to select Kafka partition for each item.
     */
    protected final @Nullable FlinkKafkaPartitioner<RowData> partitioner;

    // --------------------------------------------------------------------------------------------
    // Kafka-specific attributes
    // --------------------------------------------------------------------------------------------
    /**
     * Sink commit semantic.
     */
    protected final KafkaSinkSemantic semantic;
    /**
     * Flag to determine sink mode. In upsert mode sink transforms the delete/update-before message
     * to tombstone message.
     */
    protected final boolean upsertMode;
    /**
     * Sink buffer flush config which only supported in upsert mode now.
     */
    protected final SinkBufferFlushMode flushMode;
    /**
     * Parallelism of the physical Kafka producer. *
     */
    protected final @Nullable Integer parallelism;
    /**
     * CatalogTable for KAFKA_IGNORE_ALL_CHANGELOG
     */
    private final CatalogTable catalogTable;
    /**
     * Metric for inLong
     */
    private final String inlongMetric;
    /**
     * audit host and ports
     */
    private final String auditHostAndPorts;
    private @Nullable final String sinkMultipleFormat;
    private final DirtyOptions dirtyOptions;
    private @Nullable final DirtySink<Object> dirtySink;
    /**
     * Metadata that is appended at the end of a physical sink row.
     */
    protected List<String> metadataKeys;
    /**
     * Data type of consumed data type.
     */
    protected DataType consumedDataType;

    private boolean multipleSink;
    private Map<SchemaChangeType, SchemaChangePolicy> policyMap;

    /**
     * Constructor of KafkaDynamicSink.
     */
    public KafkaDynamicSink(
            DataType consumedDataType,
            DataType physicalDataType,
            @Nullable EncodingFormat<SerializationSchema<RowData>> keyEncodingFormat,
            EncodingFormat<SerializationSchema<RowData>> valueEncodingFormat,
            int[] keyProjection,
            int[] valueProjection,
            @Nullable String keyPrefix,
            String topic,
            Properties properties,
            CatalogTable table,
            @Nullable FlinkKafkaPartitioner<RowData> partitioner,
            KafkaSinkSemantic semantic,
            boolean upsertMode,
            SinkBufferFlushMode flushMode,
            @Nullable Integer parallelism,
            String inlongMetric,
            String auditHostAndPorts,
            @Nullable String sinkMultipleFormat,
            @Nullable String topicPattern,
            DirtyOptions dirtyOptions,
            @Nullable DirtySink<Object> dirtySink,
            boolean multipleSink,
            Map<SchemaChangeType, SchemaChangePolicy> policyMap) {
        // Format attributes
        this.consumedDataType =
                checkNotNull(consumedDataType, "Consumed data type must not be null.");
        this.physicalDataType =
                checkNotNull(physicalDataType, "Physical data type must not be null.");
        this.keyEncodingFormat = keyEncodingFormat;
        this.valueEncodingFormat =
                checkNotNull(valueEncodingFormat, "Value encoding format must not be null.");
        this.keyProjection = checkNotNull(keyProjection, "Key projection must not be null.");
        this.valueProjection = checkNotNull(valueProjection, "Value projection must not be null.");
        this.keyPrefix = keyPrefix;
        // Mutable attributes
        this.metadataKeys = Collections.emptyList();
        // Kafka-specific attributes
        this.topic = checkNotNull(topic, "Topic must not be null.");
        this.properties = checkNotNull(properties, "Properties must not be null.");
        this.catalogTable = table;
        this.partitioner = partitioner;
        this.semantic = checkNotNull(semantic, "Semantic must not be null.");
        this.upsertMode = upsertMode;
        this.flushMode = checkNotNull(flushMode);
        if (flushMode.isEnabled() && !upsertMode) {
            throw new IllegalArgumentException(
                    "Sink buffer flush is only supported in upsert-kafka.");
        }
        this.parallelism = parallelism;
        this.inlongMetric = inlongMetric;
        this.auditHostAndPorts = auditHostAndPorts;
        this.sinkMultipleFormat = sinkMultipleFormat;
        this.topicPattern = topicPattern;
        this.dirtyOptions = dirtyOptions;
        this.dirtySink = dirtySink;
        this.multipleSink = multipleSink;
        this.policyMap = policyMap;
    }

    @Override
    public ChangelogMode getChangelogMode(ChangelogMode requestedMode) {
        if (org.apache.flink.configuration.Configuration.fromMap(catalogTable.getOptions())
                .get(KAFKA_IGNORE_ALL_CHANGELOG)) {
            LOG.warn("Kafka sink receive all changelog record. "
                    + "Regard any other record as insert-only record.");
            return ChangelogMode.all();
        }
        return valueEncodingFormat.getChangelogMode();
    }

    @Override
    public SinkRuntimeProvider getSinkRuntimeProvider(Context context) {
        final SerializationSchema<RowData> keySerialization =
                createSerialization(context, keyEncodingFormat, keyProjection, keyPrefix);

        final SerializationSchema<RowData> valueSerialization =
                createSerialization(context, valueEncodingFormat, valueProjection, null);

        final FlinkKafkaProducer<RowData> kafkaProducer =
                createKafkaProducer(keySerialization, valueSerialization, sinkMultipleFormat);

        if (flushMode.isEnabled() && upsertMode) {
            BufferedUpsertSinkFunction buffedSinkFunction =
                    new BufferedUpsertSinkFunction(
                            kafkaProducer,
                            physicalDataType,
                            keyProjection,
                            context.createTypeInformation(consumedDataType),
                            flushMode);
            return SinkFunctionProvider.of(buffedSinkFunction, parallelism);
        } else {
            return SinkFunctionProvider.of(kafkaProducer, parallelism);
        }
    }

    @Override
    public Map<String, DataType> listWritableMetadata() {
        final Map<String, DataType> metadataMap = new LinkedHashMap<>();

        // according to convention, the order of the final row must be
        // PHYSICAL + FORMAT METADATA + CONNECTOR METADATA
        // where the format metadata has highest precedence

        // add value format metadata with prefix
        valueEncodingFormat
                .listWritableMetadata()
                .forEach((key, value) -> metadataMap.put(VALUE_METADATA_PREFIX + key, value));

        // add connector metadata
        Stream.of(WritableMetadata.values())
                .forEachOrdered(m -> metadataMap.put(m.key, m.dataType));
        return metadataMap;
    }

    @Override
    public void applyWritableMetadata(List<String> metadataKeys, DataType consumedDataType) {
        // separate connector and format metadata
        final List<String> formatMetadataKeys =
                metadataKeys.stream()
                        .filter(k -> k.startsWith(VALUE_METADATA_PREFIX))
                        .collect(Collectors.toList());
        final List<String> connectorMetadataKeys = new ArrayList<>(metadataKeys);
        connectorMetadataKeys.removeAll(formatMetadataKeys);

        // push down format metadata
        final Map<String, DataType> formatMetadata = valueEncodingFormat.listWritableMetadata();
        if (formatMetadata.size() > 0) {
            final List<String> requestedFormatMetadataKeys =
                    formatMetadataKeys.stream()
                            .map(k -> k.substring(VALUE_METADATA_PREFIX.length()))
                            .collect(Collectors.toList());
            valueEncodingFormat.applyWritableMetadata(requestedFormatMetadataKeys);
        }

        this.metadataKeys = connectorMetadataKeys;
        this.consumedDataType = consumedDataType;
    }

    @Override
    public DynamicTableSink copy() {
        final KafkaDynamicSink copy =
                new KafkaDynamicSink(
                        consumedDataType,
                        physicalDataType,
                        keyEncodingFormat,
                        valueEncodingFormat,
                        keyProjection,
                        valueProjection,
                        keyPrefix,
                        topic,
                        properties,
                        catalogTable,
                        partitioner,
                        semantic,
                        upsertMode,
                        flushMode,
                        parallelism,
                        inlongMetric,
                        auditHostAndPorts,
                        sinkMultipleFormat,
                        topicPattern,
                        dirtyOptions,
                        dirtySink,
                        multipleSink,
                        policyMap);
        copy.metadataKeys = metadataKeys;
        return copy;
    }

    @Override
    public String asSummaryString() {
        return "Kafka table sink";
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        final KafkaDynamicSink that = (KafkaDynamicSink) o;
        return Objects.equals(metadataKeys, that.metadataKeys)
                && Objects.equals(consumedDataType, that.consumedDataType)
                && Objects.equals(physicalDataType, that.physicalDataType)
                && Objects.equals(keyEncodingFormat, that.keyEncodingFormat)
                && Objects.equals(valueEncodingFormat, that.valueEncodingFormat)
                && Arrays.equals(keyProjection, that.keyProjection)
                && Arrays.equals(valueProjection, that.valueProjection)
                && Objects.equals(keyPrefix, that.keyPrefix)
                && Objects.equals(topic, that.topic)
                && Objects.equals(properties, that.properties)
                && Objects.equals(partitioner, that.partitioner)
                && Objects.equals(semantic, that.semantic)
                && Objects.equals(upsertMode, that.upsertMode)
                && Objects.equals(flushMode, that.flushMode)
                && Objects.equals(parallelism, that.parallelism)
                && Objects.equals(inlongMetric, that.inlongMetric)
                && Objects.equals(auditHostAndPorts, that.auditHostAndPorts);
    }

    @Override
    public int hashCode() {
        return Objects.hash(
                metadataKeys,
                consumedDataType,
                physicalDataType,
                keyEncodingFormat,
                valueEncodingFormat,
                keyProjection,
                valueProjection,
                keyPrefix,
                topic,
                properties,
                partitioner,
                semantic,
                upsertMode,
                flushMode,
                parallelism,
                inlongMetric,
                auditHostAndPorts,
                sinkMultipleFormat,
                topicPattern);
    }

    // --------------------------------------------------------------------------------------------

    protected FlinkKafkaProducer<RowData> createKafkaProducer(
            SerializationSchema<RowData> keySerialization,
            SerializationSchema<RowData> valueSerialization,
            String sinkMultipleFormat) {
        final List<LogicalType> physicalChildren = physicalDataType.getLogicalType().getChildren();

        final RowData.FieldGetter[] keyFieldGetters =
                Arrays.stream(keyProjection)
                        .mapToObj(
                                targetField -> RowData.createFieldGetter(
                                        physicalChildren.get(targetField), targetField))
                        .toArray(RowData.FieldGetter[]::new);

        final RowData.FieldGetter[] valueFieldGetters =
                Arrays.stream(valueProjection)
                        .mapToObj(
                                targetField -> RowData.createFieldGetter(
                                        physicalChildren.get(targetField), targetField))
                        .toArray(RowData.FieldGetter[]::new);

        // determine the positions of metadata in the consumed row
        final int[] metadataPositions =
                Stream.of(WritableMetadata.values())
                        .mapToInt(
                                m -> {
                                    final int pos = metadataKeys.indexOf(m.key);
                                    if (pos < 0) {
                                        return -1;
                                    }
                                    return physicalChildren.size() + pos;
                                })
                        .toArray();

        // check if metadata is used at all
        final boolean hasMetadata = metadataKeys.size() > 0;

        final DynamicKafkaSerializationSchema kafkaSerializer =
                new DynamicKafkaSerializationSchema(
                        topic,
                        partitioner,
                        keySerialization,
                        valueSerialization,
                        keyFieldGetters,
                        valueFieldGetters,
                        hasMetadata,
                        metadataPositions,
                        upsertMode,
                        sinkMultipleFormat,
                        topicPattern,
                        dirtyOptions,
                        dirtySink,
                        policyMap);

        return new FlinkKafkaProducer<>(
                topic,
                kafkaSerializer,
                properties,
                FlinkKafkaProducer.Semantic.valueOf(semantic.toString()),
                FlinkKafkaProducer.DEFAULT_KAFKA_PRODUCERS_POOL_SIZE,
                inlongMetric,
                auditHostAndPorts,
                multipleSink);
    }

    private @Nullable SerializationSchema<RowData> createSerialization(
            Context context,
            @Nullable EncodingFormat<SerializationSchema<RowData>> format,
            int[] projection,
            @Nullable String prefix) {
        if (format == null) {
            return null;
        }
        DataType physicalFormatDataType =
                DataTypeUtils.projectRow(this.physicalDataType, projection);
        if (prefix != null) {
            physicalFormatDataType = DataTypeUtils.stripRowPrefix(physicalFormatDataType, prefix);
        }
        return format.createRuntimeEncoder(context, physicalFormatDataType);
    }

    // --------------------------------------------------------------------------------------------
    // Metadata handling
    // --------------------------------------------------------------------------------------------

    enum WritableMetadata {

        HEADERS(
                "headers",
                // key and value of the map are nullable to make handling easier in queries
                DataTypes.MAP(DataTypes.STRING().nullable(), DataTypes.BYTES().nullable())
                        .nullable(),
                new MetadataConverter() {

                    private static final long serialVersionUID = 1L;

                    @Override
                    public Object read(RowData row, int pos) {
                        if (row.isNullAt(pos)) {
                            return null;
                        }
                        final MapData map = row.getMap(pos);
                        final ArrayData keyArray = map.keyArray();
                        final ArrayData valueArray = map.valueArray();
                        final List<Header> headers = new ArrayList<>();
                        for (int i = 0; i < keyArray.size(); i++) {
                            if (!keyArray.isNullAt(i) && !valueArray.isNullAt(i)) {
                                final String key = keyArray.getString(i).toString();
                                final byte[] value = valueArray.getBinary(i);
                                headers.add(new KafkaHeader(key, value));
                            }
                        }
                        return headers;
                    }
                }),

        TIMESTAMP(
                "timestamp",
                DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3).nullable(),
                new MetadataConverter() {

                    private static final long serialVersionUID = 1L;

                    @Override
                    public Object read(RowData row, int pos) {
                        if (row.isNullAt(pos)) {
                            return null;
                        }
                        return row.getTimestamp(pos, 3).getMillisecond();
                    }
                });

        final String key;

        final DataType dataType;

        final MetadataConverter converter;

        WritableMetadata(String key, DataType dataType, MetadataConverter converter) {
            this.key = key;
            this.dataType = dataType;
            this.converter = converter;
        }
    }

    // --------------------------------------------------------------------------------------------

    private static class KafkaHeader implements Header {

        private final String key;

        private final byte[] value;

        KafkaHeader(String key, byte[] value) {
            this.key = key;
            this.value = value;
        }

        @Override
        public String key() {
            return key;
        }

        @Override
        public byte[] value() {
            return value;
        }
    }
}
