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

import static org.hamcrest.Matchers.contains;
import static org.junit.Assert.assertThat;

import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.construction.Environments;
import org.apache.beam.runners.core.construction.RehydratedComponents;
import org.apache.beam.runners.core.construction.SdkComponents;
import org.apache.beam.runners.core.construction.graph.PipelineNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.runners.fnexecution.wire.LengthPrefixUnknownCoders;
import org.apache.beam.runners.local.StructuralKey;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.util.CoderUtils;
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.v20_0.com.google.common.collect.HashMultiset;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Multiset;
import org.hamcrest.BaseMatcher;
import org.hamcrest.Description;
import org.joda.time.Instant;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link GroupByKeyOnlyEvaluatorFactory}. */
@RunWith(JUnit4.class)
public class GroupByKeyOnlyEvaluatorFactoryTest {
  private BundleFactory bundleFactory = ImmutableListBundleFactory.create();

  @Test
  public void testInMemoryEvaluator() throws Exception {
    KvCoder<String, Integer> javaCoder = KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of());
    SdkComponents sdkComponents = SdkComponents.create();
    sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
    String windowingStrategyId =
        sdkComponents.registerWindowingStrategy(WindowingStrategy.globalDefault());
    String coderId = sdkComponents.registerCoder(javaCoder);
    Components.Builder builder = sdkComponents.toComponents().toBuilder();
    String javaWireCoderId =
        LengthPrefixUnknownCoders.addLengthPrefixedCoder(coderId, builder, false);
    String runnerWireCoderId =
        LengthPrefixUnknownCoders.addLengthPrefixedCoder(coderId, builder, true);
    RehydratedComponents components = RehydratedComponents.forComponents(builder.build());
    Coder<KV<String, Integer>> javaWireCoder =
        (Coder<KV<String, Integer>>) components.getCoder(javaWireCoderId);
    Coder<KV<?, ?>> runnerWireCoder = (Coder<KV<?, ?>>) components.getCoder(runnerWireCoderId);

    KV<?, ?> firstFoo = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("foo", -1));
    KV<?, ?> secondFoo = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("foo", 1));
    KV<?, ?> thirdFoo = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("foo", 3));
    KV<?, ?> firstBar = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("bar", 22));
    KV<?, ?> secondBar = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("bar", 12));
    KV<?, ?> firstBaz = asRunnerKV(javaWireCoder, runnerWireCoder, KV.of("baz", Integer.MAX_VALUE));

    PTransformNode inputTransform =
        PipelineNode.pTransform(
            "source", PTransform.newBuilder().putOutputs("out", "values").build());
    PCollectionNode values =
        PipelineNode.pCollection(
            "values",
            RunnerApi.PCollection.newBuilder()
                .setUniqueName("values")
                .setCoderId(coderId)
                .setWindowingStrategyId(windowingStrategyId)
                .build());
    PCollectionNode groupedKvs =
        PipelineNode.pCollection(
            "groupedKvs", RunnerApi.PCollection.newBuilder().setUniqueName("groupedKvs").build());
    PTransformNode groupByKeyOnly =
        PipelineNode.pTransform(
            "gbko",
            PTransform.newBuilder()
                .putInputs("input", "values")
                .putOutputs("output", "groupedKvs")
                .setSpec(FunctionSpec.newBuilder().setUrn(DirectGroupByKey.DIRECT_GBKO_URN).build())
                .build());
    Pipeline pipeline =
        Pipeline.newBuilder()
            .addRootTransformIds(inputTransform.getId())
            .addRootTransformIds(groupByKeyOnly.getId())
            .setComponents(
                builder
                    .putTransforms(inputTransform.getId(), inputTransform.getTransform())
                    .putTransforms(groupByKeyOnly.getId(), groupByKeyOnly.getTransform())
                    .putPcollections(values.getId(), values.getPCollection())
                    .putPcollections(groupedKvs.getId(), groupedKvs.getPCollection()))
            .build();

    PortableGraph graph = PortableGraph.forPipeline(pipeline);

    CommittedBundle<KV<String, Integer>> inputBundle =
        bundleFactory.<KV<String, Integer>>createBundle(values).commit(Instant.now());

    TransformEvaluator<KV<?, ?>> evaluator =
        new GroupByKeyOnlyEvaluatorFactory(graph, pipeline.getComponents(), bundleFactory)
            .forApplication(groupByKeyOnly, inputBundle);

    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(secondFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(thirdFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBar));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(secondBar));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBaz));

    TransformResult<KV<?, ?>> result = evaluator.finishBundle();

    // The input to a GroupByKey is assumed to be a KvCoder
    @SuppressWarnings({"rawtypes", "unchecked"})
    Coder runnerKeyCoder = ((KvCoder) runnerWireCoder).getKeyCoder();
    CommittedBundle<?> fooBundle = null;
    CommittedBundle<?> barBundle = null;
    CommittedBundle<?> bazBundle = null;
    StructuralKey fooKey = StructuralKey.of(firstFoo.getKey(), runnerKeyCoder);
    StructuralKey barKey = StructuralKey.of(firstBar.getKey(), runnerKeyCoder);
    StructuralKey bazKey = StructuralKey.of(firstBaz.getKey(), runnerKeyCoder);
    for (UncommittedBundle<?> groupedBundle : result.getOutputBundles()) {
      CommittedBundle<?> groupedCommitted = groupedBundle.commit(Instant.now());
      if (fooKey.equals(groupedCommitted.getKey())) {
        fooBundle = groupedCommitted;
      } else if (barKey.equals(groupedCommitted.getKey())) {
        barBundle = groupedCommitted;
      } else if (bazKey.equals(groupedCommitted.getKey())) {
        bazBundle = groupedCommitted;
      } else {
        throw new IllegalArgumentException(
            String.format("Unknown Key %s", groupedCommitted.getKey()));
      }
    }
    assertThat(
        fooBundle,
        contains(
            new KeyedWorkItemMatcher(
                KeyedWorkItems.elementsWorkItem(
                    fooKey.getKey(),
                    ImmutableSet.of(
                        WindowedValue.valueInGlobalWindow(firstFoo.getValue()),
                        WindowedValue.valueInGlobalWindow(secondFoo.getValue()),
                        WindowedValue.valueInGlobalWindow(thirdFoo.getValue()))),
                runnerKeyCoder)));
    assertThat(
        barBundle,
        contains(
            new KeyedWorkItemMatcher<>(
                KeyedWorkItems.elementsWorkItem(
                    barKey.getKey(),
                    ImmutableSet.of(
                        WindowedValue.valueInGlobalWindow(firstBar.getValue()),
                        WindowedValue.valueInGlobalWindow(secondBar.getValue()))),
                runnerKeyCoder)));
    assertThat(
        bazBundle,
        contains(
            new KeyedWorkItemMatcher<>(
                KeyedWorkItems.elementsWorkItem(
                    bazKey.getKey(),
                    ImmutableSet.of(WindowedValue.valueInGlobalWindow(firstBaz.getValue()))),
                runnerKeyCoder)));
  }

  private KV<?, ?> asRunnerKV(
      Coder<KV<String, Integer>> javaWireCoder,
      Coder<KV<?, ?>> runnerWireCoder,
      KV<String, Integer> value)
      throws org.apache.beam.sdk.coders.CoderException {
    return CoderUtils.decodeFromByteArray(
        runnerWireCoder, CoderUtils.encodeToByteArray(javaWireCoder, value));
  }

  private <K, V> KV<K, WindowedValue<V>> gwValue(KV<K, V> kv) {
    return KV.of(kv.getKey(), WindowedValue.valueInGlobalWindow(kv.getValue()));
  }

  private static class KeyedWorkItemMatcher<K, V>
      extends BaseMatcher<WindowedValue<KeyedWorkItem<K, V>>> {
    private final KeyedWorkItem<K, V> myWorkItem;
    private final Coder<K> keyCoder;

    public KeyedWorkItemMatcher(KeyedWorkItem<K, V> myWorkItem, Coder<K> keyCoder) {
      this.myWorkItem = myWorkItem;
      this.keyCoder = keyCoder;
    }

    @Override
    public boolean matches(Object item) {
      if (item == null || !(item instanceof WindowedValue)) {
        return false;
      }
      WindowedValue<KeyedWorkItem<K, V>> that = (WindowedValue<KeyedWorkItem<K, V>>) item;
      Multiset<WindowedValue<V>> myValues = HashMultiset.create();
      Multiset<WindowedValue<V>> thatValues = HashMultiset.create();
      for (WindowedValue<V> value : myWorkItem.elementsIterable()) {
        myValues.add(value);
      }
      for (WindowedValue<V> value : that.getValue().elementsIterable()) {
        thatValues.add(value);
      }
      try {
        return myValues.equals(thatValues)
            && keyCoder
                .structuralValue(myWorkItem.key())
                .equals(keyCoder.structuralValue(that.getValue().key()));
      } catch (Exception e) {
        return false;
      }
    }

    @Override
    public void describeTo(Description description) {
      description
          .appendText("KeyedWorkItem<K, V> containing key ")
          .appendValue(myWorkItem.key())
          .appendText(" and values ")
          .appendValueList("[", ", ", "]", myWorkItem.elementsIterable());
    }
  }
}