/*
 * 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.examples.snippets.transforms.io.gcp.bigquery;

import static org.junit.Assert.assertEquals;

import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQuery.DatasetDeleteOption;
import com.google.cloud.bigquery.BigQueryOptions;
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.FieldValue;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.QueryJobConfiguration;
import java.math.BigDecimal;
import java.security.SecureRandom;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.beam.examples.snippets.transforms.io.gcp.bigquery.BigQueryMyData.MyData;
import org.apache.beam.examples.snippets.transforms.io.gcp.bigquery.BigQueryMyData.MyStruct;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.FlatMapElements;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestName;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/**
 * Integration tests for BigQuery samples.
 *
 * <p>To run locally:
 *
 * <pre>{@code
 * ./gradlew integrationTest -p examples/java/ --info \
 *   --tests org.apache.beam.examples.snippets.transforms.io.gcp.bigquery.BigQuerySamplesIT \
 *   -DintegrationTestPipelineOptions='["--tempLocation=gs://YOUR-BUCKET/temp"]' \
 *   -DintegrationTestRunner=direct
 * }</pre>
 */
@RunWith(JUnit4.class)
public class BigQuerySamplesIT {
  private static final String PROJECT =
      TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
  private static final BigQuery BIGQUERY =
      BigQueryOptions.newBuilder().setProjectId(PROJECT).build().getService();
  private static final String DATASET =
      "beam_bigquery_samples_" + System.currentTimeMillis() + "_" + new SecureRandom().nextInt(32);

  @Rule public final transient TestPipeline writePipeline = TestPipeline.create();
  @Rule public final transient TestPipeline readTablePipeline = TestPipeline.create();
  @Rule public final transient TestPipeline readQueryPipeline = TestPipeline.create();
  @Rule public final transient TestPipeline readBQStorageAPIPipeline = TestPipeline.create();
  @Rule public final TestName testName = new TestName();

  @BeforeClass
  public static void beforeAll() throws Exception {
    BIGQUERY.create(DatasetInfo.newBuilder(PROJECT, DATASET).build());
  }

  @AfterClass
  public static void afterAll() {
    BIGQUERY.delete(DatasetId.of(PROJECT, DATASET), DatasetDeleteOption.deleteContents());
  }

  @Test
  public void testTableIO() throws Exception {
    String table = testName.getMethodName();

    // ===--- Test 1: createTableRow + writeToTable ---===\\
    // The rest of the tests depend on this since this is the one that writes
    // the contents into the BigQuery table, which the other tests then read.
    TableSchema schema = BigQuerySchemaCreate.createSchema();
    PCollection<TableRow> rows =
        writePipeline.apply(Create.of(Arrays.asList(BigQueryTableRowCreate.createTableRow())));
    BigQueryWriteToTable.writeToTable(PROJECT, DATASET, table, schema, rows);
    writePipeline.run().waitUntilFinish();

    // Check that the BigQuery table has the data using the BigQuery Client Library.
    String query = String.format("SELECT * FROM `%s.%s.%s`", PROJECT, DATASET, table);
    List<String> queryResults =
        StreamSupport.stream(
                BIGQUERY.query(QueryJobConfiguration.of(query)).iterateAll().spliterator(), false)
            .flatMap(values -> fieldValueListToStrings(values).stream())
            .collect(Collectors.toList());
    assertEquals(expected, queryResults);

    // ===--- Test 2: readFromTable ---=== \\
    readAndCheck(BigQueryReadFromTable.readFromTable(PROJECT, DATASET, table, readTablePipeline));
    readTablePipeline.run().waitUntilFinish();

    // ===--- Test 3: readFromQuery ---=== \\
    readAndCheck(BigQueryReadFromQuery.readFromQuery(PROJECT, DATASET, table, readQueryPipeline));
    readQueryPipeline.run().waitUntilFinish();

    // ===--- Test 4: readFromTableWithBigQueryStorageAPI ---=== \\
    readAndCheck(
        BigQueryReadFromTableWithBigQueryStorageAPI.readFromTableWithBigQueryStorageAPI(
            PROJECT, DATASET, table, readBQStorageAPIPipeline));
    readBQStorageAPIPipeline.run().waitUntilFinish();
  }

  // -- Helper methods -- \\
  private static void readAndCheck(PCollection<MyData> rows) {
    PCollection<String> contents =
        rows.apply(
            FlatMapElements.into(TypeDescriptors.strings())
                .via(BigQuerySamplesIT::myDataToStrings));
    PAssert.that(contents).containsInAnyOrder(expected);
  }

  private static List<String> expected =
      Arrays.asList(
          "string: UTF-8 strings are supported! 🌱🌳🌍",
          "int64: 432",
          "float64: 3.14159265",
          "numeric: 1234.56",
          "bool: true",
          "bytes: VVRGLTggYnl0ZSBzdHJpbmcg8J+MsfCfjLPwn4yN",
          "date: 2020-03-19",
          "datetime: 2020-03-19T20:41:25.123",
          "time: 20:41:25.123",
          "timestamp: 2020-03-20T03:41:42.123Z",
          "geography: POINT(30 10)",
          "array: [1, 2, 3, 4]",
          "struct: {string: Text 🌱🌳🌍, int64: 42}");

  private static List<String> myDataToStrings(MyData data) {
    return Arrays.asList(
        String.format("string: %s", data.myString),
        String.format("int64: %d", data.myInt64),
        String.format("float64: %.8f", data.myFloat64),
        String.format("numeric: %.2f", data.myNumeric.doubleValue()),
        String.format("bool: %s", data.myBoolean),
        String.format("bytes: %s", Base64.getEncoder().encodeToString(data.myBytes)),
        String.format("date: %s", data.myDate),
        String.format("datetime: %s", data.myDateTime),
        String.format("time: %s", data.myTime),
        String.format("timestamp: %s", data.myTimestamp),
        String.format("geography: %s", data.myGeography),
        String.format("array: %s", data.myArray),
        String.format(
            "struct: {string: %s, int64: %s}",
            data.myStruct.stringValue, data.myStruct.int64Value));
  }

  private static List<String> fieldValueListToStrings(FieldValueList row) {
    MyData data = new MyData();

    data.myString = row.get("string_field").getStringValue();
    data.myInt64 = row.get("int64_field").getLongValue();
    data.myFloat64 = row.get("float64_field").getDoubleValue();
    data.myNumeric = new BigDecimal(row.get("numeric_field").getDoubleValue());
    data.myBoolean = row.get("bool_field").getBooleanValue();
    data.myBytes = Base64.getDecoder().decode(row.get("bytes_field").getStringValue());

    data.myDate = LocalDate.parse(row.get("date_field").getStringValue()).toString();
    data.myDateTime = LocalDateTime.parse(row.get("datetime_field").getStringValue()).toString();
    data.myTime = LocalTime.parse(row.get("time_field").getStringValue()).toString();
    data.myTimestamp =
        Instant.ofEpochMilli(
                (long) (Double.parseDouble(row.get("timestamp_field").getStringValue()) * 1000.0))
            .toString();

    data.myGeography = row.get("geography_field").getStringValue();

    data.myArray =
        row.get("array_field").getRepeatedValue().stream()
            .map(FieldValue::getLongValue)
            .collect(Collectors.toList());

    FieldValueList structValues = row.get("struct_field").getRecordValue();
    data.myStruct = new MyStruct();
    data.myStruct.stringValue = structValues.get(0).getStringValue();
    data.myStruct.int64Value = structValues.get(1).getLongValue();

    return myDataToStrings(data);
  }
}
