/*
 * 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.complete.datatokenization.utils;

import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;

import com.google.api.services.bigquery.model.TableRow;
import com.google.auto.value.AutoValue;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition;
import org.apache.beam.sdk.io.gcp.bigquery.WriteResult;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.DateTimeZone;
import org.joda.time.Duration;
import org.joda.time.format.DateTimeFormat;
import org.joda.time.format.DateTimeFormatter;

/** Transforms & DoFns & Options for Teleport Error logging. */
public class ErrorConverters {

  /** Writes all Errors to GCS, place at the end of your pipeline. */
  @AutoValue
  public abstract static class WriteStringMessageErrorsAsCsv
      extends PTransform<PCollection<FailsafeElement<String, String>>, PDone> {

    public static Builder newBuilder() {
      return new AutoValue_ErrorConverters_WriteStringMessageErrorsAsCsv.Builder();
    }

    public abstract String errorWritePath();

    public abstract String csvDelimiter();

    public abstract @Nullable Duration windowDuration();

    @SuppressWarnings("argument.type.incompatible")
    @Override
    public PDone expand(PCollection<FailsafeElement<String, String>> pCollection) {

      PCollection<String> formattedErrorRows =
          pCollection.apply(
              "GetFormattedErrorRow", ParDo.of(new FailedStringToCsvRowFn(csvDelimiter())));

      if (pCollection.isBounded() == IsBounded.UNBOUNDED) {
        if (windowDuration() != null) {
          formattedErrorRows =
              formattedErrorRows.apply(Window.into(FixedWindows.of(windowDuration())));
        }
        return formattedErrorRows.apply(
            TextIO.write().to(errorWritePath()).withNumShards(1).withWindowedWrites());

      } else {
        return formattedErrorRows.apply(TextIO.write().to(errorWritePath()).withNumShards(1));
      }
    }

    /** Builder for {@link WriteStringMessageErrorsAsCsv}. */
    @AutoValue.Builder
    public abstract static class Builder {

      public abstract Builder setErrorWritePath(String errorWritePath);

      public abstract Builder setCsvDelimiter(String csvDelimiter);

      public abstract Builder setWindowDuration(@Nullable Duration duration);

      public abstract WriteStringMessageErrorsAsCsv build();
    }
  }

  /**
   * The {@link FailedStringToCsvRowFn} converts string objects which have failed processing into
   * {@link String} objects contained CSV which can be output to a filesystem.
   */
  public static class FailedStringToCsvRowFn extends DoFn<FailsafeElement<String, String>, String> {

    /**
     * The formatter used to convert timestamps into a BigQuery compatible <a
     * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#timestamp-type">format</a>.
     */
    private static final DateTimeFormatter TIMESTAMP_FORMATTER =
        DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");

    private final String csvDelimiter;

    public FailedStringToCsvRowFn(String csvDelimiter) {
      this.csvDelimiter = csvDelimiter;
    }

    public FailedStringToCsvRowFn() {
      this.csvDelimiter = ",";
    }

    @ProcessElement
    public void processElement(ProcessContext context) {
      FailsafeElement<String, String> failsafeElement = context.element();
      ArrayList<String> outputRow = new ArrayList<>();
      final String message = failsafeElement.getOriginalPayload();

      // Format the timestamp for insertion
      String timestamp =
          TIMESTAMP_FORMATTER.print(context.timestamp().toDateTime(DateTimeZone.UTC));

      outputRow.add(timestamp);
      outputRow.add(MoreObjects.firstNonNull(failsafeElement.getErrorMessage(), ""));

      // Only set the payload if it's populated on the message.
      if (message != null) {
        outputRow.add(message);
      }

      context.output(String.join(csvDelimiter, outputRow));
    }
  }

  /** Write errors as string encoded messages. */
  @AutoValue
  public abstract static class WriteStringMessageErrors
      extends PTransform<PCollection<FailsafeElement<String, String>>, WriteResult> {

    public static Builder newBuilder() {
      return new AutoValue_ErrorConverters_WriteStringMessageErrors.Builder();
    }

    public abstract String getErrorRecordsTable();

    public abstract String getErrorRecordsTableSchema();

    @Override
    public WriteResult expand(PCollection<FailsafeElement<String, String>> failedRecords) {

      return failedRecords
          .apply("FailedRecordToTableRow", ParDo.of(new FailedStringToTableRowFn()))
          .apply(
              "WriteFailedRecordsToBigQuery",
              BigQueryIO.writeTableRows()
                  .to(getErrorRecordsTable())
                  .withJsonSchema(getErrorRecordsTableSchema())
                  .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
                  .withWriteDisposition(WriteDisposition.WRITE_APPEND));
    }

    /** Builder for {@link WriteStringMessageErrors}. */
    @AutoValue.Builder
    public abstract static class Builder {

      public abstract Builder setErrorRecordsTable(String errorRecordsTable);

      public abstract Builder setErrorRecordsTableSchema(String errorRecordsTableSchema);

      public abstract WriteStringMessageErrors build();
    }
  }

  /**
   * The {@link FailedStringToTableRowFn} converts string objects which have failed processing into
   * {@link TableRow} objects which can be output to a dead-letter table.
   */
  public static class FailedStringToTableRowFn
      extends DoFn<FailsafeElement<String, String>, TableRow> {

    /**
     * The formatter used to convert timestamps into a BigQuery compatible <a
     * href="https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#timestamp-type">format</a>.
     */
    private static final DateTimeFormatter TIMESTAMP_FORMATTER =
        DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");

    @ProcessElement
    public void processElement(ProcessContext context) {
      FailsafeElement<String, String> failsafeElement = context.element();
      final String message = failsafeElement.getOriginalPayload();

      // Format the timestamp for insertion
      String timestamp =
          TIMESTAMP_FORMATTER.print(context.timestamp().toDateTime(DateTimeZone.UTC));

      // Build the table row
      @SuppressWarnings("nullness") // TableRow.set not annotated but does accept nulls
      final TableRow failedRow =
          new TableRow()
              .set("timestamp", timestamp)
              .set("errorMessage", failsafeElement.getErrorMessage())
              .set("stacktrace", failsafeElement.getStacktrace());

      // Only set the payload if it's populated on the message.
      if (message != null) {
        failedRow
            .set("payloadString", message)
            .set("payloadBytes", message.getBytes(StandardCharsets.UTF_8));
      }

      context.output(failedRow);
    }
  }

  /**
   * {@link WriteErrorsToTextIO} is a {@link PTransform} that writes strings error messages to file
   * system using TextIO and custom line format {@link SerializableFunction} to convert errors in
   * necessary format. <br>
   * Example of usage in pipeline:
   *
   * <pre>{@code
   * pCollection.apply("Write to TextIO",
   *   WriteErrorsToTextIO.<String,String>newBuilder()
   *     .setErrorWritePath("errors.txt")
   *     .setTranslateFunction((FailsafeElement<String,String> failsafeElement) -> {
   *       ArrayList<String> outputRow  = new ArrayList<>();
   *       final String message = failsafeElement.getOriginalPayload();
   *       String timestamp = Instant.now().toString();
   *       outputRow.add(timestamp);
   *       outputRow.add(failsafeElement.getErrorMessage());
   *       outputRow.add(failsafeElement.getStacktrace());
   *       // Only set the payload if it's populated on the message.
   *       if (failsafeElement.getOriginalPayload() != null) {
   *         outputRow.add(message);
   *       }
   *
   *       return String.join(",",outputRow);
   *     })
   * }</pre>
   */
  @AutoValue
  public abstract static class WriteErrorsToTextIO<T, V>
      extends PTransform<PCollection<FailsafeElement<T, V>>, PDone> {

    public static <T, V> WriteErrorsToTextIO.Builder<T, V> newBuilder() {
      return new AutoValue_ErrorConverters_WriteErrorsToTextIO.Builder<>();
    }

    public abstract String errorWritePath();

    public abstract SerializableFunction<FailsafeElement<T, V>, String> translateFunction();

    public abstract @Nullable Duration windowDuration();

    @Override
    @SuppressWarnings("argument.type.incompatible")
    public PDone expand(PCollection<FailsafeElement<T, V>> pCollection) {

      PCollection<String> formattedErrorRows =
          pCollection.apply(
              "GetFormattedErrorRow",
              MapElements.into(TypeDescriptors.strings()).via(translateFunction()));

      if (pCollection.isBounded() == PCollection.IsBounded.UNBOUNDED) {
        if (windowDuration() == null) {
          throw new RuntimeException("Unbounded input requires window interval to be set");
        }
        return formattedErrorRows
            .apply(Window.into(FixedWindows.of(windowDuration())))
            .apply(TextIO.write().to(errorWritePath()).withNumShards(1).withWindowedWrites());
      }

      return formattedErrorRows.apply(TextIO.write().to(errorWritePath()).withNumShards(1));
    }

    /** Builder for {@link WriteErrorsToTextIO}. */
    @AutoValue.Builder
    public abstract static class Builder<T, V> {

      public abstract WriteErrorsToTextIO.Builder<T, V> setErrorWritePath(String errorWritePath);

      public abstract WriteErrorsToTextIO.Builder<T, V> setTranslateFunction(
          SerializableFunction<FailsafeElement<T, V>, String> translateFunction);

      public abstract WriteErrorsToTextIO.Builder<T, V> setWindowDuration(
          @Nullable Duration duration);

      abstract SerializableFunction<FailsafeElement<T, V>, String> translateFunction();

      abstract WriteErrorsToTextIO<T, V> autoBuild();

      public WriteErrorsToTextIO<T, V> build() {
        checkNotNull(translateFunction(), "translateFunction is required.");
        return autoBuild();
      }
    }
  }
}
