/*
 * 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.beam.sdk.io.clickhouse;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;

import java.sql.ResultSet;
import java.util.Arrays;
import java.util.Objects;
import org.apache.beam.sdk.io.clickhouse.TableSchema.ColumnType;
import org.apache.beam.sdk.schemas.JavaFieldSchema;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
import org.apache.beam.sdk.schemas.logicaltypes.FixedBytes;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.Row;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link ClickHouseIO}. */
@RunWith(JUnit4.class)
public class ClickHouseIOTest extends BaseClickHouseTest {

  @Rule public TestPipeline pipeline = TestPipeline.create();

  @Test
  public void testInt64() throws Exception {
    Schema schema =
        Schema.of(Schema.Field.of("f0", FieldType.INT64), Schema.Field.of("f1", FieldType.INT64));
    Row row1 = Row.withSchema(schema).addValue(1L).addValue(2L).build();
    Row row2 = Row.withSchema(schema).addValue(2L).addValue(4L).build();
    Row row3 = Row.withSchema(schema).addValue(3L).addValue(6L).build();

    executeSql("CREATE TABLE test_int64 (f0 Int64, f1 Int64) ENGINE=Log");

    pipeline.apply(Create.of(row1, row2, row3).withRowSchema(schema)).apply(write("test_int64"));

    pipeline.run().waitUntilFinish();

    long sum0 = executeQueryAsLong("SELECT SUM(f0) FROM test_int64");
    long sum1 = executeQueryAsLong("SELECT SUM(f1) FROM test_int64");

    assertEquals(6L, sum0);
    assertEquals(12L, sum1);
  }

  @Test
  public void testNullableInt64() throws Exception {
    Schema schema = Schema.of(Schema.Field.nullable("f0", FieldType.INT64));
    Row row1 = Row.withSchema(schema).addValue(1L).build();
    Row row2 = Row.withSchema(schema).addValue(null).build();
    Row row3 = Row.withSchema(schema).addValue(3L).build();

    executeSql("CREATE TABLE test_nullable_int64 (f0 Nullable(Int64)) ENGINE=Log");

    pipeline
        .apply(Create.of(row1, row2, row3).withRowSchema(schema))
        .apply(write("test_nullable_int64"));

    pipeline.run().waitUntilFinish();

    long sum = executeQueryAsLong("SELECT SUM(f0) FROM test_nullable_int64");
    long count0 = executeQueryAsLong("SELECT COUNT(*) FROM test_nullable_int64");
    long count1 = executeQueryAsLong("SELECT COUNT(f0) FROM test_nullable_int64");

    assertEquals(4L, sum);
    assertEquals(3L, count0);
    assertEquals(2L, count1);
  }

  @Test
  public void testInt64WithDefault() throws Exception {
    Schema schema = Schema.of(Schema.Field.nullable("f0", FieldType.INT64));
    Row row1 = Row.withSchema(schema).addValue(1L).build();
    Row row2 = Row.withSchema(schema).addValue(null).build();
    Row row3 = Row.withSchema(schema).addValue(3L).build();

    executeSql("CREATE TABLE test_int64_with_default (f0 Int64 DEFAULT -1) ENGINE=Log");

    pipeline
        .apply(Create.of(row1, row2, row3).withRowSchema(schema))
        .apply(write("test_int64_with_default"));

    pipeline.run().waitUntilFinish();

    long sum = executeQueryAsLong("SELECT SUM(f0) FROM test_int64_with_default");

    assertEquals(3L, sum);
  }

  @Test
  public void testArrayOfArrayOfInt64() throws Exception {
    Schema schema =
        Schema.of(Schema.Field.of("f0", FieldType.array(FieldType.array(FieldType.INT64))));
    Row row1 =
        Row.withSchema(schema)
            .addValue(
                Arrays.asList(Arrays.asList(1L, 2L), Arrays.asList(2L, 3L), Arrays.asList(3L, 4L)))
            .build();

    executeSql("CREATE TABLE test_array_of_array_of_int64 (f0 Array(Array(Int64))) ENGINE=Log");

    pipeline
        .apply(Create.of(row1).withRowSchema(schema))
        .apply(write("test_array_of_array_of_int64"));

    pipeline.run().waitUntilFinish();

    long sum0 =
        executeQueryAsLong(
            "SELECT SUM(arraySum(arrayMap(x -> arraySum(x), f0))) "
                + "FROM test_array_of_array_of_int64");

    assertEquals(15L, sum0);
  }

  @Test
  public void testPrimitiveTypes() throws Exception {
    Schema schema =
        Schema.of(
            Schema.Field.of("f0", FieldType.DATETIME),
            Schema.Field.of("f1", FieldType.DATETIME),
            Schema.Field.of("f2", FieldType.FLOAT),
            Schema.Field.of("f3", FieldType.DOUBLE),
            Schema.Field.of("f4", FieldType.BYTE),
            Schema.Field.of("f5", FieldType.INT16),
            Schema.Field.of("f6", FieldType.INT32),
            Schema.Field.of("f7", FieldType.INT64),
            Schema.Field.of("f8", FieldType.STRING),
            Schema.Field.of("f9", FieldType.INT16),
            Schema.Field.of("f10", FieldType.INT32),
            Schema.Field.of("f11", FieldType.INT64),
            Schema.Field.of("f12", FieldType.INT64),
            Schema.Field.of("f13", FieldType.STRING),
            Schema.Field.of("f14", FieldType.STRING),
            Schema.Field.of("f15", FieldType.STRING),
            Schema.Field.of("f16", FieldType.BYTES),
            Schema.Field.of("f17", FieldType.logicalType(FixedBytes.of(3))));
    Row row1 =
        Row.withSchema(schema)
            .addValue(new DateTime(2030, 10, 1, 0, 0, 0, DateTimeZone.UTC))
            .addValue(new DateTime(2030, 10, 9, 8, 7, 6, DateTimeZone.UTC))
            .addValue(2.2f)
            .addValue(3.3)
            .addValue((byte) 4)
            .addValue((short) 5)
            .addValue(6)
            .addValue(7L)
            .addValue("eight")
            .addValue((short) 9)
            .addValue(10)
            .addValue(11L)
            .addValue(12L)
            .addValue("abc")
            .addValue("cde")
            .addValue("qwe")
            .addValue(new byte[] {'a', 's', 'd'})
            .addValue(new byte[] {'z', 'x', 'c'})
            .build();

    executeSql(
        "CREATE TABLE test_primitive_types ("
            + "f0  Date,"
            + "f1  DateTime,"
            + "f2  Float32,"
            + "f3  Float64,"
            + "f4  Int8,"
            + "f5  Int16,"
            + "f6  Int32,"
            + "f7  Int64,"
            + "f8  String,"
            + "f9  UInt8,"
            + "f10 UInt16,"
            + "f11 UInt32,"
            + "f12 UInt64,"
            + "f13 Enum8('abc' = 1, 'cde' = 2),"
            + "f14 Enum16('abc' = -1, 'cde' = -2),"
            + "f15 FixedString(3),"
            + "f16 FixedString(3),"
            + "f17 FixedString(3)"
            + ") ENGINE=Log");

    pipeline.apply(Create.of(row1).withRowSchema(schema)).apply(write("test_primitive_types"));

    pipeline.run().waitUntilFinish();

    try (ResultSet rs = executeQuery("SELECT * FROM test_primitive_types")) {
      rs.next();

      assertEquals("2030-10-01", rs.getString("f0"));
      assertEquals("2030-10-09 08:07:06", rs.getString("f1"));
      assertEquals("2.2", rs.getString("f2"));
      assertEquals("3.3", rs.getString("f3"));
      assertEquals("4", rs.getString("f4"));
      assertEquals("5", rs.getString("f5"));
      assertEquals("6", rs.getString("f6"));
      assertEquals("7", rs.getString("f7"));
      assertEquals("eight", rs.getString("f8"));
      assertEquals("9", rs.getString("f9"));
      assertEquals("10", rs.getString("f10"));
      assertEquals("11", rs.getString("f11"));
      assertEquals("12", rs.getString("f12"));
      assertEquals("abc", rs.getString("f13"));
      assertEquals("cde", rs.getString("f14"));
      assertArrayEquals(new byte[] {'q', 'w', 'e'}, rs.getBytes("f15"));
      assertArrayEquals(new byte[] {'a', 's', 'd'}, rs.getBytes("f16"));
      assertArrayEquals(new byte[] {'z', 'x', 'c'}, rs.getBytes("f17"));
    }
  }

  @Test
  public void testArrayOfPrimitiveTypes() throws Exception {
    Schema schema =
        Schema.of(
            Schema.Field.of("f0", FieldType.array(FieldType.DATETIME)),
            Schema.Field.of("f1", FieldType.array(FieldType.DATETIME)),
            Schema.Field.of("f2", FieldType.array(FieldType.FLOAT)),
            Schema.Field.of("f3", FieldType.array(FieldType.DOUBLE)),
            Schema.Field.of("f4", FieldType.array(FieldType.BYTE)),
            Schema.Field.of("f5", FieldType.array(FieldType.INT16)),
            Schema.Field.of("f6", FieldType.array(FieldType.INT32)),
            Schema.Field.of("f7", FieldType.array(FieldType.INT64)),
            Schema.Field.of("f8", FieldType.array(FieldType.STRING)),
            Schema.Field.of("f9", FieldType.array(FieldType.INT16)),
            Schema.Field.of("f10", FieldType.array(FieldType.INT32)),
            Schema.Field.of("f11", FieldType.array(FieldType.INT64)),
            Schema.Field.of("f12", FieldType.array(FieldType.INT64)),
            Schema.Field.of("f13", FieldType.array(FieldType.STRING)),
            Schema.Field.of("f14", FieldType.array(FieldType.STRING)));
    Row row1 =
        Row.withSchema(schema)
            .addArray(
                new DateTime(2030, 10, 1, 0, 0, 0, DateTimeZone.UTC),
                new DateTime(2031, 10, 1, 0, 0, 0, DateTimeZone.UTC))
            .addArray(
                new DateTime(2030, 10, 9, 8, 7, 6, DateTimeZone.UTC),
                new DateTime(2031, 10, 9, 8, 7, 6, DateTimeZone.UTC))
            .addArray(2.2f, 3.3f)
            .addArray(3.3, 4.4)
            .addArray((byte) 4, (byte) 5)
            .addArray((short) 5, (short) 6)
            .addArray(6, 7)
            .addArray(7L, 8L)
            .addArray("eight", "nine")
            .addArray((short) 9, (short) 10)
            .addArray(10, 11)
            .addArray(11L, 12L)
            .addArray(12L, 13L)
            .addArray("abc", "cde")
            .addArray("cde", "abc")
            .build();

    executeSql(
        "CREATE TABLE test_array_of_primitive_types ("
            + "f0  Array(Date),"
            + "f1  Array(DateTime),"
            + "f2  Array(Float32),"
            + "f3  Array(Float64),"
            + "f4  Array(Int8),"
            + "f5  Array(Int16),"
            + "f6  Array(Int32),"
            + "f7  Array(Int64),"
            + "f8  Array(String),"
            + "f9  Array(UInt8),"
            + "f10 Array(UInt16),"
            + "f11 Array(UInt32),"
            + "f12 Array(UInt64),"
            + "f13 Array(Enum8('abc' = 1, 'cde' = 2)),"
            + "f14 Array(Enum16('abc' = -1, 'cde' = -2))"
            + ") ENGINE=Log");

    pipeline
        .apply(Create.of(row1).withRowSchema(schema))
        .apply(write("test_array_of_primitive_types"));

    pipeline.run().waitUntilFinish();

    try (ResultSet rs = executeQuery("SELECT * FROM test_array_of_primitive_types")) {
      rs.next();

      assertEquals("['2030-10-01','2031-10-01']", rs.getString("f0"));
      assertEquals("['2030-10-09 08:07:06','2031-10-09 08:07:06']", rs.getString("f1"));
      assertEquals("[2.2,3.3]", rs.getString("f2"));
      assertEquals("[3.3,4.4]", rs.getString("f3"));
      assertEquals("[4,5]", rs.getString("f4"));
      assertEquals("[5,6]", rs.getString("f5"));
      assertEquals("[6,7]", rs.getString("f6"));
      assertEquals("[7,8]", rs.getString("f7"));
      assertEquals("['eight','nine']", rs.getString("f8"));
      assertEquals("[9,10]", rs.getString("f9"));
      assertEquals("[10,11]", rs.getString("f10"));
      assertEquals("[11,12]", rs.getString("f11"));
      assertEquals("[12,13]", rs.getString("f12"));
      assertEquals("['abc','cde']", rs.getString("f13"));
      assertEquals("['cde','abc']", rs.getString("f14"));
    }
  }

  @Test
  public void testInsertSql() {
    TableSchema tableSchema =
        TableSchema.of(
            TableSchema.Column.of("f0", ColumnType.INT64),
            TableSchema.Column.of("f1", ColumnType.INT64));

    String expected = "INSERT INTO \"test_table\" (\"f0\", \"f1\")";

    assertEquals(expected, ClickHouseIO.WriteFn.insertSql(tableSchema, "test_table"));
  }

  /** POJO used to test . */
  @DefaultSchema(JavaFieldSchema.class)
  public static final class POJO {
    public int f0;
    public long f1;

    public POJO(int f0, long f1) {
      this.f0 = f0;
      this.f1 = f1;
    }

    public POJO() {}

    @Override
    public boolean equals(@Nullable Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      final POJO pojo = (POJO) o;
      return f0 == pojo.f0 && f1 == pojo.f1;
    }

    @Override
    public int hashCode() {
      return Objects.hash(f0, f1);
    }
  }

  @Ignore
  // FIXME java.lang.ClassNotFoundException: javax.annotation.Nullable
  public void testPojo() throws Exception {
    POJO pojo1 = new POJO(1, 2L);
    POJO pojo2 = new POJO(2, 4L);
    POJO pojo3 = new POJO(3, 6L);

    executeSql("CREATE TABLE test_pojo(f0 Int32, f1 Int64) ENGINE=Log");

    pipeline.apply(Create.of(pojo1, pojo2, pojo3)).apply(write("test_pojo"));

    pipeline.run().waitUntilFinish();

    long sum0 = executeQueryAsLong("SELECT SUM(f0) FROM test_pojo");
    long sum1 = executeQueryAsLong("SELECT SUM(f1) FROM test_pojo");

    assertEquals(6L, sum0);
    assertEquals(12L, sum1);
  }

  private <T> ClickHouseIO.Write<T> write(String table) {
    return ClickHouseIO.<T>write(clickHouse.getJdbcUrl(), table).withMaxRetries(0);
  }
}
