/*
 * Copyright 2018 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.ksql.integration;

import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.hasSize;

import io.confluent.common.utils.IntegrationTest;
import io.confluent.ksql.GenericRow;
import io.confluent.ksql.KsqlContextTestUtil;
import io.confluent.ksql.KsqlEngine;
import io.confluent.ksql.KsqlEngineTestUtil;
import io.confluent.ksql.function.InternalFunctionRegistry;
import io.confluent.ksql.logging.processing.ProcessingLogContext;
import io.confluent.ksql.metastore.MetaStore;
import io.confluent.ksql.query.QueryId;
import io.confluent.ksql.services.DefaultServiceContext;
import io.confluent.ksql.services.KafkaTopicClient;
import io.confluent.ksql.services.ServiceContext;
import io.confluent.ksql.test.util.EmbeddedSingleNodeKafkaCluster;
import io.confluent.ksql.util.KsqlConfig;
import io.confluent.ksql.util.OrderDataProvider;
import io.confluent.ksql.util.PersistentQueryMetadata;
import io.confluent.ksql.util.QueryMetadata;
import io.confluent.ksql.util.SchemaUtil;
import io.confluent.ksql.util.TopicConsumer;
import io.confluent.ksql.util.TopicProducer;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.data.SchemaBuilder;
import org.apache.kafka.test.TestUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.experimental.categories.Category;

@Category({IntegrationTest.class})
public class JsonFormatTest {
  private static final String inputTopic = "orders_topic";
  private static final String inputStream = "ORDERS";
  private static final String usersTopic = "users_topic";
  private static final String usersTable = "USERS";
  private static final String messageLogTopic = "log_topic";
  private static final String messageLogStream = "message_log";
  private static final AtomicInteger COUNTER = new AtomicInteger();

  @ClassRule
  public static final EmbeddedSingleNodeKafkaCluster CLUSTER = EmbeddedSingleNodeKafkaCluster.build();

  private MetaStore metaStore;
  private KsqlConfig ksqlConfig;
  private KsqlEngine ksqlEngine;
  private ServiceContext serviceContext;
  private ProcessingLogContext processingLogContext;
  private final TopicProducer topicProducer = new TopicProducer(CLUSTER);
  private final TopicConsumer topicConsumer = new TopicConsumer(CLUSTER);

  private QueryId queryId;
  private KafkaTopicClient topicClient;
  private String streamName;

  @Before
  public void before() throws Exception {
    streamName = "STREAM_" + COUNTER.getAndIncrement();

    ksqlConfig = KsqlContextTestUtil.createKsqlConfig(CLUSTER);
    serviceContext = DefaultServiceContext.create(ksqlConfig);
    processingLogContext = ProcessingLogContext.create();
    ksqlEngine = new KsqlEngine(
        serviceContext,
        processingLogContext,
        new InternalFunctionRegistry(),
        ksqlConfig.getString(KsqlConfig.KSQL_SERVICE_ID_CONFIG));
    topicClient = serviceContext.getTopicClient();
    metaStore = ksqlEngine.getMetaStore();

    createInitTopics();
    produceInitData();
    execInitCreateStreamQueries();
  }

  private void createInitTopics() {
    topicClient.createTopic(inputTopic, 1, (short) 1);
    topicClient.createTopic(usersTopic, 1, (short) 1);
    topicClient.createTopic(messageLogTopic, 1, (short) 1);
  }

  private void produceInitData() throws Exception {
    final OrderDataProvider orderDataProvider = new OrderDataProvider();

    topicProducer
            .produceInputData(inputTopic, orderDataProvider.data(), orderDataProvider.schema());

    final Schema messageSchema = SchemaBuilder.struct().field("MESSAGE", SchemaBuilder.OPTIONAL_STRING_SCHEMA).build();

    final GenericRow messageRow = new GenericRow(Collections.singletonList(
        "{\"log\":{\"@timestamp\":\"2017-05-30T16:44:22.175Z\",\"@version\":\"1\","
        + "\"caasVersion\":\"0.0.2\",\"cloud\":\"aws\",\"logs\":[{\"entry\":\"first\"}],\"clusterId\":\"cp99\",\"clusterName\":\"kafka\",\"cpComponentId\":\"kafka\",\"host\":\"kafka-1-wwl0p\",\"k8sId\":\"k8s13\",\"k8sName\":\"perf\",\"level\":\"ERROR\",\"logger\":\"kafka.server.ReplicaFetcherThread\",\"message\":\"Found invalid messages during fetch for partition [foo512,172] offset 0 error Record is corrupt (stored crc = 1321230880, computed crc = 1139143803)\",\"networkId\":\"vpc-d8c7a9bf\",\"region\":\"us-west-2\",\"serverId\":\"1\",\"skuId\":\"sku5\",\"source\":\"kafka\",\"tenantId\":\"t47\",\"tenantName\":\"perf-test\",\"thread\":\"ReplicaFetcherThread-0-2\",\"zone\":\"us-west-2a\"},\"stream\":\"stdout\",\"time\":2017}"));

    final Map<String, GenericRow> records = new HashMap<>();
    records.put("1", messageRow);

    topicProducer.produceInputData(messageLogTopic, records, messageSchema);
  }

  private void execInitCreateStreamQueries() {
    final String ordersStreamStr = String.format("CREATE STREAM %s (ORDERTIME bigint, ORDERID varchar, "
        + "ITEMID varchar, ORDERUNITS double, PRICEARRAY array<double>, KEYVALUEMAP "
        + "map<varchar, double>) WITH (value_format = 'json', "
        + "kafka_topic='%s' , "
        + "key='ordertime');", inputStream, inputTopic);

    final String usersTableStr = String.format("CREATE TABLE %s (userid varchar, age integer) WITH "
                                         + "(value_format = 'json', kafka_topic='%s', "
                                         + "KEY='userid');",
                                         usersTable, usersTopic);

    final String messageStreamStr = String.format("CREATE STREAM %s (message varchar) WITH (value_format = 'json', "
        + "kafka_topic='%s');", messageLogStream, messageLogTopic);

    KsqlEngineTestUtil.execute(ksqlEngine, ordersStreamStr, ksqlConfig, Collections.emptyMap());
    KsqlEngineTestUtil.execute(ksqlEngine, usersTableStr, ksqlConfig, Collections.emptyMap());
    KsqlEngineTestUtil.execute(ksqlEngine, messageStreamStr, ksqlConfig, Collections.emptyMap());
  }

  @After
  public void after() {
    terminateQuery();
    ksqlEngine.close();
    serviceContext.close();
  }

  //@Test
  public void testSelectDateTimeUDFs() {
    final String streamName = "SelectDateTimeUDFsStream".toUpperCase();

    final String selectColumns =
        "(ORDERTIME+1500962514806) , TIMESTAMPTOSTRING(ORDERTIME+1500962514806, "
        + "'yyyy-MM-dd HH:mm:ss.SSS'), "
        + "STRINGTOTIMESTAMP"
            + "(TIMESTAMPTOSTRING"
            + "(ORDERTIME+1500962514806, 'yyyy-MM-dd HH:mm:ss.SSS'), 'yyyy-MM-dd HH:mm:ss.SSS')";
    final String whereClause = "ORDERUNITS > 20 AND ITEMID LIKE '%_8'";

    final String queryString = String.format(
        "CREATE STREAM %s AS SELECT %s FROM %s WHERE %s;",
        streamName,
        selectColumns,
        inputStream,
        whereClause
    );

    executePersistentQuery(queryString);

    final Schema resultSchema = SchemaUtil
        .removeImplicitRowTimeRowKeyFromSchema(metaStore.getSource(streamName).getSchema());

    final Map<String, GenericRow> expectedResults = new HashMap<>();
    expectedResults.put("8", new GenericRow(Arrays.asList(1500962514814L,
        "2017-07-24 23:01:54.814",
        1500962514814L)));

    final Map<String, GenericRow> results = readNormalResults(streamName, resultSchema, expectedResults.size());

    assertThat(results, equalTo(expectedResults));
  }

  @Test
  public void testSinkProperties() throws Exception {
    final int resultPartitionCount = 3;
    final String queryString = String.format("CREATE STREAM %s WITH (PARTITIONS = %d) AS SELECT * "
            + "FROM %s;",
        streamName, resultPartitionCount, inputStream);

    executePersistentQuery(queryString);

    TestUtils.waitForCondition(
        () -> topicClient.isTopicExists(streamName),
        "Wait for async topic creation"
    );

    assertThat(
        topicClient.describeTopic(streamName).partitions(),
        hasSize(3));
    assertThat(topicClient.getTopicCleanupPolicy(streamName), equalTo(
        KafkaTopicClient.TopicCleanupPolicy.DELETE));
  }

  @Test
  public void testTableSinkCleanupProperty() throws Exception {
    final String queryString = String.format("CREATE TABLE %s AS SELECT * "
                                             + "FROM %s;",
        streamName, usersTable);
    executePersistentQuery(queryString);

    TestUtils.waitForCondition(
        () -> topicClient.isTopicExists(streamName),
        "Wait for async topic creation"
    );

    assertThat(topicClient.getTopicCleanupPolicy(streamName), equalTo(
        KafkaTopicClient.TopicCleanupPolicy.COMPACT));
  }

  @Test
  public void testJsonStreamExtractor() {
    final String queryString = String.format("CREATE STREAM %s AS SELECT EXTRACTJSONFIELD"
            + "(message, '$.log.cloud') "
            + "FROM %s;",
        streamName, messageLogStream);

    executePersistentQuery(queryString);

    final Schema resultSchema = SchemaUtil
            .removeImplicitRowTimeRowKeyFromSchema(metaStore.getSource(streamName).getSchema());

    final Map<String, GenericRow> expectedResults = new HashMap<>();
    expectedResults.put("1", new GenericRow(Collections.singletonList("aws")));

    final Map<String, GenericRow> results = readNormalResults(streamName, resultSchema, expectedResults.size());

    assertThat(results, equalTo(expectedResults));
  }

  @Test
  public void testJsonStreamExtractorNested() {
    final String queryString = String.format("CREATE STREAM %s AS SELECT EXTRACTJSONFIELD"
                    + "(message, '$.log.logs[0].entry') "
                    + "FROM %s;",
            streamName, messageLogStream);

    executePersistentQuery(queryString);

    final Schema resultSchema = SchemaUtil
            .removeImplicitRowTimeRowKeyFromSchema(metaStore.getSource(streamName).getSchema());

    final Map<String, GenericRow> expectedResults = new HashMap<>();
    expectedResults.put("1", new GenericRow(Collections.singletonList("first")));

    final Map<String, GenericRow> results = readNormalResults(streamName, resultSchema, expectedResults.size());

    assertThat(results, equalTo(expectedResults));
  }

  private void executePersistentQuery(final String queryString) {
    final QueryMetadata queryMetadata = KsqlEngineTestUtil
        .execute(ksqlEngine, queryString, ksqlConfig, Collections.emptyMap()).get(0);

    queryMetadata.start();
    queryId = ((PersistentQueryMetadata)queryMetadata).getQueryId();
  }

  private Map<String, GenericRow> readNormalResults(final String resultTopic, final Schema resultSchema, final int expectedNumMessages) {
    return topicConsumer.readResults(resultTopic, resultSchema, expectedNumMessages, new StringDeserializer());
  }

  private void terminateQuery() {
    ksqlEngine.getPersistentQuery(queryId)
        .ifPresent(QueryMetadata::close);
  }

}
