package com.sonymusic;

import com.google.bigtable.v2.*;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO;
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.ValueProvider;
import org.apache.beam.sdk.transforms.*;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Dataflow pipeline that exports data from a Cloud Bigtable table to Avro files in GCS. Currently,
 * filtering on Cloud Bigtable table is not supported.
 */
public class BigtableRemoveByRegex {
  private static final Logger LOG = LoggerFactory.getLogger(BigtableRemoveByRegex.class);

  private static final Boolean DEFAULT_DELETE_ROWS = false;

  /** Options for the export pipeline. */
  public interface Options extends DataflowPipelineOptions {
    @Description("The project that contains the table to export.")
    ValueProvider<String> getBigtableProjectId();

    @SuppressWarnings("unused")
    void setBigtableProjectId(ValueProvider<String> projectId);

    @Description("The Bigtable instance id that contains the table to export.")
    ValueProvider<String> getBigtableInstanceId();

    @SuppressWarnings("unused")
    void setBigtableInstanceId(ValueProvider<String> instanceId);

    @Description("The Bigtable table id to export.")
    ValueProvider<String> getBigtableTableId();

    @SuppressWarnings("unused")
    void setBigtableTableId(ValueProvider<String> tableId);

    @Description("The regexp value to filter rowkeys by")
    ValueProvider<String> getRowKeyRegex();

    @SuppressWarnings("unused")
    void setRowKeyRegex(ValueProvider<String> rowKeyRegex);

    @Description("Determines if the copied values should be removed from BigTable")
    ValueProvider<Boolean> getDeleteRows();

    @SuppressWarnings("unused")
    void setDeleteRows(ValueProvider<Boolean> deleteRows);
  }

  /**
   * Runs a pipeline to export data from a Cloud Bigtable table to Avro files in GCS.
   *
   * @param args arguments to the pipeline
   */
  public static void main(String[] args) {
    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);

    PipelineResult result = run(options);

    // Wait for pipeline to finish only if it is not constructing a template.
    if (options.as(DataflowPipelineOptions.class).getTemplateLocation() == null) {
      result.waitUntilFinish();
    }
  }

  public static PipelineResult run(Options options) {
    Pipeline pipeline = Pipeline.create(options);

    ValueProvider<RowFilter> rowFilterValueProvider = ValueProvider.NestedValueProvider.of(
      options.getRowKeyRegex(),
      new SerializableFunction<String, RowFilter>() {
        @Override
        public RowFilter apply(String input) {
          return RowFilter.newBuilder()
                  .setRowKeyRegexFilter(ByteString.copyFromUtf8(input))
                  .build();
        }});

    BigtableIO.Read read =
        BigtableIO.read()
            .withProjectId(options.getBigtableProjectId())
            .withInstanceId(options.getBigtableInstanceId())
            .withTableId(options.getBigtableTableId())
            .withRowFilter(rowFilterValueProvider);

    BigtableIO.Write write =
      BigtableIO.write()
        .withProjectId(options.getBigtableProjectId())
        .withInstanceId(options.getBigtableInstanceId())
        .withTableId(options.getBigtableTableId());

    // Do not validate input fields if it is running as a template.
    if (options.as(DataflowPipelineOptions.class).getTemplateLocation() != null) {
      read = read.withoutValidation();
    }

    PCollection<Row> btRows = pipeline
        .apply("Read from Bigtable", read);

    btRows
      .apply("Transform to Bigtable mutations.",
        ParDo.of(DeleteBigtableRowsFn.create(options.getDeleteRows())))
      .apply("Delete from Bigtable", write);

    return pipeline.run();
  }

  static class DeleteBigtableRowsFn extends DoFn<Row, KV<ByteString, Iterable<Mutation>>> {
    private final ValueProvider<Boolean> deleteRowsFlag;
    private Boolean deleteRows;

    public static BigtableRemoveByRegex.DeleteBigtableRowsFn create(
      ValueProvider<Boolean> deleteRowsFlag) {
      return new BigtableRemoveByRegex.DeleteBigtableRowsFn(deleteRowsFlag);
    }

    private DeleteBigtableRowsFn(
      ValueProvider<Boolean> deleteRowsFlag) {
      this.deleteRowsFlag = deleteRowsFlag;
    }

    @Setup
    public void setup() {
      if (deleteRowsFlag != null) {
        deleteRows = deleteRowsFlag.get();
      }
      deleteRows = MoreObjects.firstNonNull(deleteRows, DEFAULT_DELETE_ROWS);
      LOG.info("DeleteRows set to: " + deleteRows);
    }

    @ProcessElement
    public void processElement(
      @Element Row row, OutputReceiver<KV<ByteString, Iterable<Mutation>>> out) {
      if (deleteRows) {
        Mutation deleteRowMutation = Mutation.newBuilder().setDeleteFromRow(
          Mutation.DeleteFromRow.getDefaultInstance()).build();

        out.output(KV.of(row.getKey(), ImmutableList.of(deleteRowMutation)));
      }
    }
  }
}