/*
 * Copyright (C) 2018 Google LLC
 *
 * Licensed 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 com.sonymusic;

import com.google.bigtable.v2.*;
import com.google.cloud.teleport.bigtable.BigtableCell;
import com.google.cloud.teleport.bigtable.BigtableRow;
import com.google.cloud.teleport.bigtable.PipelineUtils;
import com.google.cloud.teleport.util.DualInputNestedValueProvider;
import com.google.cloud.teleport.util.DualInputNestedValueProvider.TranslatorInput;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteOutput;
import com.google.protobuf.ByteString;
import com.google.protobuf.UnsafeByteOperations;
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.AvroIO;
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;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

/**
 * 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 BigtableToAvroRegexp {
  private static final Logger LOG = LoggerFactory.getLogger(BigtableToAvroRegexp.class);

  private static final Boolean DEFAULT_DELETE_ROWS = false;

  /** Options for the export pipeline. */
  public interface Options extends PipelineOptions {
    @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 output location to write to (e.g. gs://mybucket/somefolder/)")
    ValueProvider<String> getOutputDirectory();

    @SuppressWarnings("unused")
    void setOutputDirectory(ValueProvider<String> outputDirectory);

    @Description("The prefix for each exported file in outputDirectory")
    ValueProvider<String> getFilenamePrefix();

    @SuppressWarnings("unused")
    void setFilenamePrefix(ValueProvider<String> filenamePrefix);

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

    @SuppressWarnings("unused")
    void setRowKeyRegexp(ValueProvider<String> rowKeyRegexp);

    @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(PipelineUtils.tweakPipelineOptions(options));

    ValueProvider<RowFilter> rowFilterValueProvider = ValueProvider.NestedValueProvider.of(
      options.getRowKeyRegexp(),
      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();
    }

    ValueProvider<String> filePathPrefix =
        DualInputNestedValueProvider.of(
            options.getOutputDirectory(),
            options.getFilenamePrefix(),
            new SerializableFunction<TranslatorInput<String, String>, String>() {
              @Override
              public String apply(TranslatorInput<String, String> input) {
                return new StringBuilder(input.getX()).append(input.getY()).toString();
              }
            });

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

    btRows
        .apply("Transform to Avro", MapElements.via(new BigtableToAvroFn()))
        .apply(
            "Write to Avro in GCS",
            AvroIO.write(BigtableRow.class).to(filePathPrefix).withSuffix(".avro"));

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

    return pipeline.run();
  }

  /** Translates Bigtable {@link Row} to Avro {@link BigtableRow}. */
  static class BigtableToAvroFn extends SimpleFunction<Row, BigtableRow> {
    @Override
    public BigtableRow apply(Row row) {
      ByteBuffer key = ByteBuffer.wrap(toByteArray(row.getKey()));
      List<BigtableCell> cells = new ArrayList<>();
      for (Family family : row.getFamiliesList()) {
        String familyName = family.getName();
        for (Column column : family.getColumnsList()) {
          ByteBuffer qualifier = ByteBuffer.wrap(toByteArray(column.getQualifier()));
          for (Cell cell : column.getCellsList()) {
            long timestamp = cell.getTimestampMicros();
            ByteBuffer value = ByteBuffer.wrap(toByteArray(cell.getValue()));
            cells.add(new BigtableCell(familyName, qualifier, timestamp, value));
          }
        }
      }
      return new BigtableRow(key, cells);
    }
  }

  /**
   * Extracts the byte array from the given {@link ByteString} without copy.
   *
   * @param byteString A {@link ByteString} from which to extract the array.
   * @return an array of byte.
   */
  protected static byte[] toByteArray(final ByteString byteString) {
    try {
      ZeroCopyByteOutput byteOutput = new ZeroCopyByteOutput();
      UnsafeByteOperations.unsafeWriteTo(byteString, byteOutput);
      return byteOutput.bytes;
    } catch (IOException e) {
      return byteString.toByteArray();
    }
  }

  private static final class ZeroCopyByteOutput extends ByteOutput {
    private byte[] bytes;

    @Override
    public void writeLazy(byte[] value, int offset, int length) {
      if (offset != 0 || length != value.length) {
        throw new UnsupportedOperationException();
      }
      bytes = value;
    }

    @Override
    public void write(byte value) {
      throw new UnsupportedOperationException();
    }

    @Override
    public void write(byte[] value, int offset, int length) {
      throw new UnsupportedOperationException();
    }

    @Override
    public void write(ByteBuffer value) {
      throw new UnsupportedOperationException();
    }

    @Override
    public void writeLazy(ByteBuffer value) {
      throw new UnsupportedOperationException();
    }
  }

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

    public static BigtableToAvroRegexp.DeleteBigtableRowsFn create(
      ValueProvider<Boolean> deleteRowsFlag) {
      return new BigtableToAvroRegexp.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)));
      }
    }
  }
}