/*
 * Copyright 2020 Confluent Inc.
 *
 * Licensed under the Confluent Community License (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.confluent.io/confluent-community-license
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OF ANY KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations under the License.
 */

package io.confluent.connect.json;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.collect.ImmutableMap;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.data.SchemaAndValue;
import org.apache.kafka.connect.data.SchemaBuilder;
import org.apache.kafka.connect.data.Struct;
import org.junit.Before;
import org.junit.Test;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import io.confluent.kafka.schemaregistry.client.MockSchemaRegistryClient;
import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient;
import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException;
import io.confluent.kafka.schemaregistry.json.JsonSchema;
import io.confluent.kafka.schemaregistry.json.JsonSchemaUtils;
import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig;
import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer;

import static org.junit.Assert.assertEquals;

public class JsonSchemaConverterTest {

  private final ResourceLoader loader = ResourceLoader.DEFAULT;

  private static final String TOPIC = "topic";

  private static final Map<String, ?> SR_CONFIG =
      Collections.singletonMap(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG,
      "localhost"
  );

  private final SchemaRegistryClient schemaRegistry;
  private final JsonSchemaConverter converter;

  public JsonSchemaConverterTest() {
    schemaRegistry = new MockSchemaRegistryClient();
    converter = new JsonSchemaConverter(schemaRegistry);
  }

  @Before
  public void setUp() {
    Map<String, String> config = new HashMap<>();
    config.put("schema.registry.url", "http://fake-url");
    converter.configure(config, false);
  }

  @Test
  public void testPrimitive() {
    SchemaAndValue original = new SchemaAndValue(Schema.BOOLEAN_SCHEMA, true);
    byte[] converted = converter.fromConnectData(TOPIC, original.schema(), original.value());
    SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, converted);
    // Because of registration in schema registry and lookup, we'll have added a version number
    SchemaAndValue expected = new SchemaAndValue(SchemaBuilder.bool().version(1).build(), true);
    assertEquals(expected, schemaAndValue);
  }

  @Test
  public void testComplex() {
    SchemaBuilder builder = SchemaBuilder.struct()
        .field("int8", Schema.INT8_SCHEMA)
        .field("int16", Schema.INT16_SCHEMA)
        .field("int32", Schema.INT32_SCHEMA)
        .field("int64", Schema.INT64_SCHEMA)
        .field("float32", Schema.FLOAT32_SCHEMA)
        .field("float64", Schema.FLOAT64_SCHEMA)
        .field("boolean", Schema.BOOLEAN_SCHEMA)
        .field("string", Schema.STRING_SCHEMA)
        .field("bytes", Schema.BYTES_SCHEMA)
        .field("array", SchemaBuilder.array(Schema.STRING_SCHEMA).build())
        .field("map", SchemaBuilder.map(Schema.STRING_SCHEMA, Schema.INT32_SCHEMA).build());
    Schema schema = builder.build();
    Struct original = new Struct(schema).put("int8", (byte) 12)
        .put("int16", (short) 12)
        .put("int32", 12)
        .put("int64", 12L)
        .put("float32", 12.2f)
        .put("float64", 12.2)
        .put("boolean", true)
        .put("string", "foo")
        .put("bytes", "foo".getBytes())
        .put("array", Arrays.asList("a", "b", "c"))
        .put("map", Collections.singletonMap("field", 1));
    // Because of registration in schema registry and lookup, we'll have added a version number
    Schema expectedSchema = builder.version(1).build();
    Struct expected = new Struct(expectedSchema).put("int8", (byte) 12)
        .put("int16", (short) 12)
        .put("int32", 12)
        .put("int64", 12L)
        .put("float32", 12.2f)
        .put("float64", 12.2)
        .put("boolean", true)
        .put("string", "foo")
        .put("bytes", "foo".getBytes())
        .put("array", Arrays.asList("a", "b", "c"))
        .put("map", Collections.singletonMap("field", 1));

    byte[] converted = converter.fromConnectData(TOPIC, original.schema(), original);
    SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, converted);
    assertEquals(expected, schemaAndValue.value());
  }

  @Test
  public void testNull() {
    byte[] converted = converter.fromConnectData(TOPIC, Schema.OPTIONAL_BOOLEAN_SCHEMA, null);
    SchemaAndValue schemaAndValue = converter.toConnectData(TOPIC, converted);
    assertEquals(new SchemaAndValue(SchemaBuilder.bool().version(1).optional().build(), null),
        schemaAndValue
    );
  }

  @Test
  public void testVersionExtractedForDefaultSubjectNameStrategy() throws Exception {
    // Version info should be extracted even if the data was not created with Copycat. Manually
    // register a few compatible schemas and validate that data serialized with our normal
    // serializer can be read and gets version info inserted
    String subject = TOPIC + "-value";
    KafkaJsonSchemaSerializer serializer = new KafkaJsonSchemaSerializer(schemaRegistry,
        ImmutableMap.of("schema.registry.url", "http://fake-url")
    );
    JsonSchemaConverter jsonConverter = new JsonSchemaConverter(schemaRegistry);
    jsonConverter.configure(Collections.singletonMap("schema.registry.url", "http://fake-url"),
        false
    );
    testVersionExtracted(subject, serializer, jsonConverter);
  }

  private void testVersionExtracted(
      String subject, KafkaJsonSchemaSerializer serializer, JsonSchemaConverter jsonConverter
  ) throws IOException, RestClientException {
    JsonNode rawSchemaJson1 = loader.readJsonNode("key.json");
    JsonNode rawSchemaJson2 = loader.readJsonNode("keyvalue.json");
    schemaRegistry.register(subject, new JsonSchema(rawSchemaJson1));
    schemaRegistry.register(subject, new JsonSchema(rawSchemaJson2));

    ObjectMapper mapper = new ObjectMapper();

    ObjectNode objectNode1 = mapper.createObjectNode();
    objectNode1.put("key", 15);

    ObjectNode objectNode2 = mapper.createObjectNode();
    objectNode2.put("key", 15);
    objectNode2.put("value", "bar");

    // Get serialized data
    byte[] serializedRecord1 = serializer.serialize(TOPIC,
        JsonSchemaUtils.envelope(rawSchemaJson1, objectNode1)
    );
    byte[] serializedRecord2 = serializer.serialize(TOPIC,
        JsonSchemaUtils.envelope(rawSchemaJson2, objectNode2)
    );

    SchemaAndValue converted1 = jsonConverter.toConnectData(TOPIC, serializedRecord1);
    assertEquals(1L, (long) converted1.schema().version());

    SchemaAndValue converted2 = jsonConverter.toConnectData(TOPIC, serializedRecord2);
    assertEquals(2L, (long) converted2.schema().version());
  }

  @Test
  public void testVersionMaintained() {
    // Version info provided from the Copycat schema should be maintained. This should be true
    // regardless of any underlying schema registry versioning since the versions are explicitly
    // specified by the connector.

    // Use newer schema first
    Schema newerSchema = SchemaBuilder.struct()
        .version(2)
        .field("orig", Schema.OPTIONAL_INT16_SCHEMA)
        .field("new", Schema.OPTIONAL_INT16_SCHEMA)
        .build();
    SchemaAndValue newer = new SchemaAndValue(newerSchema,
        new Struct(newerSchema).put("orig", (short) 1).put("new", (short) 2)
    );
    byte[] newerSerialized = converter.fromConnectData(TOPIC, newer.schema(), newer.value());

    Schema olderSchema = SchemaBuilder.struct()
        .version(1)
        .field("orig", Schema.OPTIONAL_INT16_SCHEMA)
        .build();
    SchemaAndValue older = new SchemaAndValue(olderSchema,
        new Struct(olderSchema).put("orig", (short) 1)
    );
    byte[] olderSerialized = converter.fromConnectData(TOPIC, older.schema(), older.value());

    assertEquals(2L, (long) converter.toConnectData(TOPIC, newerSerialized).schema().version());
    assertEquals(1L, (long) converter.toConnectData(TOPIC, olderSerialized).schema().version());
  }

  @Test
  public void testSameSchemaMultipleTopicForValue() throws IOException, RestClientException {
    SchemaRegistryClient schemaRegistry = new MockSchemaRegistryClient();
    JsonSchemaConverter jsonConverter = new JsonSchemaConverter(schemaRegistry);
    jsonConverter.configure(SR_CONFIG, false);
    assertSameSchemaMultipleTopic(jsonConverter, schemaRegistry, false);
  }

  @Test
  public void testSameSchemaMultipleTopicForKey() throws IOException, RestClientException {
    SchemaRegistryClient schemaRegistry = new MockSchemaRegistryClient();
    JsonSchemaConverter jsonConverter = new JsonSchemaConverter(schemaRegistry);
    jsonConverter.configure(SR_CONFIG, true);
    assertSameSchemaMultipleTopic(jsonConverter, schemaRegistry, true);
  }

  @Test
  public void testExplicitlyNamedNestedMapsWithNonStringKeys() {
    final Schema schema = SchemaBuilder.map(Schema.OPTIONAL_STRING_SCHEMA,
        SchemaBuilder.map(Schema.OPTIONAL_STRING_SCHEMA, Schema.OPTIONAL_INT32_SCHEMA)
            .name("foo.bar")
            .build()
    ).name("biz.baz").version(1).build();
    final JsonSchemaConverter jsonConverter =
        new JsonSchemaConverter(new MockSchemaRegistryClient());
    jsonConverter.configure(Collections.singletonMap(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG,
        "localhost"
    ), false);
    final Object value = Collections.singletonMap("foo", Collections.singletonMap("bar", 1));

    final byte[] bytes = jsonConverter.fromConnectData("topic", schema, value);
    final SchemaAndValue schemaAndValue = jsonConverter.toConnectData("topic", bytes);

    assertEquals(schemaAndValue.schema(), schema);
    assertEquals(schemaAndValue.value(), value);
  }

  private void assertSameSchemaMultipleTopic(
      JsonSchemaConverter converter, SchemaRegistryClient schemaRegistry, boolean isKey
  ) throws IOException, RestClientException {
    JsonNode rawSchemaJson1 = loader.readJsonNode("key.json");
    JsonNode rawSchemaJson2_1 = loader.readJsonNode("keyvalue.json");
    JsonNode rawSchemaJson2_2 = loader.readJsonNode("keyvalue.json");
    String subjectSuffix = isKey ? "key" : "value";
    schemaRegistry.register("topic1-" + subjectSuffix, new JsonSchema(rawSchemaJson2_1));
    schemaRegistry.register("topic2-" + subjectSuffix, new JsonSchema(rawSchemaJson1));
    schemaRegistry.register("topic2-" + subjectSuffix, new JsonSchema(rawSchemaJson2_2));

    ObjectMapper mapper = new ObjectMapper();

    ObjectNode objectNode1 = mapper.createObjectNode();
    objectNode1.put("key", 15);
    objectNode1.put("value", "bar");

    ObjectNode objectNode2 = mapper.createObjectNode();
    objectNode2.put("key", 15);
    objectNode2.put("value", "bar");

    KafkaJsonSchemaSerializer serializer = new KafkaJsonSchemaSerializer(schemaRegistry,
        ImmutableMap.of("schema.registry.url", "http://fake-url")
    );
    byte[] serializedRecord1 = serializer.serialize("topic1",
        JsonSchemaUtils.envelope(rawSchemaJson2_1, objectNode1)
    );
    byte[] serializedRecord2 = serializer.serialize("topic2",
        JsonSchemaUtils.envelope(rawSchemaJson2_2, objectNode2)
    );

    SchemaAndValue converted1 = converter.toConnectData("topic1", serializedRecord1);
    assertEquals(1L, (long) converted1.schema().version());

    SchemaAndValue converted2 = converter.toConnectData("topic2", serializedRecord2);
    assertEquals(2L, (long) converted2.schema().version());

    converted2 = converter.toConnectData("topic2", serializedRecord2);
    assertEquals(2L, (long) converted2.schema().version());
  }
}
