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

import com.hazelcast.jet.core.Processor;
import com.hazelcast.jet.core.Watermark;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import javax.annotation.Nonnull;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.DoFnRunners;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.InMemoryTimerInternals;
import org.apache.beam.runners.core.SideInputReader;
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.StepContext;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.joda.time.Instant;

/**
 * Jet {@link com.hazelcast.jet.core.Processor} implementation for Beam's stateful ParDo primitive.
 */
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
public class StatefulParDoP<OutputT>
    extends AbstractParDoP<KV<?, ?>, OutputT> { // todo: unify with ParDoP?

  private KeyedStepContext keyedStepContext;

  private StatefulParDoP(
      DoFn<KV<?, ?>, OutputT> doFn,
      WindowingStrategy<?, ?> windowingStrategy,
      DoFnSchemaInformation doFnSchemaInformation,
      Map<TupleTag<?>, int[]> outputCollToOrdinals,
      SerializablePipelineOptions pipelineOptions,
      TupleTag<OutputT> mainOutputTag,
      Coder<KV<?, ?>> inputCoder,
      Map<PCollectionView<?>, Coder<?>> sideInputCoders,
      Map<TupleTag<?>, Coder<?>> outputCoders,
      Coder<KV<?, ?>> inputValueCoder,
      Map<TupleTag<?>, Coder<?>> outputValueCoders,
      Map<Integer, PCollectionView<?>> ordinalToSideInput,
      Map<String, PCollectionView<?>> sideInputMapping,
      String ownerId,
      String stepId) {
    super(
        doFn,
        windowingStrategy,
        doFnSchemaInformation,
        outputCollToOrdinals,
        pipelineOptions,
        mainOutputTag,
        inputCoder,
        sideInputCoders,
        outputCoders,
        inputValueCoder,
        outputValueCoders,
        ordinalToSideInput,
        sideInputMapping,
        ownerId,
        stepId);
  }

  private static void fireTimer(
      Object key, TimerInternals.TimerData timer, DoFnRunner<KV<?, ?>, ?> doFnRunner) {
    StateNamespace namespace = timer.getNamespace();
    BoundedWindow window = ((StateNamespaces.WindowNamespace) namespace).getWindow();
    doFnRunner.onTimer(
        timer.getTimerId(),
        timer.getTimerFamilyId(),
        key,
        window,
        timer.getTimestamp(),
        timer.getOutputTimestamp(),
        timer.getDomain());
  }

  @Override
  protected DoFnRunner<KV<?, ?>, OutputT> getDoFnRunner(
      PipelineOptions pipelineOptions,
      DoFn<KV<?, ?>, OutputT> doFn,
      SideInputReader sideInputReader,
      JetOutputManager outputManager,
      TupleTag<OutputT> mainOutputTag,
      List<TupleTag<?>> additionalOutputTags,
      Coder<KV<?, ?>> inputValueCoder,
      Map<TupleTag<?>, Coder<?>> outputValueCoders,
      WindowingStrategy<?, ?> windowingStrategy,
      DoFnSchemaInformation doFnSchemaInformation,
      Map<String, PCollectionView<?>> sideInputMapping) {
    keyedStepContext = new KeyedStepContext();
    return DoFnRunners.simpleRunner(
        pipelineOptions,
        doFn,
        sideInputReader,
        outputManager,
        mainOutputTag,
        additionalOutputTags,
        keyedStepContext,
        inputValueCoder,
        outputValueCoders,
        windowingStrategy,
        doFnSchemaInformation,
        sideInputMapping);
  }

  @Override
  protected void startRunnerBundle(DoFnRunner<KV<?, ?>, OutputT> runner) {
    keyedStepContext.advanceProcessingTimes();
    super.startRunnerBundle(runner);
  }

  @Override
  protected void processElementWithRunner(
      DoFnRunner<KV<?, ?>, OutputT> runner, WindowedValue<KV<?, ?>> windowedValue) {
    KV<?, ?> kv = windowedValue.getValue();
    Object key = kv.getKey();

    keyedStepContext.setKey(key);
    super.processElementWithRunner(runner, windowedValue);
    keyedStepContext.clearKey();
  }

  @Override
  public boolean tryProcessWatermark(@Nonnull Watermark watermark) {
    return flushTimers(watermark.timestamp()) && super.tryProcessWatermark(watermark);
  }

  @Override
  public boolean complete() {
    return flushTimers(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) && super.complete();
  }

  private boolean flushTimers(long watermark) {
    keyedStepContext.flushTimers(watermark);
    return outputManager.tryFlush();
  }

  private void fireEligibleTimers(Object key, InMemoryTimerInternals timerInternals) {
    while (true) {
      TimerInternals.TimerData timer;
      boolean hasFired = false;

      while ((timer = timerInternals.removeNextEventTimer()) != null) {
        hasFired = true;
        fireTimer(key, timer, doFnRunner);
      }

      while ((timer = timerInternals.removeNextProcessingTimer()) != null) {
        hasFired = true;
        fireTimer(key, timer, doFnRunner);
      }

      while ((timer = timerInternals.removeNextSynchronizedProcessingTimer()) != null) {
        hasFired = true;
        fireTimer(key, timer, doFnRunner);
      }

      if (!hasFired) {
        break;
      }
    }
  }

  /**
   * Jet {@link Processor} supplier that will provide instances of {@link StatefulParDoP}.
   *
   * @param <OutputT> the type of main output elements of the DoFn being used
   */
  public static class Supplier<OutputT> extends AbstractSupplier<KV<?, ?>, OutputT> {

    public Supplier(
        String stepId,
        String ownerId,
        DoFn<KV<?, ?>, OutputT> doFn,
        WindowingStrategy<?, ?> windowingStrategy,
        DoFnSchemaInformation doFnSchemaInformation,
        SerializablePipelineOptions pipelineOptions,
        TupleTag<OutputT> mainOutputTag,
        Set<TupleTag<OutputT>> allOutputTags,
        Coder<KV<?, ?>> inputCoder,
        Map<PCollectionView<?>, Coder<?>> sideInputCoders,
        Map<TupleTag<?>, Coder<?>> outputCoders,
        Coder<KV<?, ?>> inputValueCoder,
        Map<TupleTag<?>, Coder<?>> outputValueCoders,
        Collection<PCollectionView<?>> sideInputs,
        Map<String, PCollectionView<?>> sideInputMapping) {
      super(
          stepId,
          ownerId,
          doFn,
          windowingStrategy,
          doFnSchemaInformation,
          pipelineOptions,
          mainOutputTag,
          allOutputTags,
          inputCoder,
          sideInputCoders,
          outputCoders,
          inputValueCoder,
          outputValueCoders,
          sideInputs,
          sideInputMapping);
    }

    @Override
    Processor getEx(
        DoFn<KV<?, ?>, OutputT> doFn,
        WindowingStrategy<?, ?> windowingStrategy,
        DoFnSchemaInformation doFnSchemaInformation,
        Map<TupleTag<?>, int[]> outputCollToOrdinals,
        SerializablePipelineOptions pipelineOptions,
        TupleTag<OutputT> mainOutputTag,
        Coder<KV<?, ?>> inputCoder,
        Map<PCollectionView<?>, Coder<?>> sideInputCoders,
        Map<TupleTag<?>, Coder<?>> outputCoders,
        Coder<KV<?, ?>> inputValueCoder,
        Map<TupleTag<?>, Coder<?>> outputValueCoders,
        Map<Integer, PCollectionView<?>> ordinalToSideInput,
        Map<String, PCollectionView<?>> sideInputMapping,
        String ownerId,
        String stepId) {
      return new StatefulParDoP<>(
          doFn,
          windowingStrategy,
          doFnSchemaInformation,
          outputCollToOrdinals,
          pipelineOptions,
          mainOutputTag,
          inputCoder,
          sideInputCoders,
          outputCoders,
          inputValueCoder,
          outputValueCoders,
          ordinalToSideInput,
          sideInputMapping,
          ownerId,
          stepId);
    }
  }

  private class KeyedStepContext implements StepContext {

    private final Object nullKey = new Object();

    private final ConcurrentHashMap<Object, InMemoryStateInternals> keyedStateInternals;
    private final ConcurrentHashMap<Object, InMemoryTimerInternals> keyedTimerInternals;

    @SuppressWarnings("ThreadLocalUsage")
    private final ThreadLocal<Object> currentKey = new ThreadLocal<>();

    KeyedStepContext() {
      this.keyedStateInternals = new ConcurrentHashMap<>();
      this.keyedTimerInternals = new ConcurrentHashMap<>();
    }

    void setKey(Object key) {
      Object normalizedKey = key == null ? nullKey : key;
      currentKey.set(normalizedKey);
      keyedStateInternals.computeIfAbsent(normalizedKey, InMemoryStateInternals::forKey);
      keyedTimerInternals.computeIfAbsent(normalizedKey, k -> new InMemoryTimerInternals());
    }

    void clearKey() {
      currentKey.remove();
    }

    @Override
    public StateInternals stateInternals() {
      Object key = currentKey.get();
      if (key == null) {
        throw new IllegalStateException("Active key should be set");
      }
      return keyedStateInternals.get(key);
    }

    @Override
    public TimerInternals timerInternals() {
      Object key = currentKey.get();
      if (key == null) {
        throw new IllegalStateException("Active key should be set");
      }
      return keyedTimerInternals.get(key);
    }

    public void advanceProcessingTimes() {
      Instant now = Instant.now();
      keyedTimerInternals
          .values()
          .forEach(
              timerInternals -> {
                try {
                  timerInternals.advanceProcessingTime(now);
                  timerInternals.advanceSynchronizedProcessingTime(now);
                } catch (Exception e) {
                  throw new RuntimeException("Failed advancing time!");
                }
              });
    }

    public void flushTimers(long watermark) {
      Instant watermarkInstant = new Instant(watermark);
      keyedTimerInternals
          .entrySet()
          .forEach(
              (entry) -> {
                InMemoryTimerInternals timerInternals = entry.getValue();
                if (timerInternals.currentInputWatermarkTime().isBefore(watermark)) {
                  try {
                    timerInternals.advanceInputWatermark(watermarkInstant);
                    if (watermarkInstant.equals(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
                      timerInternals.advanceProcessingTime(watermarkInstant);
                      timerInternals.advanceSynchronizedProcessingTime(watermarkInstant);
                    }
                    Object key = entry.getKey();
                    setKey(key);
                    fireEligibleTimers(key, timerInternals);
                    clearKey();
                  } catch (Exception e) {
                    throw new RuntimeException("Failed advancing processing time", e);
                  }
                }
              });
    }
  }
}
