/*
 * 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.gcp.bigquery;

import static org.junit.Assert.assertEquals;

import com.google.api.client.util.BackOff;
import com.google.api.client.util.BackOffUtils;
import com.google.api.client.util.Sleeper;
import com.google.api.services.bigquery.model.QueryResponse;
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableCell;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableReference;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import java.security.SecureRandom;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.extensions.gcp.util.BackOffAdapter;
import org.apache.beam.sdk.io.gcp.testing.BigqueryClient;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.testing.DataflowPortabilityApiUnsupported;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap;
import org.joda.time.Duration;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Integration test for BigqueryIO with DataflowRunner and DirectRunner. */
@RunWith(JUnit4.class)
public class BigQueryToTableIT {
  private static final Logger LOG = LoggerFactory.getLogger(BigQueryToTableIT.class);
  private static String project;

  private static final BigqueryClient BQ_CLIENT = new BigqueryClient("BigQueryToTableIT");

  private static final String BIG_QUERY_DATASET_ID =
      "bq_query_to_table_" + System.currentTimeMillis() + "_" + (new SecureRandom().nextInt(32));

  private static final TableSchema LEGACY_QUERY_TABLE_SCHEMA =
      new TableSchema()
          .setFields(ImmutableList.of(new TableFieldSchema().setName("fruit").setType("STRING")));
  private static final TableSchema NEW_TYPES_QUERY_TABLE_SCHEMA =
      new TableSchema()
          .setFields(
              ImmutableList.of(
                  new TableFieldSchema().setName("bytes").setType("BYTES"),
                  new TableFieldSchema().setName("date").setType("DATE"),
                  new TableFieldSchema().setName("time").setType("TIME")));
  private static final String NEW_TYPES_QUERY_TABLE_NAME = "types";
  private static final List<Map<String, Object>> NEW_TYPES_QUERY_TABLE_DATA =
      ImmutableList.of(
          ImmutableMap.of("bytes", "abc=", "date", "2000-01-01", "time", "00:00:00"),
          ImmutableMap.of("bytes", "dec=", "date", "3000-12-31", "time", "23:59:59.990000"),
          ImmutableMap.of("bytes", "xyw=", "date", "2011-01-01", "time", "23:59:59.999999"));
  private static final int MAX_RETRY = 5;

  private void runBigQueryToTablePipeline(BigQueryToTableOptions options) {
    Pipeline p = Pipeline.create(options);
    BigQueryIO.Read bigQueryRead = BigQueryIO.read().fromQuery(options.getQuery());
    if (options.getUsingStandardSql()) {
      bigQueryRead = bigQueryRead.usingStandardSql();
    }
    PCollection<TableRow> input = p.apply(bigQueryRead);
    if (options.getReshuffle()) {
      input =
          input
              .apply(WithKeys.<Void, TableRow>of((Void) null))
              .setCoder(KvCoder.of(VoidCoder.of(), TableRowJsonCoder.of()))
              .apply(Reshuffle.<Void, TableRow>of())
              .apply(Values.<TableRow>create());
    }
    input.apply(
        BigQueryIO.writeTableRows()
            .to(options.getOutput())
            .withSchema(options.getOutputSchema())
            .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED));

    p.run().waitUntilFinish();
  }

  private BigQueryToTableOptions setupLegacyQueryTest(String outputTable) {
    BigQueryToTableOptions options =
        TestPipeline.testingPipelineOptions().as(BigQueryToTableOptions.class);
    options.setTempLocation(options.getTempRoot() + "/bq_it_temp");
    options.setQuery("SELECT * FROM (SELECT \"apple\" as fruit), (SELECT \"orange\" as fruit),");
    options.setOutput(outputTable);
    options.setOutputSchema(BigQueryToTableIT.LEGACY_QUERY_TABLE_SCHEMA);
    return options;
  }

  private BigQueryToTableOptions setupNewTypesQueryTest(String outputTable) {
    BigQueryToTableOptions options =
        TestPipeline.testingPipelineOptions().as(BigQueryToTableOptions.class);
    options.setTempLocation(options.getTempRoot() + "/bq_it_temp");
    options.setQuery(
        String.format(
            "SELECT bytes, date, time FROM [%s:%s.%s]",
            project, BIG_QUERY_DATASET_ID, BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_NAME));
    options.setOutput(outputTable);
    options.setOutputSchema(BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_SCHEMA);
    return options;
  }

  private BigQueryToTableOptions setupStandardQueryTest(String outputTable) {
    BigQueryToTableOptions options = this.setupLegacyQueryTest(outputTable);
    options.setTempLocation(options.getTempRoot() + "/bq_it_temp");
    options.setQuery(
        "SELECT * FROM (SELECT \"apple\" as fruit) UNION ALL (SELECT \"orange\" as fruit)");
    options.setUsingStandardSql(true);
    return options;
  }

  private List<TableRow> getTableRowsFromQuery(String query, int maxRetry) throws Exception {
    FluentBackoff backoffFactory =
        FluentBackoff.DEFAULT
            .withMaxRetries(maxRetry)
            .withInitialBackoff(Duration.standardSeconds(1L));
    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = BackOffAdapter.toGcpBackOff(backoffFactory.backoff());
    do {
      LOG.info("Starting querying {}", query);
      QueryResponse response = BQ_CLIENT.queryWithRetries(query, project);
      if (response.getRows() != null) {
        LOG.info("Got table content with query {}", query);
        return response.getRows();
      }
    } while (BackOffUtils.next(sleeper, backoff));
    LOG.info("Got empty table for query {} with retry {}", query, maxRetry);
    return Collections.emptyList();
  }

  private void verifyLegacyQueryRes(String outputTable) throws Exception {
    List<String> legacyQueryExpectedRes = ImmutableList.of("apple", "orange");
    List<TableRow> tableRows =
        getTableRowsFromQuery(String.format("SELECT fruit from [%s];", outputTable), MAX_RETRY);
    List<String> tableResult =
        tableRows.stream()
            .flatMap(row -> row.getF().stream().map(cell -> cell.getV().toString()))
            .sorted()
            .collect(Collectors.toList());
    assertEquals(legacyQueryExpectedRes, tableResult);
  }

  private void verifyNewTypesQueryRes(String outputTable) throws Exception {
    List<String> newTypeQueryExpectedRes =
        ImmutableList.of(
            "abc=,2000-01-01,00:00:00",
            "dec=,3000-12-31,23:59:59.990000",
            "xyw=,2011-01-01,23:59:59.999999");
    QueryResponse response =
        BQ_CLIENT.queryWithRetries(
            String.format("SELECT bytes, date, time FROM [%s];", outputTable), project);
    List<TableRow> tableRows =
        getTableRowsFromQuery(
            String.format("SELECT bytes, date, time FROM [%s];", outputTable), MAX_RETRY);
    List<String> tableResult =
        tableRows.stream()
            .map(
                row -> {
                  String res = "";
                  for (TableCell cell : row.getF()) {
                    if (res.isEmpty()) {
                      res = cell.getV().toString();
                    } else {
                      res = res + "," + cell.getV().toString();
                    }
                  }
                  return res;
                })
            .sorted()
            .collect(Collectors.toList());
    assertEquals(newTypeQueryExpectedRes, tableResult);
  }

  private void verifyStandardQueryRes(String outputTable) throws Exception {
    this.verifyLegacyQueryRes(outputTable);
  }

  /** Customized PipelineOption for BigQueryToTable Pipeline. */
  public interface BigQueryToTableOptions extends TestPipelineOptions, ExperimentalOptions {

    @Description("The BigQuery query to be used for creating the source")
    @Validation.Required
    String getQuery();

    void setQuery(String query);

    @Description(
        "BigQuery table to write to, specified as "
            + "<project_id>:<dataset_id>.<table_id>. The dataset must already exist.")
    @Validation.Required
    String getOutput();

    void setOutput(String value);

    @Description("BigQuery output table schema.")
    @Validation.Required
    TableSchema getOutputSchema();

    void setOutputSchema(TableSchema value);

    @Description("Whether to force reshuffle.")
    @Default.Boolean(false)
    boolean getReshuffle();

    void setReshuffle(boolean reshuffle);

    @Description("Whether to use the Standard SQL dialect when querying BigQuery.")
    @Default.Boolean(false)
    boolean getUsingStandardSql();

    void setUsingStandardSql(boolean usingStandardSql);
  }

  @BeforeClass
  public static void setupTestEnvironment() throws Exception {
    PipelineOptionsFactory.register(BigQueryToTableOptions.class);
    project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
    // Create one BQ dataset for all test cases.
    BQ_CLIENT.createNewDataset(project, BIG_QUERY_DATASET_ID);

    // Create table and insert data for new type query test cases.
    BQ_CLIENT.createNewTable(
        project,
        BIG_QUERY_DATASET_ID,
        new Table()
            .setSchema(BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_SCHEMA)
            .setTableReference(
                new TableReference()
                    .setTableId(BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_NAME)
                    .setDatasetId(BIG_QUERY_DATASET_ID)
                    .setProjectId(project)));
    BQ_CLIENT.insertDataToTable(
        project,
        BIG_QUERY_DATASET_ID,
        BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_NAME,
        BigQueryToTableIT.NEW_TYPES_QUERY_TABLE_DATA);
  }

  @AfterClass
  public static void cleanup() {
    LOG.info("Start to clean up tables and datasets.");
    BQ_CLIENT.deleteDataset(project, BIG_QUERY_DATASET_ID);
  }

  @Test
  public void testLegacyQueryWithoutReshuffle() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testLegacyQueryWithoutReshuffle";

    this.runBigQueryToTablePipeline(setupLegacyQueryTest(outputTable));

    this.verifyLegacyQueryRes(outputTable);
  }

  @Test
  public void testNewTypesQueryWithoutReshuffle() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testNewTypesQueryWithoutReshuffle";

    this.runBigQueryToTablePipeline(setupNewTypesQueryTest(outputTable));

    this.verifyNewTypesQueryRes(outputTable);
  }

  @Test
  public void testNewTypesQueryWithReshuffle() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testNewTypesQueryWithReshuffle";
    BigQueryToTableOptions options = setupNewTypesQueryTest(outputTable);
    options.setReshuffle(true);

    this.runBigQueryToTablePipeline(options);

    this.verifyNewTypesQueryRes(outputTable);
  }

  @Test
  public void testStandardQueryWithoutCustom() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testStandardQueryWithoutCustom";

    this.runBigQueryToTablePipeline(setupStandardQueryTest(outputTable));

    this.verifyStandardQueryRes(outputTable);
  }

  @Test
  @Category(DataflowPortabilityApiUnsupported.class)
  public void testNewTypesQueryWithoutReshuffleWithCustom() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testNewTypesQueryWithoutReshuffleWithCustom";
    BigQueryToTableOptions options = this.setupNewTypesQueryTest(outputTable);
    options.setExperiments(
        ImmutableList.of("enable_custom_bigquery_sink", "enable_custom_bigquery_source"));

    this.runBigQueryToTablePipeline(options);

    this.verifyNewTypesQueryRes(outputTable);
  }

  @Test
  @Category(DataflowPortabilityApiUnsupported.class)
  public void testLegacyQueryWithoutReshuffleWithCustom() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testLegacyQueryWithoutReshuffleWithCustom";
    BigQueryToTableOptions options = this.setupLegacyQueryTest(outputTable);
    options.setExperiments(
        ImmutableList.of("enable_custom_bigquery_sink", "enable_custom_bigquery_source"));

    this.runBigQueryToTablePipeline(options);

    this.verifyLegacyQueryRes(outputTable);
  }

  @Test
  @Category(DataflowPortabilityApiUnsupported.class)
  public void testStandardQueryWithoutReshuffleWithCustom() throws Exception {
    final String outputTable =
        project + ":" + BIG_QUERY_DATASET_ID + "." + "testStandardQueryWithoutReshuffleWithCustom";
    BigQueryToTableOptions options = this.setupStandardQueryTest(outputTable);
    options.setExperiments(
        ImmutableList.of("enable_custom_bigquery_sink", "enable_custom_bigquery_source"));

    this.runBigQueryToTablePipeline(options);

    this.verifyStandardQueryRes(outputTable);
  }
}