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

import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Matchers.anyList;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.ReadyCheckingSideInputReader;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.construction.TransformInputs;
import org.apache.beam.runners.direct.ParDoMultiOverrideFactory.StatefulParDo;
import org.apache.beam.runners.direct.WatermarkManager.TimerUpdate;
import org.apache.beam.runners.local.StructuralKey;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor;
import org.mockito.Matchers;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;

/** Tests for {@link StatefulParDoEvaluatorFactory}. */
@RunWith(JUnit4.class)
public class StatefulParDoEvaluatorFactoryTest implements Serializable {
  @Mock private transient EvaluationContext mockEvaluationContext;
  @Mock private transient DirectExecutionContext mockExecutionContext;
  @Mock private transient DirectExecutionContext.DirectStepContext mockStepContext;
  @Mock private transient ReadyCheckingSideInputReader mockSideInputReader;
  @Mock private transient UncommittedBundle<Integer> mockUncommittedBundle;

  private static final String KEY = "any-key";
  private final transient PipelineOptions options = PipelineOptionsFactory.create();
  private final transient StateInternals stateInternals =
      CopyOnAccessInMemoryStateInternals.<Object>withUnderlying(KEY, null);

  private static final BundleFactory BUNDLE_FACTORY = ImmutableListBundleFactory.create();

  @Rule
  public transient TestPipeline pipeline =
      TestPipeline.create().enableAbandonedNodeEnforcement(false);

  @Before
  public void setup() {
    MockitoAnnotations.initMocks(this);
    when((StateInternals) mockStepContext.stateInternals()).thenReturn(stateInternals);
    when(mockEvaluationContext.createSideInputReader(anyList()))
        .thenReturn(
            SideInputContainer.create(mockEvaluationContext, Collections.emptyList())
                .createReaderForViews(Collections.emptyList()));
  }

  @Test
  public void windowCleanupScheduled() throws Exception {
    // To test the factory, first we set up a pipeline and then we use the constructed
    // pipeline to create the right parameters to pass to the factory

    final String stateId = "my-state-id";

    // For consistency, window it into FixedWindows. Actually we will fabricate an input bundle.
    PCollection<KV<String, Integer>> input =
        pipeline
            .apply(Create.of(KV.of("hello", 1), KV.of("hello", 2)))
            .apply(Window.into(FixedWindows.of(Duration.millis(10))));

    TupleTag<Integer> mainOutput = new TupleTag<>();
    PCollection<Integer> produced =
        input
            .apply(
                new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>(
                    new DoFn<KV<String, Integer>, Integer>() {
                      @StateId(stateId)
                      private final StateSpec<ValueState<String>> spec =
                          StateSpecs.value(StringUtf8Coder.of());

                      @ProcessElement
                      public void process(ProcessContext c) {}
                    },
                    mainOutput,
                    TupleTagList.empty(),
                    Collections.emptyList(),
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap()))
            .get(mainOutput)
            .setCoder(VarIntCoder.of());

    StatefulParDoEvaluatorFactory<String, Integer, Integer> factory =
        new StatefulParDoEvaluatorFactory<>(mockEvaluationContext, options);

    AppliedPTransform<
            PCollection<? extends KeyedWorkItem<String, KV<String, Integer>>>,
            PCollectionTuple,
            StatefulParDo<String, Integer, Integer>>
        producingTransform = (AppliedPTransform) DirectGraphs.getProducer(produced);

    // Then there will be a digging down to the step context to get the state internals
    when(mockEvaluationContext.getExecutionContext(
            eq(producingTransform), Mockito.<StructuralKey>any()))
        .thenReturn(mockExecutionContext);
    when(mockExecutionContext.getStepContext(any())).thenReturn(mockStepContext);

    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9));
    IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19));

    StateNamespace firstWindowNamespace =
        StateNamespaces.window(IntervalWindow.getCoder(), firstWindow);
    StateNamespace secondWindowNamespace =
        StateNamespaces.window(IntervalWindow.getCoder(), secondWindow);
    StateTag<ValueState<String>> tag =
        StateTags.tagForSpec(stateId, StateSpecs.value(StringUtf8Coder.of()));

    // Set up non-empty state. We don't mock + verify calls to clear() but instead
    // check that state is actually empty. We musn't care how it is accomplished.
    stateInternals.state(firstWindowNamespace, tag).write("first");
    stateInternals.state(secondWindowNamespace, tag).write("second");

    // A single bundle with some elements in the global window; it should register cleanup for the
    // global window state merely by having the evaluator created. The cleanup logic does not
    // depend on the window.
    CommittedBundle<KV<String, Integer>> inputBundle =
        BUNDLE_FACTORY
            .createBundle(input)
            .add(
                WindowedValue.of(
                    KV.of("hello", 1), new Instant(3), firstWindow, PaneInfo.NO_FIRING))
            .add(
                WindowedValue.of(
                    KV.of("hello", 2), new Instant(11), secondWindow, PaneInfo.NO_FIRING))
            .commit(Instant.now());

    // Merely creating the evaluator should suffice to register the cleanup callback
    factory.forApplication(producingTransform, inputBundle);

    ArgumentCaptor<Runnable> argumentCaptor = ArgumentCaptor.forClass(Runnable.class);
    verify(mockEvaluationContext)
        .scheduleAfterWindowExpiration(
            eq(producingTransform), eq(firstWindow), Mockito.any(), argumentCaptor.capture());

    // Should actually clear the state for the first window
    argumentCaptor.getValue().run();
    assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue());
    assertThat(stateInternals.state(secondWindowNamespace, tag).read(), equalTo("second"));

    verify(mockEvaluationContext)
        .scheduleAfterWindowExpiration(
            eq(producingTransform), eq(secondWindow), Mockito.any(), argumentCaptor.capture());

    // Should actually clear the state for the second window
    argumentCaptor.getValue().run();
    assertThat(stateInternals.state(secondWindowNamespace, tag).read(), nullValue());
  }

  /**
   * A test that explicitly delays a side input so that the main input will have to be reprocessed,
   * testing that {@code finishBundle()} re-assembles the GBK outputs correctly.
   */
  @Test
  public void testUnprocessedElements() throws Exception {
    // To test the factory, first we set up a pipeline and then we use the constructed
    // pipeline to create the right parameters to pass to the factory

    final String stateId = "my-state-id";

    // For consistency, window it into FixedWindows. Actually we will fabricate an input bundle.
    PCollection<KV<String, Integer>> mainInput =
        pipeline
            .apply(Create.of(KV.of("hello", 1), KV.of("hello", 2)))
            .apply(Window.into(FixedWindows.of(Duration.millis(10))));

    final PCollectionView<List<Integer>> sideInput =
        pipeline
            .apply("Create side input", Create.of(42))
            .apply("Window side input", Window.into(FixedWindows.of(Duration.millis(10))))
            .apply("View side input", View.asList());

    TupleTag<Integer> mainOutput = new TupleTag<>();
    PCollection<Integer> produced =
        mainInput
            .apply(
                new ParDoMultiOverrideFactory.GbkThenStatefulParDo<>(
                    new DoFn<KV<String, Integer>, Integer>() {
                      @StateId(stateId)
                      private final StateSpec<ValueState<String>> spec =
                          StateSpecs.value(StringUtf8Coder.of());

                      @ProcessElement
                      public void process(ProcessContext c) {}
                    },
                    mainOutput,
                    TupleTagList.empty(),
                    Collections.singletonList(sideInput),
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap()))
            .get(mainOutput)
            .setCoder(VarIntCoder.of());

    StatefulParDoEvaluatorFactory<String, Integer, Integer> factory =
        new StatefulParDoEvaluatorFactory<>(mockEvaluationContext, options);

    // This will be the stateful ParDo from the expansion
    AppliedPTransform<
            PCollection<KeyedWorkItem<String, KV<String, Integer>>>,
            PCollectionTuple,
            StatefulParDo<String, Integer, Integer>>
        producingTransform = (AppliedPTransform) DirectGraphs.getProducer(produced);

    // Then there will be a digging down to the step context to get the state internals
    when(mockEvaluationContext.getExecutionContext(
            eq(producingTransform), Mockito.<StructuralKey>any()))
        .thenReturn(mockExecutionContext);
    when(mockExecutionContext.getStepContext(any())).thenReturn(mockStepContext);
    when(mockEvaluationContext.createBundle(Matchers.<PCollection<Integer>>any()))
        .thenReturn(mockUncommittedBundle);
    when(mockStepContext.getTimerUpdate()).thenReturn(TimerUpdate.empty());

    // And digging to check whether the window is ready
    when(mockEvaluationContext.createSideInputReader(anyList())).thenReturn(mockSideInputReader);
    when(mockSideInputReader.isReady(Matchers.any(), Matchers.any())).thenReturn(false);

    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9));

    // A single bundle with some elements in the global window; it should register cleanup for the
    // global window state merely by having the evaluator created. The cleanup logic does not
    // depend on the window.
    String key = "hello";
    WindowedValue<KV<String, Integer>> firstKv =
        WindowedValue.of(KV.of(key, 1), new Instant(3), firstWindow, PaneInfo.NO_FIRING);

    WindowedValue<KeyedWorkItem<String, KV<String, Integer>>> gbkOutputElement =
        firstKv.withValue(
            KeyedWorkItems.elementsWorkItem(
                "hello",
                ImmutableList.of(
                    firstKv,
                    firstKv.withValue(KV.of(key, 13)),
                    firstKv.withValue(KV.of(key, 15)))));

    CommittedBundle<KeyedWorkItem<String, KV<String, Integer>>> inputBundle =
        BUNDLE_FACTORY
            .createBundle(
                (PCollection<KeyedWorkItem<String, KV<String, Integer>>>)
                    Iterables.getOnlyElement(
                        TransformInputs.nonAdditionalInputs(producingTransform)))
            .add(gbkOutputElement)
            .commit(Instant.now());
    TransformEvaluator<KeyedWorkItem<String, KV<String, Integer>>> evaluator =
        factory.forApplication(producingTransform, inputBundle);

    evaluator.processElement(gbkOutputElement);

    // This should push back every element as a KV<String, Iterable<Integer>>
    // in the appropriate window. Since the keys are equal they are single-threaded
    TransformResult<KeyedWorkItem<String, KV<String, Integer>>> result = evaluator.finishBundle();

    List<Integer> pushedBackInts = new ArrayList<>();

    for (WindowedValue<? extends KeyedWorkItem<String, KV<String, Integer>>> unprocessedElement :
        result.getUnprocessedElements()) {

      assertThat(
          Iterables.getOnlyElement(unprocessedElement.getWindows()),
          equalTo((BoundedWindow) firstWindow));

      assertThat(unprocessedElement.getValue().key(), equalTo("hello"));
      for (WindowedValue<KV<String, Integer>> windowedKv :
          unprocessedElement.getValue().elementsIterable()) {
        pushedBackInts.add(windowedKv.getValue().getValue());
      }
    }
    assertThat(pushedBackInts, containsInAnyOrder(1, 13, 15));
  }
}