/*
 * 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.kafka.source.reader.deserializer;

import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.connector.kafka.testutils.SimpleCollector;
import org.apache.flink.connector.kafka.util.JacksonMapperFactory;
import org.apache.flink.connector.testutils.formats.DummyInitializationContext;
import org.apache.flink.connector.testutils.source.deserialization.TestingDeserializationContext;
import org.apache.flink.formats.json.JsonDeserializationSchema;
import org.apache.flink.streaming.util.serialization.JSONKeyValueDeserializationSchema;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.Configurable;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

/** Unit tests for KafkaRecordDeserializationSchema. */
class KafkaRecordDeserializationSchemaTest {

    private static final ObjectMapper OBJECT_MAPPER = JacksonMapperFactory.createObjectMapper();

    private static Map<String, ?> configurableConfiguration;
    private static Map<String, ?> configuration;
    private static boolean isKeyDeserializer;

    @BeforeEach
    void setUp() {
        configurableConfiguration = new HashMap<>(1);
        configuration = new HashMap<>(1);
        isKeyDeserializer = false;
    }

    @Test
    void testKafkaDeserializationSchemaWrapper() throws Exception {
        final ConsumerRecord<byte[], byte[]> consumerRecord = getConsumerRecord();
        KafkaRecordDeserializationSchema<ObjectNode> schema =
                new JSONKeyValueDeserializationSchema(true);
        schema.open(new DummyInitializationContext());
        SimpleCollector<ObjectNode> collector = new SimpleCollector<>();
        schema.deserialize(consumerRecord, collector);

        assertThat(collector.getList()).hasSize(1);
        ObjectNode deserializedValue = collector.getList().get(0);

        assertThat(deserializedValue.get("key").get("index").asInt()).isEqualTo(4);
        assertThat(deserializedValue.get("value").get("word").asText()).isEqualTo("world");
        assertThat(deserializedValue.get("metadata").get("topic").asText()).isEqualTo("topic#1");
        assertThat(deserializedValue.get("metadata").get("offset").asInt()).isEqualTo(4);
        assertThat(deserializedValue.get("metadata").get("partition").asInt()).isEqualTo(3);
    }

    @Test
    void testKafkaValueDeserializationSchemaWrapper() throws Exception {
        final ConsumerRecord<byte[], byte[]> consumerRecord = getConsumerRecord();
        KafkaRecordDeserializationSchema<Map<String, Object>> schema =
                KafkaRecordDeserializationSchema.valueOnly(
                        new JsonDeserializationSchema<>(
                                TypeInformation.of(new TypeHint<Map<String, Object>>() {})));
        schema.open(new DummyInitializationContext());
        SimpleCollector<Map<String, Object>> collector = new SimpleCollector<>();
        schema.deserialize(consumerRecord, collector);

        assertThat(collector.getList()).hasSize(1);
        Map<String, Object> deserializedValue = collector.getList().get(0);
        assertThat(deserializedValue.get("word")).isEqualTo("world");
        assertThat(deserializedValue.get("key")).isNull();
        assertThat(deserializedValue.get("metadata")).isNull();
    }

    @Test
    void testKafkaValueDeserializerWrapper() throws Exception {
        final String topic = "Topic";
        byte[] value = new StringSerializer().serialize(topic, "world");
        final ConsumerRecord<byte[], byte[]> consumerRecord =
                new ConsumerRecord<>(topic, 0, 0L, null, value);
        KafkaRecordDeserializationSchema<String> schema =
                KafkaRecordDeserializationSchema.valueOnly(StringDeserializer.class);
        schema.open(new TestingDeserializationContext());

        SimpleCollector<String> collector = new SimpleCollector<>();
        schema.deserialize(consumerRecord, collector);

        assertThat(collector.getList()).hasSize(1);
        assertThat(collector.getList().get(0)).isEqualTo("world");
    }

    @Test
    void testKafkaValueDeserializerWrapperWithoutConfigurable() throws Exception {
        final Map<String, String> config = Collections.singletonMap("simpleKey", "simpleValue");
        KafkaRecordDeserializationSchema<String> schema =
                KafkaRecordDeserializationSchema.valueOnly(SimpleStringSerializer.class, config);
        schema.open(new TestingDeserializationContext());
        assertThat(config).isEqualTo(configuration);
        assertThat(isKeyDeserializer).isFalse();
        assertThat(configurableConfiguration).isEmpty();
    }

    @Test
    void testKafkaValueDeserializerWrapperWithConfigurable() throws Exception {
        final Map<String, String> config = Collections.singletonMap("configKey", "configValue");
        KafkaRecordDeserializationSchema<String> schema =
                KafkaRecordDeserializationSchema.valueOnly(
                        ConfigurableStringSerializer.class, config);
        schema.open(new TestingDeserializationContext());
        assertThat(config).isEqualTo(configurableConfiguration);
        assertThat(isKeyDeserializer).isFalse();
        assertThat(configuration).isEmpty();
    }

    private ConsumerRecord<byte[], byte[]> getConsumerRecord() throws JsonProcessingException {
        ObjectNode initialKey = OBJECT_MAPPER.createObjectNode();
        initialKey.put("index", 4);
        byte[] serializedKey = OBJECT_MAPPER.writeValueAsBytes(initialKey);

        ObjectNode initialValue = OBJECT_MAPPER.createObjectNode();
        initialValue.put("word", "world");
        byte[] serializedValue = OBJECT_MAPPER.writeValueAsBytes(initialValue);

        return new ConsumerRecord<>("topic#1", 3, 4L, serializedKey, serializedValue);
    }

    /**
     * Serializer based on Kafka's serialization stack. This is the special case that implements
     * {@link Configurable}
     *
     * <p>This class must be public to make it instantiable by the tests.
     */
    public static class ConfigurableStringSerializer extends StringDeserializer
            implements Configurable {
        @Override
        public void configure(Map<String, ?> configs) {
            configurableConfiguration = configs;
        }
    }

    /**
     * Serializer based on Kafka's serialization stack.
     *
     * <p>This class must be public to make it instantiable by the tests.
     */
    public static class SimpleStringSerializer extends StringDeserializer {
        @Override
        public void configure(Map<String, ?> configs, boolean isKey) {
            configuration = configs;
            isKeyDeserializer = isKey;
        }
    }
}
