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

import static org.hamcrest.MatcherAssert.assertThat;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.Set;
import javax.annotation.Nullable;
import org.apache.beam.runners.dataflow.TestDataflowPipelineOptions;
import org.apache.beam.runners.dataflow.TestDataflowRunner;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.extensions.ordered.StringBufferOrderedProcessingHandler.StringBufferOrderedProcessingWithGlobalSequenceHandler;
import org.apache.beam.sdk.extensions.ordered.UnprocessedEvent.Reason;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.SerializableMatcher;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestStream;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.Repeatedly;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PCollectionView;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.UnknownKeyFor;
import org.hamcrest.BaseMatcher;
import org.hamcrest.Description;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Rule;

/**
 * Ordered Processing tests use the same testing scenario. Events are sent in or out of sequence.
 * Each event is a string for a particular key. The output is a concatenation of all strings.
 */
public class OrderedEventProcessorTestBase {

  public static final boolean LAST_EVENT_RECEIVED = true;
  public static final int EMISSION_FREQUENCY_ON_EVERY_ELEMENT = 1;
  public static final int INITIAL_SEQUENCE_OF_0 = 0;
  public static final boolean DONT_PRODUCE_STATUS_ON_EVERY_EVENT = false;
  public static final int LARGE_MAX_RESULTS_PER_OUTPUT = 1000;
  public static final int EMISSION_FREQUENCY_ON_EVERY_OTHER_EVENT = 2;
  public static final boolean PRODUCE_STATUS_ON_EVERY_EVENT = true;
  public static final boolean STREAMING = true;
  public static final boolean BATCH = false;
  public static final Set<KV<String, KV<Long, UnprocessedEvent<String>>>> NO_EXPECTED_DLQ_EVENTS =
      Collections.emptySet();
  @Rule public final transient TestPipeline streamingPipeline = TestPipeline.create();
  @Rule public final transient TestPipeline batchPipeline = TestPipeline.create();

  protected boolean runTestsOnDataflowRunner() {
    return Boolean.getBoolean("run-tests-on-dataflow");
  }

  protected String getSystemProperty(String name) {
    String property = System.getProperty(name);
    if (property == null) {
      throw new IllegalStateException("Unable to find system property '" + name + "'");
    }
    return property;
  }

  static class MapEventsToKV extends DoFn<Event, KV<String, KV<Long, String>>> {

    @ProcessElement
    public void convert(
        @Element Event event, OutputReceiver<KV<String, KV<Long, String>>> outputReceiver) {
      outputReceiver.output(KV.of(event.getKey(), KV.of(event.getSequence(), event.getValue())));
    }
  }

  static class MapStringBufferStateToString
      extends DoFn<KV<String, StringBuilderState>, KV<String, String>> {

    @ProcessElement
    public void map(
        @Element KV<String, StringBuilderState> element,
        OutputReceiver<KV<String, String>> outputReceiver) {
      outputReceiver.output(KV.of(element.getKey(), element.getValue().toString()));
    }
  }

  /**
   * The majority of the tests use this method. Testing is done in the global window.
   *
   * @throws CannotProvideCoderException
   */
  protected void doTest(
      Event[] events,
      @Nullable Collection<KV<String, OrderedProcessingStatus>> expectedStatuses,
      Collection<KV<String, String>> expectedOutput,
      Collection<KV<String, KV<Long, UnprocessedEvent<String>>>> expectedUnprocessedEvents,
      int emissionFrequency,
      long initialSequence,
      int maxResultsPerOutput,
      boolean produceStatusOnEveryEvent,
      boolean streaming,
      boolean isGlobalSequence,
      @Nullable ContiguousSequenceRange expectedLastCompletedSequence)
      throws @UnknownKeyFor @NonNull @Initialized CannotProvideCoderException {

    Pipeline pipeline = streaming ? streamingPipeline : batchPipeline;
    if (runTestsOnDataflowRunner()) {
      pipeline.getOptions().setRunner(TestDataflowRunner.class);
      TestDataflowPipelineOptions options =
          pipeline.getOptions().as(TestDataflowPipelineOptions.class);
      options.setExperiments(Arrays.asList("disable_runner_v2"));
      options.setTempRoot("gs://" + getSystemProperty("temp_dataflow_bucket"));
    }
    PCollection<Event> rawInput =
        streaming
            ? createStreamingPCollection(pipeline, events)
            : createBatchPCollection(pipeline, events);
    PCollection<KV<String, KV<Long, String>>> input =
        rawInput.apply("To KV", ParDo.of(new MapEventsToKV()));

    OrderedProcessingHandler<String, String, StringBuilderState, String> handler =
        isGlobalSequence
            ? new StringBufferOrderedProcessingWithGlobalSequenceHandler(
                emissionFrequency, initialSequence)
            : new StringBufferOrderedProcessingHandler(emissionFrequency, initialSequence);
    handler.setMaxOutputElementsPerBundle(maxResultsPerOutput);
    if (produceStatusOnEveryEvent) {
      handler.setProduceStatusUpdateOnEveryEvent(true);
      // This disables status updates emitted on timers.
      handler.setStatusUpdateFrequency(null);
    } else {
      handler.setStatusUpdateFrequency(
          streaming ? Duration.standardMinutes(5) : Duration.standardSeconds(1));
    }

    OrderedEventProcessor<String, String, String, StringBuilderState> orderedEventProcessor =
        OrderedEventProcessor.create(handler);

    OrderedEventProcessorResult<String, String, String> processingResult =
        input.apply("Process Events", orderedEventProcessor);

    PAssert.that("Output matches", processingResult.output()).containsInAnyOrder(expectedOutput);

    if (streaming && expectedStatuses != null) {
      // Only in a streaming pipeline the events will arrive in a pre-determined order and the
      // statuses
      // will be deterministic. In batch pipelines events can be processed in any order,
      // so we skip status verification and rely on the output and unprocessed event matches.
      PAssert.that("Statuses match", processingResult.processingStatuses())
          .containsInAnyOrder(expectedStatuses);
    }

    // This is a temporary workaround until PAssert changes.
    boolean unprocessedEventsHaveExceptionStackTrace = false;
    for (KV<String, KV<Long, UnprocessedEvent<String>>> event : expectedUnprocessedEvents) {
      if (event.getValue().getValue().getReason() == Reason.exception_thrown) {
        unprocessedEventsHaveExceptionStackTrace = true;
        break;
      }
    }

    if (unprocessedEventsHaveExceptionStackTrace) {
      PAssert.thatSingleton(
              "Unprocessed event count",
              processingResult
                  .unprocessedEvents()
                  .apply(
                      "Window",
                      Window.<KV<String, KV<Long, UnprocessedEvent<String>>>>into(
                              new GlobalWindows())
                          .triggering(Repeatedly.forever(AfterWatermark.pastEndOfWindow()))
                          .discardingFiredPanes())
                  .apply("Count", Count.globally()))
          .isEqualTo((long) expectedUnprocessedEvents.size());
    } else {
      PAssert.that("Unprocessed events match", processingResult.unprocessedEvents())
          .containsInAnyOrder(expectedUnprocessedEvents);
    }

    if (expectedLastCompletedSequence != null && processingResult.latestContiguousRange() != null) {
      PCollection<ContiguousSequenceRange> globalSequences =
          rawInput.apply(
              "Publish Global Sequences",
              new GlobalSequenceRangePublisher(
                  processingResult.latestContiguousRange(),
                  handler.getKeyCoder(pipeline, input.getCoder()),
                  handler.getEventCoder(pipeline, input.getCoder())));
      PAssert.that("CompletedSequenceRange verification", globalSequences)
          .satisfies(new LastExpectedGlobalSequenceRangeMatcher(expectedLastCompletedSequence));
    }
    pipeline.run();
  }

  static class LastExpectedGlobalSequenceRangeMatcher
      implements SerializableFunction<Iterable<ContiguousSequenceRange>, Void> {

    private final long expectedStart;
    private final long expectedEnd;

    LastExpectedGlobalSequenceRangeMatcher(ContiguousSequenceRange expected) {
      this.expectedStart = expected.getStart();
      this.expectedEnd = expected.getEnd();
    }

    @Override
    public Void apply(Iterable<ContiguousSequenceRange> input) {
      StringBuilder listOfRanges = new StringBuilder("[");
      Iterator<ContiguousSequenceRange> iterator = input.iterator();
      ContiguousSequenceRange lastRange = null;
      while (iterator.hasNext()) {
        lastRange = iterator.next();

        if (listOfRanges.length() > 1) {
          listOfRanges.append(", ");
        }
        listOfRanges.append(lastRange);
      }
      listOfRanges.append(']');
      boolean foundExpectedRange =
          lastRange != null
              && lastRange.getStart() == expectedStart
              && lastRange.getEnd() == expectedEnd;

      assertThat(
          "Expected range not found: ["
              + expectedStart
              + '-'
              + expectedEnd
              + "], received ranges: "
              + listOfRanges,
          foundExpectedRange);
      return null;
    }
  }

  private @UnknownKeyFor @NonNull @Initialized PCollection<Event> createBatchPCollection(
      Pipeline pipeline, Event[] events) {
    return pipeline
        .apply("Create Batch Events", Create.of(Arrays.asList(events)))
        .apply("Reshuffle", Reshuffle.viaRandomKey());
  }

  private @UnknownKeyFor @NonNull @Initialized PCollection<Event> createStreamingPCollection(
      Pipeline pipeline, Event[] events)
      throws @UnknownKeyFor @NonNull @Initialized CannotProvideCoderException {
    Instant now = Instant.now().minus(Duration.standardMinutes(20));
    TestStream.Builder<Event> messageFlow =
        TestStream.create(pipeline.getCoderRegistry().getCoder(Event.class))
            .advanceWatermarkTo(now);

    int delayInMilliseconds = 0;
    for (Event e : events) {
      messageFlow =
          messageFlow
              .advanceWatermarkTo(now.plus(Duration.millis(++delayInMilliseconds)))
              .addElements(e);
    }

    // Needed to force the processing time based timers.
    messageFlow = messageFlow.advanceProcessingTime(Duration.standardMinutes(15));
    return pipeline.apply("Create Streaming Events", messageFlow.advanceWatermarkToInfinity());
  }

  /**
   * Unprocessed event's explanation contains stacktraces which makes tests very brittle because it
   * requires hardcoding the line numbers in the code. We use this matcher to only compare on the
   * first line of the explanation.
   */
  static class UnprocessedEventMatcher
      extends BaseMatcher<KV<String, KV<Long, UnprocessedEvent<String>>>>
      implements SerializableMatcher<KV<String, KV<Long, UnprocessedEvent<String>>>> {

    private KV<String, KV<Long, UnprocessedEvent<String>>> element;

    public UnprocessedEventMatcher(KV<String, KV<Long, UnprocessedEvent<String>>> element) {
      this.element = element;
    }

    @Override
    public boolean matches(Object actual) {
      KV<String, KV<Long, UnprocessedEvent<String>>> toMatch =
          (KV<String, KV<Long, UnprocessedEvent<String>>>) actual;

      UnprocessedEvent<String> originalEvent = element.getValue().getValue();
      UnprocessedEvent<String> eventToMatch = toMatch.getValue().getValue();

      return element.getKey().equals(toMatch.getKey())
          && element.getValue().getKey().equals(toMatch.getValue().getKey())
          && originalEvent.getEvent().equals(eventToMatch.getEvent())
          && originalEvent.getReason() == eventToMatch.getReason()
          && normalizeExplanation(originalEvent.getExplanation())
              .equals(normalizeExplanation(eventToMatch.getExplanation()));
    }

    @Override
    public void describeTo(Description description) {
      description.appendText("Just some text...");
    }

    static String normalizeExplanation(String value) {
      if (value == null) {
        return "";
      }
      String firstLine = value.split("\n", 1)[0];
      if (firstLine.contains("Exception")) {
        return firstLine;
      }
      return value;
    }
  }

  static class GlobalSequenceRangePublisher
      extends PTransform<PCollection<Event>, PCollection<ContiguousSequenceRange>> {

    private final PCollectionView<Iterable<ContiguousSequenceRange>> lastCompletedSequenceRangeView;
    private final Coder<String> keyCoder;
    private final Coder<String> eventCoder;

    public GlobalSequenceRangePublisher(
        PCollectionView<Iterable<ContiguousSequenceRange>> latestCompletedSequenceRange,
        Coder<String> keyCoder,
        Coder<String> eventCoder) {
      this.lastCompletedSequenceRangeView = latestCompletedSequenceRange;
      this.keyCoder = keyCoder;
      this.eventCoder = eventCoder;
    }

    @Override
    public PCollection<ContiguousSequenceRange> expand(PCollection<Event> input) {
      PCollection<KV<String, KV<Long, String>>> events =
          input
              // In production pipelines the global sequence will typically be obtained
              // by using GenerateSequence. But GenerateSequence doesn't work well with TestStream,
              // That's why we use the input events here.
              //              .apply("Create Ticker",
              //                  GenerateSequence.from(0).to(2).withRate(1,
              // Duration.standardSeconds(5)))
              .apply("To KV", ParDo.of(new MapEventsToKV()));
      if (input.isBounded() == IsBounded.BOUNDED) {
        return events.apply(
            "Emit SideInput",
            ParDo.of(new SideInputEmitter())
                .withSideInput("lastCompletedSequence", lastCompletedSequenceRangeView));
      } else {
        PCollection<KV<String, KV<Long, String>>> tickers =
            events.apply(
                "Create Tickers",
                new PerKeyTickerGenerator<>(keyCoder, eventCoder, Duration.standardSeconds(1)));
        return tickers.apply(
            "Emit SideInput",
            ParDo.of(new SideInputEmitter())
                .withSideInput("lastCompletedSequence", lastCompletedSequenceRangeView));
      }
    }

    static class SideInputEmitter
        extends DoFn<KV<String, KV<Long, String>>, ContiguousSequenceRange> {

      @ProcessElement
      public void produceCompletedRange(
          @SideInput("lastCompletedSequence") Iterable<ContiguousSequenceRange> sideInput,
          OutputReceiver<ContiguousSequenceRange> outputReceiver) {
        outputReceiver.output(ContiguousSequenceRange.largestRange(sideInput));
      }
    }
  }
}
