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

import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.TypedRead;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFn.Element;
import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * An example that reads the public samples of weather data from BigQuery, counts the number of
 * tornadoes that occur in each month, and writes the results to BigQuery.
 *
 * <p>Concepts: Reading/writing BigQuery; counting a PCollection; user-defined PTransforms
 *
 * <p>Note: Before running this example, you must create a BigQuery dataset to contain your output
 * table.
 *
 * <p>To execute this pipeline locally, specify the BigQuery table for the output with the form:
 *
 * <pre>{@code
 * --output=YOUR_PROJECT_ID:DATASET_ID.TABLE_ID
 * }</pre>
 *
 * <p>To change the runner, specify:
 *
 * <pre>{@code
 * --runner=YOUR_SELECTED_RUNNER
 * }</pre>
 *
 * See examples/java/README.md for instructions about how to configure different runners.
 *
 * <p>The BigQuery input table defaults to {@code apache-beam-testing.samples.weather_stations} and
 * can be overridden with {@code --input}.
 */
public class BigQueryTornadoes {
  private static final Logger LOG = LoggerFactory.getLogger(BigQueryTornadoes.class);

  // Default to using a 1000 row subset of the public weather station table publicdata:samples.gsod.
  private static final String WEATHER_SAMPLES_TABLE =
      "apache-beam-testing.samples.weather_stations";

  /**
   * Examines each row in the input table. If a tornado was recorded in that sample, the month in
   * which it occurred is output.
   */
  static class ExtractTornadoesFn extends DoFn<TableRow, Integer> {
    @ProcessElement
    public void processElement(@Element TableRow row, OutputReceiver<Integer> receiver) {
      if ((Boolean) row.get("tornado")) {
        receiver.output(Integer.parseInt((String) row.get("month")));
      }
    }
  }

  /**
   * Prepares the data for writing to BigQuery by building a TableRow object containing an integer
   * representation of month and the number of tornadoes that occurred in each month.
   */
  static class FormatCountsFn extends DoFn<KV<Integer, Long>, TableRow> {
    @ProcessElement
    public void processElement(
        @Element KV<Integer, Long> element, OutputReceiver<TableRow> receiver) {
      TableRow row =
          new TableRow().set("month", element.getKey()).set("tornado_count", element.getValue());
      receiver.output(row);
    }
  }

  /**
   * Takes rows from a table and generates a table of counts.
   *
   * <p>The input schema is described by https://developers.google.com/bigquery/docs/dataset-gsod .
   * The output contains the total number of tornadoes found in each month in the following schema:
   *
   * <ul>
   *   <li>month: integer
   *   <li>tornado_count: integer
   * </ul>
   */
  static class CountTornadoes extends PTransform<PCollection<TableRow>, PCollection<TableRow>> {
    @Override
    public PCollection<TableRow> expand(PCollection<TableRow> rows) {

      // row... => month...
      PCollection<Integer> tornadoes = rows.apply(ParDo.of(new ExtractTornadoesFn()));

      // month... => <month,count>...
      PCollection<KV<Integer, Long>> tornadoCounts = tornadoes.apply(Count.perElement());

      // <month,count>... => row...
      PCollection<TableRow> results = tornadoCounts.apply(ParDo.of(new FormatCountsFn()));

      return results;
    }
  }

  /**
   * Options supported by {@link BigQueryTornadoes}.
   *
   * <p>Inherits standard configuration options.
   */
  public interface Options extends PipelineOptions {
    @Description("Table to read from, specified as <project_id>:<dataset_id>.<table_id>")
    @Default.String(WEATHER_SAMPLES_TABLE)
    String getInput();

    void setInput(String value);

    @Description("SQL Query to read from, will be used if Input is not set.")
    @Default.String("")
    String getInputQuery();

    void setInputQuery(String value);

    @Description("Read method to use to read from BigQuery")
    @Default.Enum("EXPORT")
    TypedRead.Method getReadMethod();

    void setReadMethod(TypedRead.Method value);

    @Description("Write method to use to write to BigQuery")
    @Default.Enum("DEFAULT")
    BigQueryIO.Write.Method getWriteMethod();

    void setWriteMethod(BigQueryIO.Write.Method value);

    @Description("Write disposition to use to write to BigQuery")
    @Default.Enum("WRITE_TRUNCATE")
    BigQueryIO.Write.WriteDisposition getWriteDisposition();

    void setWriteDisposition(BigQueryIO.Write.WriteDisposition value);

    @Description("Create disposition to use to write to BigQuery")
    @Default.Enum("CREATE_IF_NEEDED")
    BigQueryIO.Write.CreateDisposition getCreateDisposition();

    void setCreateDisposition(BigQueryIO.Write.CreateDisposition value);

    @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);
  }

  public static void applyBigQueryTornadoes(Pipeline p, Options options) {
    // Build the table schema for the output table.
    List<TableFieldSchema> fields = new ArrayList<>();
    fields.add(new TableFieldSchema().setName("month").setType("INTEGER"));
    fields.add(new TableFieldSchema().setName("tornado_count").setType("INTEGER"));
    TableSchema schema = new TableSchema().setFields(fields);

    TypedRead<TableRow> bigqueryIO;
    if (!options.getInputQuery().isEmpty()) {
      bigqueryIO =
          BigQueryIO.readTableRows()
              .fromQuery(options.getInputQuery())
              .usingStandardSql()
              .withMethod(options.getReadMethod());
    } else {
      bigqueryIO =
          BigQueryIO.readTableRows().from(options.getInput()).withMethod(options.getReadMethod());

      // Selected fields only applies when using Method.DIRECT_READ and
      // when reading directly from a table.
      if (options.getReadMethod() == TypedRead.Method.DIRECT_READ) {
        bigqueryIO = bigqueryIO.withSelectedFields(Lists.newArrayList("month", "tornado"));
      }
    }

    PCollection<TableRow> rowsFromBigQuery = p.apply(bigqueryIO);

    rowsFromBigQuery
        .apply(new CountTornadoes())
        .apply(
            BigQueryIO.writeTableRows()
                .to(options.getOutput())
                .withSchema(schema)
                .withCreateDisposition(options.getCreateDisposition())
                .withWriteDisposition(options.getWriteDisposition())
                .withMethod(options.getWriteMethod()));
  }

  public static void runBigQueryTornadoes(Options options) {
    LOG.info("Running BigQuery Tornadoes with options {}", options.toString());
    Pipeline p = Pipeline.create(options);
    applyBigQueryTornadoes(p, options);
    p.run().waitUntilFinish();
  }

  public static void main(String[] args) {
    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);

    runBigQueryTornadoes(options);
  }
}
