/*
 * 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 com.google.auto.value.AutoValue;
import java.util.Collections;
import java.util.NavigableSet;
import javax.annotation.Nullable;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItems;
import org.apache.beam.runners.core.StateNamespaces.WindowNamespace;
import org.apache.beam.runners.core.StateTag;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.TimerInternals.TimerData;
import org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext;
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.options.PipelineOptions;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.state.WatermarkHoldState;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
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.vendor.guava.v26_0_jre.com.google.common.cache.CacheLoader;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.joda.time.Instant;

/** A {@link TransformEvaluatorFactory} for stateful {@link ParDo}. */
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
final class StatefulParDoEvaluatorFactory<K, InputT, OutputT> implements TransformEvaluatorFactory {

  private final ParDoEvaluatorFactory<KV<K, InputT>, OutputT> delegateFactory;

  private final EvaluationContext evaluationContext;

  StatefulParDoEvaluatorFactory(EvaluationContext evaluationContext, PipelineOptions options) {
    this.delegateFactory =
        new ParDoEvaluatorFactory<>(
            evaluationContext,
            ParDoEvaluator.<KV<K, InputT>, OutputT>defaultRunnerFactory(),
            new CacheLoader<AppliedPTransform<?, ?, ?>, DoFnLifecycleManager>() {
              @Override
              public DoFnLifecycleManager load(AppliedPTransform<?, ?, ?> appliedStatefulParDo)
                  throws Exception {
                // StatefulParDo is overridden after the portable pipeline is received, so we
                // do not go through the portability translation layers
                StatefulParDo<?, ?, ?> statefulParDo =
                    (StatefulParDo<?, ?, ?>) appliedStatefulParDo.getTransform();
                return DoFnLifecycleManager.of(statefulParDo.getDoFn(), options);
              }
            },
            options);
    this.evaluationContext = evaluationContext;
  }

  @Override
  public <T> TransformEvaluator<T> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) throws Exception {
    @SuppressWarnings({"unchecked", "rawtypes"})
    TransformEvaluator<T> evaluator =
        (TransformEvaluator<T>)
            createEvaluator((AppliedPTransform) application, (CommittedBundle) inputBundle);
    return evaluator;
  }

  @Override
  public void cleanup() throws Exception {
    delegateFactory.cleanup();
  }

  @SuppressWarnings({"unchecked", "rawtypes"})
  private TransformEvaluator<KeyedWorkItem<K, KV<K, InputT>>> createEvaluator(
      AppliedPTransform<
              PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>,
              PCollectionTuple,
              StatefulParDo<K, InputT, OutputT>>
          application,
      CommittedBundle<KeyedWorkItem<K, KV<K, InputT>>> inputBundle)
      throws Exception {

    DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator =
        delegateFactory.createEvaluator(
            (AppliedPTransform) application,
            (PCollection) inputBundle.getPCollection(),
            inputBundle.getKey(),
            application.getTransform().getSideInputs(),
            application.getTransform().getMainOutputTag(),
            application.getTransform().getAdditionalOutputTags().getAll(),
            application.getTransform().getSchemaInformation(),
            application.getTransform().getSideInputMapping());

    DirectStepContext stepContext =
        evaluationContext
            .getExecutionContext(application, inputBundle.getKey())
            .getStepContext(evaluationContext.getStepName(application));

    stepContext.stateInternals().commit();
    return new StatefulParDoEvaluator<>(delegateEvaluator, stepContext);
  }

  @AutoValue
  abstract static class AppliedPTransformOutputKeyAndWindow<K, InputT, OutputT> {
    abstract AppliedPTransform<
            PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>,
            PCollectionTuple,
            StatefulParDo<K, InputT, OutputT>>
        getTransform();

    abstract StructuralKey<K> getKey();

    abstract BoundedWindow getWindow();

    static <K, InputT, OutputT> AppliedPTransformOutputKeyAndWindow<K, InputT, OutputT> create(
        AppliedPTransform<
                PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>,
                PCollectionTuple,
                StatefulParDo<K, InputT, OutputT>>
            transform,
        StructuralKey<K> key,
        BoundedWindow w) {
      return new AutoValue_StatefulParDoEvaluatorFactory_AppliedPTransformOutputKeyAndWindow<>(
          transform, key, w);
    }
  }

  private static class StatefulParDoEvaluator<K, InputT>
      implements TransformEvaluator<KeyedWorkItem<K, KV<K, InputT>>> {

    private final DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator;
    private final DirectTimerInternals timerInternals;

    DirectStepContext stepContext;

    public StatefulParDoEvaluator(
        DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator,
        DirectStepContext stepContext) {
      this.delegateEvaluator = delegateEvaluator;
      this.timerInternals = delegateEvaluator.getParDoEvaluator().getStepContext().timerInternals();
      this.stepContext = stepContext;
    }

    @Override
    public void processElement(WindowedValue<KeyedWorkItem<K, KV<K, InputT>>> gbkResult)
        throws Exception {
      for (WindowedValue<KV<K, InputT>> windowedValue : gbkResult.getValue().elementsIterable()) {
        delegateEvaluator.processElement(windowedValue);
      }

      for (TimerData timerData : gbkResult.getValue().timersIterable()) {
        // Get any new or modified timers that are earlier than the current one. In order to
        // maintain timer ordering,
        // we need to fire these timers first.
        NavigableSet<TimerData> earlierTimers =
            timerInternals.getModifiedTimersOrdered(timerData.getDomain()).headSet(timerData, true);
        while (!earlierTimers.isEmpty()) {
          TimerData insertedTimer = earlierTimers.pollFirst();
          if (timerModified(insertedTimer)) {
            continue;
          }
          // Make sure to register this timer as deleted. This could be a timer that was originally
          // set for the future
          // and not in the bundle but was reset to an earlier time in this bundle. If we don't
          // explicity delete the
          // future timer, then it will still fire.
          timerInternals.deleteTimer(insertedTimer);
          processTimer(insertedTimer, gbkResult.getValue().key());
        }

        // As long as the timer hasn't been modified or deleted earlier in the bundle, fire it.
        if (!timerModified(timerData)) {
          processTimer(timerData, gbkResult.getValue().key());
        }
      }
    }

    // Check to see if a timer has been modified inside this bundle.
    private boolean timerModified(TimerData timerData) {
      @Nullable
      TimerData modifiedTimer = timerInternals.getModifiedTimerIds().get(timerData.stringKey());
      return modifiedTimer != null && !modifiedTimer.equals(timerData);
    }

    private void processTimer(TimerData timerData, K key) throws Exception {
      WindowNamespace<?> windowNamespace = (WindowNamespace) timerData.getNamespace();
      BoundedWindow timerWindow = windowNamespace.getWindow();
      delegateEvaluator.onTimer(timerData, key, timerWindow);
      clearWatermarkHold(timerData);
    }

    private void clearWatermarkHold(TimerData timer) {
      StateTag<WatermarkHoldState> timerWatermarkHoldTag = setTimerTag(timer);
      stepContext.stateInternals().state(timer.getNamespace(), timerWatermarkHoldTag).clear();
      stepContext.stateInternals().commit();
    }

    private void setWatermarkHold(TimerData timer) {
      StateTag<WatermarkHoldState> timerWatermarkHoldTag = setTimerTag(timer);
      stepContext
          .stateInternals()
          .state(timer.getNamespace(), timerWatermarkHoldTag)
          .add(timer.getOutputTimestamp());
    }

    @Override
    public TransformResult<KeyedWorkItem<K, KV<K, InputT>>> finishBundle() throws Exception {

      TransformResult<KV<K, InputT>> delegateResult = delegateEvaluator.finishBundle();
      boolean isTimerDeclared = false;
      for (TimerData timerData : delegateResult.getTimerUpdate().getSetTimers()) {
        setWatermarkHold(timerData);
        isTimerDeclared = true;
      }
      for (TimerData timerData : delegateResult.getTimerUpdate().getDeletedTimers()) {
        clearWatermarkHold(timerData);
      }

      CopyOnAccessInMemoryStateInternals state;
      Instant watermarkHold;

      if (isTimerDeclared && delegateResult.getState() != null) { // For both State and Timer Holds
        state = delegateResult.getState();
        watermarkHold = stepContext.commitState().getEarliestWatermarkHold();
      } else if (isTimerDeclared) { // For only Timer holds
        state = stepContext.commitState();
        watermarkHold = state.getEarliestWatermarkHold();
      } else { // For only State ( non Timer ) holds
        state = delegateResult.getState();
        watermarkHold = delegateResult.getWatermarkHold();
      }

      TimerUpdate timerUpdate = delegateResult.getTimerUpdate();
      StepTransformResult.Builder<KeyedWorkItem<K, KV<K, InputT>>> regroupedResult =
          StepTransformResult.<KeyedWorkItem<K, KV<K, InputT>>>withHold(
                  delegateResult.getTransform(), watermarkHold)
              .withTimerUpdate(timerUpdate)
              .withState(state)
              .withMetricUpdates(delegateResult.getLogicalMetricUpdates())
              .addOutput(Lists.newArrayList(delegateResult.getOutputBundles()))
              .withBundleFinalizations(delegateResult.getBundleFinalizations());

      // The delegate may have pushed back unprocessed elements across multiple keys and windows.
      // Since processing is single-threaded per key and window, we don't need to regroup the
      // outputs, but just make a bunch of singletons
      for (WindowedValue<?> untypedUnprocessed : delegateResult.getUnprocessedElements()) {
        WindowedValue<KV<K, InputT>> windowedKv = (WindowedValue<KV<K, InputT>>) untypedUnprocessed;
        WindowedValue<KeyedWorkItem<K, KV<K, InputT>>> pushedBack =
            windowedKv.withValue(
                KeyedWorkItems.elementsWorkItem(
                    windowedKv.getValue().getKey(), Collections.singleton(windowedKv)));

        regroupedResult.addUnprocessedElements(pushedBack);
      }

      return regroupedResult.build();
    }
  }

  private static StateTag<WatermarkHoldState> setTimerTag(TimerData timerData) {
    return StateTags.makeSystemTagInternal(
        StateTags.watermarkStateInternal(
            "timer-" + timerData.getTimerId() + "+" + timerData.getTimerFamilyId(),
            TimestampCombiner.EARLIEST));
  }
}
