/*
 * 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.spark.translation;

import static org.junit.Assert.assertEquals;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.CombineWithContext;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.CombineFnUtil;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Test;

/** * Test suite for {@link SparkCombineFn}. */
public class SparkCombineFnTest {

  final SerializablePipelineOptions opts =
      new SerializablePipelineOptions(PipelineOptionsFactory.create());
  CombineWithContext.CombineFnWithContext<Integer, Long, Long> combineFn;

  @Before
  public void setUp() {
    combineFn =
        (CombineWithContext.CombineFnWithContext<Integer, Long, Long>)
            CombineFnUtil.toFnWithContext(getSumFn());
  }

  @Test
  public void testGlobalWindowCombineFn() throws Exception {
    SparkCombineFn<KV<String, Integer>, Integer, Long, Long> sparkCombineFn =
        SparkCombineFn.keyed(
            combineFn, opts, Collections.emptyMap(), WindowingStrategy.globalDefault());

    WindowedValue<KV<String, Integer>> first = input("key", 1, Instant.now());
    WindowedValue<KV<String, Integer>> second = input("key", 2, Instant.now());
    WindowedValue<KV<String, Integer>> third = input("key", 3, Instant.now());
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c1 =
        sparkCombineFn.createCombiner(first);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c2 =
        sparkCombineFn.createCombiner(third);
    sparkCombineFn.mergeValue(c1, second);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c3 =
        sparkCombineFn.mergeCombiners(c1, c2);
    assertEquals(6, (long) Iterables.getOnlyElement(sparkCombineFn.extractOutput(c3)).getValue());
  }

  @Test
  public void testGlobalCombineFn() throws Exception {
    SparkCombineFn<Integer, Integer, Long, Long> sparkCombineFn =
        SparkCombineFn.globally(
            combineFn, opts, Collections.emptyMap(), WindowingStrategy.globalDefault());

    WindowedValue<Integer> first = inputValue(1, Instant.now());
    WindowedValue<Integer> second = inputValue(2, Instant.now());
    WindowedValue<Integer> third = inputValue(3, Instant.now());
    SparkCombineFn.WindowedAccumulator<Integer, Integer, Long, ?> c1 =
        sparkCombineFn.createCombiner(first);
    SparkCombineFn.WindowedAccumulator<Integer, Integer, Long, ?> c2 =
        sparkCombineFn.createCombiner(third);
    sparkCombineFn.mergeValue(c1, second);
    SparkCombineFn.WindowedAccumulator<Integer, Integer, Long, ?> c3 =
        sparkCombineFn.mergeCombiners(c1, c2);
    assertEquals(6, (long) Iterables.getOnlyElement(sparkCombineFn.extractOutput(c3)).getValue());
  }

  @Test
  public void testSessionCombineFn() throws Exception {
    WindowingStrategy<Object, IntervalWindow> strategy =
        WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(1000)));

    SparkCombineFn<KV<String, Integer>, Integer, Long, Long> sparkCombineFn =
        SparkCombineFn.keyed(combineFn, opts, Collections.emptyMap(), strategy);

    Instant now = Instant.ofEpochMilli(0);
    WindowedValue<KV<String, Integer>> first =
        input("key", 1, now.plus(5000), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> second =
        input("key", 2, now.plus(1000), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> third =
        input("key", 3, now.plus(500), strategy.getWindowFn());
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c1 =
        sparkCombineFn.createCombiner(first);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c2 =
        sparkCombineFn.createCombiner(third);
    sparkCombineFn.mergeValue(c1, second);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c3 =
        sparkCombineFn.mergeCombiners(c1, c2);
    Iterable<WindowedValue<Long>> output = sparkCombineFn.extractOutput(c3);
    assertEquals(2, Iterables.size(output));
    List<String> format =
        StreamSupport.stream(output.spliterator(), false)
            .map(val -> val.getValue() + ":" + val.getTimestamp().getMillis())
            .collect(Collectors.toList());
    assertEquals(Lists.newArrayList("5:1999", "1:5999"), format);
  }

  @Test
  public void testSlidingCombineFnNonMerging() throws Exception {
    WindowingStrategy<Object, IntervalWindow> strategy =
        WindowingStrategy.of(SlidingWindows.of(Duration.millis(3000)).every(Duration.millis(1000)));

    SparkCombineFn<KV<String, Integer>, Integer, Long, Long> sparkCombineFn =
        SparkCombineFn.keyed(
            combineFn,
            opts,
            Collections.emptyMap(),
            strategy,
            SparkCombineFn.WindowedAccumulator.Type.NON_MERGING);

    Instant now = Instant.ofEpochMilli(0);
    WindowedValue<KV<String, Integer>> first =
        input("key", 1, now.plus(5000), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> second =
        input("key", 2, now.plus(1500), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> third =
        input("key", 3, now.plus(500), strategy.getWindowFn());
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c1 =
        sparkCombineFn.createCombiner(first);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c2 =
        sparkCombineFn.createCombiner(third);
    sparkCombineFn.mergeValue(c1, second);
    SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> c3 =
        sparkCombineFn.mergeCombiners(c1, c2);
    Iterable<WindowedValue<Long>> output = sparkCombineFn.extractOutput(c3);
    assertEquals(7, Iterables.size(output));
    List<String> format =
        StreamSupport.stream(output.spliterator(), false)
            .map(val -> val.getValue() + ":" + val.getTimestamp().getMillis())
            .collect(Collectors.toList());
    assertUnorderedEquals(
        Lists.newArrayList("3:999", "5:1999", "5:2999", "2:3999", "1:5999", "1:6999", "1:7999"),
        format);
  }

  @Test
  public void testSlidingCombineFnExplode() throws Exception {
    WindowingStrategy<Object, IntervalWindow> strategy =
        WindowingStrategy.of(SlidingWindows.of(Duration.millis(3000)).every(Duration.millis(1000)));

    SparkCombineFn<KV<String, Integer>, Integer, Long, Long> sparkCombineFn =
        SparkCombineFn.keyed(
            combineFn,
            opts,
            Collections.emptyMap(),
            strategy,
            SparkCombineFn.WindowedAccumulator.Type.EXPLODE_WINDOWS);

    Instant now = Instant.ofEpochMilli(0);
    WindowedValue<KV<String, Integer>> first =
        input("key", 1, now.plus(5000), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> second =
        input("key", 2, now.plus(1500), strategy.getWindowFn());
    WindowedValue<KV<String, Integer>> third =
        input("key", 3, now.plus(500), strategy.getWindowFn());

    Map<KV<String, BoundedWindow>, List<WindowedValue<KV<String, Integer>>>> groupByKeyAndWindow;
    groupByKeyAndWindow =
        Stream.of(first, second, third)
            .flatMap(e -> StreamSupport.stream(e.explodeWindows().spliterator(), false))
            .collect(
                Collectors.groupingBy(
                    e -> KV.of(e.getValue().getKey(), Iterables.getOnlyElement(e.getWindows()))));

    List<String> result = new ArrayList<>();
    for (Map.Entry<KV<String, BoundedWindow>, List<WindowedValue<KV<String, Integer>>>> e :
        groupByKeyAndWindow.entrySet()) {

      SparkCombineFn.WindowedAccumulator<KV<String, Integer>, Integer, Long, ?> combiner = null;
      for (WindowedValue<KV<String, Integer>> v : e.getValue()) {
        if (combiner == null) {
          combiner = sparkCombineFn.createCombiner(v);
        } else {
          combiner.add(v, sparkCombineFn);
        }
      }
      WindowedValue<Long> combined = Iterables.getOnlyElement(combiner.extractOutput());
      result.add(combined.getValue() + ":" + combined.getTimestamp().getMillis());
    }

    assertUnorderedEquals(
        Lists.newArrayList("3:999", "5:1999", "5:2999", "2:3999", "1:5999", "1:6999", "1:7999"),
        result);
  }

  private static Combine.CombineFn<Integer, Long, Long> getSumFn() {
    return new Combine.CombineFn<Integer, Long, Long>() {

      @Override
      public Long createAccumulator() {
        return 0L;
      }

      @Override
      public Long addInput(Long mutableAccumulator, Integer input) {
        return mutableAccumulator + input;
      }

      @Override
      public Long mergeAccumulators(Iterable<Long> accumulators) {
        return StreamSupport.stream(accumulators.spliterator(), false).mapToLong(e -> e).sum();
      }

      @Override
      public Long extractOutput(Long accumulator) {
        return accumulator;
      }
    };
  }

  private <K, V> WindowedValue<KV<K, V>> input(K key, V value, Instant timestamp) throws Exception {

    return input(key, value, timestamp, WindowingStrategy.globalDefault().getWindowFn());
  }

  private <K, V> WindowedValue<KV<K, V>> input(
      K key, V value, Instant timestamp, WindowFn<?, ?> windowFn) throws Exception {

    return inputValue(KV.of(key, value), timestamp, windowFn);
  }

  private <V> WindowedValue<V> inputValue(V value, Instant timestamp) throws Exception {
    return inputValue(value, timestamp, WindowingStrategy.globalDefault().getWindowFn());
  }

  private <V> WindowedValue<V> inputValue(V value, Instant timestamp, WindowFn<?, ?> windowFn)
      throws Exception {

    @SuppressWarnings("unchecked")
    WindowFn<V, BoundedWindow> cast = (WindowFn<V, BoundedWindow>) windowFn;
    return WindowedValue.of(
        value,
        timestamp,
        cast.assignWindows(assignContext(cast, value, timestamp)),
        PaneInfo.NO_FIRING);
  }

  <V> WindowFn<V, BoundedWindow>.AssignContext assignContext(
      WindowFn<V, BoundedWindow> windowFn, V value, Instant timestamp) {
    return windowFn.new AssignContext() {
      @Override
      public V element() {
        return value;
      }

      @Override
      public Instant timestamp() {
        return timestamp;
      }

      @Override
      public BoundedWindow window() {
        return GlobalWindow.INSTANCE;
      }
    };
  }

  private <T> void assertUnorderedEquals(List<T> actual, List<T> expected) {
    assertEquals(
        actual.stream().collect(Collectors.toSet()), expected.stream().collect(Collectors.toSet()));
  }
}