/*
 * 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.extensions.sql.zetasql;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.DateType;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimeType;
import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;

/** TestInput. */
class TestInput {
  public static final FieldType DATE = FieldType.logicalType(new DateType());
  public static final FieldType TIME = FieldType.logicalType(new TimeType());

  public static final TestBoundedTable BASIC_TABLE_ONE =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("Key")
                  .addStringField("Value")
                  .addDateTimeField("ts")
                  .build())
          .addRows(
              14L,
              "KeyValue234",
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:06"),
              15L,
              "KeyValue235",
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:07"));

  public static final TestBoundedTable BASIC_TABLE_TWO =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("RowKey")
                  .addStringField("Value")
                  .addDateTimeField("ts")
                  .build())
          .addRows(
              15L,
              "BigTable235",
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:07"),
              16L,
              "BigTable236",
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:08"));

  public static final TestBoundedTable BASIC_TABLE_THREE =
      TestBoundedTable.of(Schema.builder().addInt64Field("ColId").addStringField("Value").build())
          .addRows(15L, "Spanner235", 16L, "Spanner236", 17L, "Spanner237");

  public static final TestBoundedTable AGGREGATE_TABLE_ONE =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("Key")
                  .addInt64Field("Key2")
                  .addInt64Field("f_int_1")
                  .addStringField("f_str_1")
                  .addDoubleField("f_double_1")
                  .build())
          .addRows(1L, 10L, 1L, "1", 1.0)
          .addRows(1L, 11L, 2L, "2", 2.0)
          .addRows(2L, 11L, 3L, "3", 3.0)
          .addRows(2L, 11L, 4L, "4", 4.0)
          .addRows(2L, 12L, 5L, "5", 5.0)
          .addRows(3L, 13L, 6L, "6", 6.0)
          .addRows(3L, 13L, 7L, "7", 7.0);

  public static final TestBoundedTable AGGREGATE_TABLE_TWO =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("Key")
                  .addInt64Field("Key2")
                  .addInt64Field("f_int_1")
                  .addStringField("f_str_1")
                  .build())
          .addRows(1L, 10L, 1L, "1")
          .addRows(2L, 11L, 3L, "3")
          .addRows(2L, 11L, 4L, "4")
          .addRows(2L, 12L, 5L, "5")
          .addRows(2L, 13L, 6L, "6")
          .addRows(3L, 13L, 7L, "7");

  public static final TestBoundedTable TABLE_ALL_TYPES =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("row_id")
                  .addBooleanField("bool_col")
                  .addInt64Field("int64_col")
                  .addDoubleField("double_col")
                  .addStringField("str_col")
                  .addByteArrayField("bytes_col")
                  .build())
          .addRows(1L, true, -1L, 0.125d, "1", stringToBytes("1"))
          .addRows(2L, false, -2L, Math.pow(0.1, 324.0), "2", stringToBytes("2"))
          .addRows(3L, true, -3L, 0.375d, "3", stringToBytes("3"))
          .addRows(4L, false, -4L, 0.5d, "4", stringToBytes("4"))
          .addRows(5L, false, -5L, 0.5d, "5", stringToBytes("5"));

  public static final TestBoundedTable TABLE_ALL_TYPES_2 =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("row_id")
                  .addBooleanField("bool_col")
                  .addInt64Field("int64_col")
                  .addDoubleField("double_col")
                  .addStringField("str_col")
                  .addByteArrayField("bytes_col")
                  .build())
          .addRows(6L, true, -6L, 0.125d, "6", stringToBytes("6"))
          .addRows(7L, false, -7L, Math.pow(0.1, 324.0), "7", stringToBytes("7"))
          .addRows(8L, true, -8L, 0.375d, "8", stringToBytes("8"))
          .addRows(9L, false, -9L, 0.5d, "9", stringToBytes("9"))
          .addRows(10L, false, -10L, 0.5d, "10", stringToBytes("10"));

  public static final TestBoundedTable TIMESTAMP_TABLE_ONE =
      TestBoundedTable.of(Schema.builder().addDateTimeField("ts").addInt64Field("value").build())
          .addRows(
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:06"),
              3L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:07"),
              4L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:08"),
              6L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:09"),
              7L);

  public static final TestBoundedTable TIMESTAMP_TABLE_TWO =
      TestBoundedTable.of(Schema.builder().addDateTimeField("ts").addInt64Field("value").build())
          .addRows(
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:06"),
              3L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:07"),
              4L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:12"),
              6L,
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-01 21:26:13"),
              7L);

  public static final TestBoundedTable TIME_TABLE =
      TestBoundedTable.of(
              Schema.builder()
                  .addNullableField("f_date", DATE)
                  .addNullableField("f_time", TIME)
                  .addNullableField("f_timestamp", FieldType.DATETIME)
                  .addNullableField("f_timestamp_with_time_zone", FieldType.DATETIME)
                  .build())
          .addRows(
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-07-11 00:00:00"),
              DateTimeUtils.parseTimestampWithUTCTimeZone("1970-01-01 12:33:59.348"),
              DateTimeUtils.parseTimestampWithUTCTimeZone("2018-12-20 23:59:59.999"),
              DateTimeUtils.parseTimestampWithTimeZone("2018-12-10 10:38:59-1000"));

  public static final TestBoundedTable TABLE_ALL_NULL =
      TestBoundedTable.of(
              Schema.builder()
                  .addNullableField("primary_key", FieldType.INT64)
                  .addNullableField("bool_val", FieldType.BOOLEAN)
                  .addNullableField("double_val", FieldType.DOUBLE)
                  .addNullableField("int64_val", FieldType.INT64)
                  .addNullableField("str_val", FieldType.STRING)
                  .build())
          .addRows(1L, null, null, null, null);

  public static final Schema TABLE_WITH_STRUCT_ROW_SCHEMA =
      Schema.builder().addInt32Field("struct_col_int").addStringField("struct_col_str").build();

  public static final TestBoundedTable TABLE_WITH_STRUCT =
      TestBoundedTable.of(
              Schema.builder()
                  .addField("id", FieldType.INT64)
                  .addField("struct_col", FieldType.row(TABLE_WITH_STRUCT_ROW_SCHEMA))
                  .build())
          .addRows(
              1L,
              Row.withSchema(TABLE_WITH_STRUCT_ROW_SCHEMA).addValues(16, "row_one").build(),
              2L,
              Row.withSchema(TABLE_WITH_STRUCT_ROW_SCHEMA).addValues(17, "row_two").build());

  public static final TestBoundedTable TABLE_WITH_STRUCT_TIMESTAMP_STRING =
      TestBoundedTable.of(
              Schema.builder()
                  .addField("struct_col", FieldType.row(TABLE_WITH_STRUCT_ROW_SCHEMA))
                  .build())
          .addRows(
              Row.withSchema(TABLE_WITH_STRUCT_ROW_SCHEMA)
                  .addValues(3, "2019-01-15 13:21:03")
                  .build());

  private static final Schema structSchema =
      Schema.builder().addInt64Field("row_id").addStringField("data").build();
  private static final Schema structTableSchema =
      Schema.builder().addRowField("rowCol", structSchema).build();

  public static final TestBoundedTable TABLE_WITH_STRUCT_TWO =
      TestBoundedTable.of(structTableSchema)
          .addRows(Row.withSchema(structSchema).addValues(1L, "data1").build())
          .addRows(Row.withSchema(structSchema).addValues(2L, "data2").build())
          .addRows(Row.withSchema(structSchema).addValues(3L, "data2").build())
          .addRows(Row.withSchema(structSchema).addValues(3L, "data3").build());

  public static final TestBoundedTable TABLE_WITH_ARRAY =
      TestBoundedTable.of(Schema.builder().addArrayField("array_col", FieldType.STRING).build())
          .addRows(Arrays.asList("1", "2", "3"), ImmutableList.of());

  public static final TestBoundedTable TABLE_WITH_ARRAY_FOR_UNNEST =
      TestBoundedTable.of(
              Schema.builder()
                  .addInt64Field("int_col")
                  .addArrayField("int_array_col", FieldType.INT64)
                  .build())
          .addRows(14L, Arrays.asList(14L, 18L))
          .addRows(18L, Arrays.asList(22L, 24L));

  public static final TestBoundedTable TABLE_FOR_CASE_WHEN =
      TestBoundedTable.of(
              Schema.builder().addInt64Field("f_int").addStringField("f_string").build())
          .addRows(1L, "20181018");

  public static final TestBoundedTable TABLE_EMPTY =
      TestBoundedTable.of(Schema.builder().addInt64Field("ColId").addStringField("Value").build());

  private static final Schema TABLE_WTH_MAP_SCHEMA =
      Schema.builder()
          .addMapField("map_field", FieldType.STRING, FieldType.STRING)
          .addRowField("row_field", structSchema)
          .build();
  public static final TestBoundedTable TABLE_WITH_MAP =
      TestBoundedTable.of(TABLE_WTH_MAP_SCHEMA)
          .addRows(
              ImmutableMap.of("MAP_KEY_1", "MAP_VALUE_1"),
              Row.withSchema(structSchema).addValues(1L, "data1").build());

  public static byte[] stringToBytes(String s) {
    return s.getBytes(StandardCharsets.UTF_8);
  }
}