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

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Set;

import scala.Tuple2;

import com.google.common.collect.Sets;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.VoidFunction;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.util.ManualClock;
import org.junit.Assert;
import org.junit.Test;

import org.apache.spark.HashPartitioner;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.Optional;
import org.apache.spark.api.java.function.Function3;
import org.apache.spark.api.java.function.Function4;
import org.apache.spark.streaming.api.java.JavaPairDStream;
import org.apache.spark.streaming.api.java.JavaMapWithStateDStream;

public class JavaMapWithStateSuite extends LocalJavaStreamingContext implements Serializable {

  /**
   * This test is only for testing the APIs. It's not necessary to run it.
   */
  public void testAPI() {
    JavaPairRDD<String, Boolean> initialRDD = null;
    JavaPairDStream<String, Integer> wordsDstream = null;

    Function4<Time, String, Optional<Integer>, State<Boolean>, Optional<Double>> mappingFunc =
        new Function4<Time, String, Optional<Integer>, State<Boolean>, Optional<Double>>() {
          @Override
          public Optional<Double> call(
              Time time, String word, Optional<Integer> one, State<Boolean> state) {
            // Use all State's methods here
            state.exists();
            state.get();
            state.isTimingOut();
            state.remove();
            state.update(true);
            return Optional.of(2.0);
          }
        };

    JavaMapWithStateDStream<String, Integer, Boolean, Double> stateDstream =
        wordsDstream.mapWithState(
            StateSpec.function(mappingFunc)
                .initialState(initialRDD)
                .numPartitions(10)
                .partitioner(new HashPartitioner(10))
                .timeout(Durations.seconds(10)));

    stateDstream.stateSnapshots();

    Function3<String, Optional<Integer>, State<Boolean>, Double> mappingFunc2 =
        new Function3<String, Optional<Integer>, State<Boolean>, Double>() {
          @Override
          public Double call(String key, Optional<Integer> one, State<Boolean> state) {
            // Use all State's methods here
            state.exists();
            state.get();
            state.isTimingOut();
            state.remove();
            state.update(true);
            return 2.0;
          }
        };

    JavaMapWithStateDStream<String, Integer, Boolean, Double> stateDstream2 =
        wordsDstream.mapWithState(
            StateSpec.function(mappingFunc2)
                .initialState(initialRDD)
                .numPartitions(10)
                .partitioner(new HashPartitioner(10))
                .timeout(Durations.seconds(10)));

    stateDstream2.stateSnapshots();
  }

  @Test
  public void testBasicFunction() {
    List<List<String>> inputData = Arrays.asList(
        Collections.<String>emptyList(),
        Arrays.asList("a"),
        Arrays.asList("a", "b"),
        Arrays.asList("a", "b", "c"),
        Arrays.asList("a", "b"),
        Arrays.asList("a"),
        Collections.<String>emptyList()
    );

    List<Set<Integer>> outputData = Arrays.asList(
        Collections.<Integer>emptySet(),
        Sets.newHashSet(1),
        Sets.newHashSet(2, 1),
        Sets.newHashSet(3, 2, 1),
        Sets.newHashSet(4, 3),
        Sets.newHashSet(5),
        Collections.<Integer>emptySet()
    );

    @SuppressWarnings("unchecked")
    List<Set<Tuple2<String, Integer>>> stateData = Arrays.asList(
        Collections.<Tuple2<String, Integer>>emptySet(),
        Sets.newHashSet(new Tuple2<>("a", 1)),
        Sets.newHashSet(new Tuple2<>("a", 2), new Tuple2<>("b", 1)),
        Sets.newHashSet(new Tuple2<>("a", 3), new Tuple2<>("b", 2), new Tuple2<>("c", 1)),
        Sets.newHashSet(new Tuple2<>("a", 4), new Tuple2<>("b", 3), new Tuple2<>("c", 1)),
        Sets.newHashSet(new Tuple2<>("a", 5), new Tuple2<>("b", 3), new Tuple2<>("c", 1)),
        Sets.newHashSet(new Tuple2<>("a", 5), new Tuple2<>("b", 3), new Tuple2<>("c", 1))
    );

    Function3<String, Optional<Integer>, State<Integer>, Integer> mappingFunc =
        new Function3<String, Optional<Integer>, State<Integer>, Integer>() {
          @Override
          public Integer call(String key, Optional<Integer> value, State<Integer> state) {
            int sum = value.orElse(0) + (state.exists() ? state.get() : 0);
            state.update(sum);
            return sum;
          }
        };
    testOperation(
        inputData,
        StateSpec.function(mappingFunc),
        outputData,
        stateData);
  }

  private <K, S, T> void testOperation(
      List<List<K>> input,
      StateSpec<K, Integer, S, T> mapWithStateSpec,
      List<Set<T>> expectedOutputs,
      List<Set<Tuple2<K, S>>> expectedStateSnapshots) {
    int numBatches = expectedOutputs.size();
    JavaDStream<K> inputStream = JavaTestUtils.attachTestInputStream(ssc, input, 2);
    JavaMapWithStateDStream<K, Integer, S, T> mapWithStateDStream =
        JavaPairDStream.fromJavaDStream(inputStream.map(new Function<K, Tuple2<K, Integer>>() {
          @Override
          public Tuple2<K, Integer> call(K x) {
            return new Tuple2<>(x, 1);
          }
        })).mapWithState(mapWithStateSpec);

    final List<Set<T>> collectedOutputs =
        Collections.synchronizedList(new ArrayList<Set<T>>());
    mapWithStateDStream.foreachRDD(new VoidFunction<JavaRDD<T>>() {
      @Override
      public void call(JavaRDD<T> rdd) {
        collectedOutputs.add(Sets.newHashSet(rdd.collect()));
      }
    });
    final List<Set<Tuple2<K, S>>> collectedStateSnapshots =
        Collections.synchronizedList(new ArrayList<Set<Tuple2<K, S>>>());
    mapWithStateDStream.stateSnapshots().foreachRDD(new VoidFunction<JavaPairRDD<K, S>>() {
      @Override
      public void call(JavaPairRDD<K, S> rdd) {
        collectedStateSnapshots.add(Sets.newHashSet(rdd.collect()));
      }
    });
    BatchCounter batchCounter = new BatchCounter(ssc.ssc());
    ssc.start();
    ((ManualClock) ssc.ssc().scheduler().clock())
        .advance(ssc.ssc().progressListener().batchDuration() * numBatches + 1);
    batchCounter.waitUntilBatchesCompleted(numBatches, 10000);

    Assert.assertEquals(expectedOutputs, collectedOutputs);
    Assert.assertEquals(expectedStateSnapshots, collectedStateSnapshots);
  }
}
