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

import static org.apache.beam.runners.dataflow.util.CloudObject.fromSpec;
import static org.apache.beam.runners.dataflow.util.CloudObjects.coderFromCloudObject;
import static org.apache.beam.runners.dataflow.util.Structs.getBytes;
import static org.apache.beam.runners.dataflow.util.Structs.getObject;
import static org.apache.beam.sdk.util.SerializableUtils.deserializeFromByteArray;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executors;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.DoFnRunners.OutputManager;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItemCoder;
import org.apache.beam.runners.core.OutputAndTimeBoundedSplittableProcessElementInvoker;
import org.apache.beam.runners.core.OutputWindowedValue;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.SimpleDoFnRunner;
import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.dataflow.util.CloudObject;
import org.apache.beam.runners.dataflow.util.PropertyNames;
import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.util.DoFnInfo;
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.Duration;
import org.joda.time.Instant;

/**
 * A {@link ParDoFnFactory} to create instances of user {@link ProcessFn} according to
 * specifications from the Dataflow service.
 */
class SplittableProcessFnFactory {
  static final ParDoFnFactory createDefault() {
    return new UserParDoFnFactory(new ProcessFnExtractor(), new SplittableDoFnRunnerFactory());
  }

  private static class ProcessFnExtractor implements UserParDoFnFactory.DoFnExtractor {
    @Override
    public DoFnInfo<?, ?> getDoFnInfo(CloudObject cloudUserFn) throws Exception {
      DoFnInfo<?, ?> doFnInfo =
          (DoFnInfo<?, ?>)
              deserializeFromByteArray(
                  getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN), "Serialized DoFnInfo");
      Coder restrictionCoder =
          coderFromCloudObject(
              fromSpec(getObject(cloudUserFn, WorkerPropertyNames.RESTRICTION_CODER)));

      ProcessFn processFn =
          new ProcessFn(
              doFnInfo.getDoFn(),
              doFnInfo.getInputCoder(),
              restrictionCoder,
              doFnInfo.getWindowingStrategy());

      return DoFnInfo.forFn(
          processFn,
          doFnInfo.getWindowingStrategy(),
          doFnInfo.getSideInputViews(),
          KeyedWorkItemCoder.of(
              ByteArrayCoder.of(),
              KvCoder.of(doFnInfo.getInputCoder(), restrictionCoder),
              doFnInfo.getWindowingStrategy().getWindowFn().windowCoder()),
          doFnInfo.getOutputCoders(),
          doFnInfo.getMainOutput(),
          doFnInfo.getDoFnSchemaInformation(),
          doFnInfo.getSideInputMapping());
    }
  }

  private static class SplittableDoFnRunnerFactory<
          InputT,
          OutputT,
          RestrictionT,
          PositionT,
          TrackerT extends RestrictionTracker<RestrictionT, PositionT>>
      implements DoFnRunnerFactory<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> {
    @Override
    public DoFnRunner<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> createRunner(
        DoFn<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> fn,
        PipelineOptions options,
        TupleTag<OutputT> mainOutputTag,
        List<TupleTag<?>> sideOutputTags,
        Iterable<PCollectionView<?>> sideInputViews,
        SideInputReader sideInputReader,
        Coder<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>> inputCoder,
        Map<TupleTag<?>, Coder<?>> outputCoders,
        WindowingStrategy<?, ?> windowingStrategy,
        DataflowExecutionContext.DataflowStepContext stepContext,
        DataflowExecutionContext.DataflowStepContext userStepContext,
        OutputManager outputManager,
        DoFnSchemaInformation doFnSchemaInformation,
        Map<String, PCollectionView<?>> sideInputMapping) {
      ProcessFn<InputT, OutputT, RestrictionT, TrackerT> processFn =
          (ProcessFn<InputT, OutputT, RestrictionT, TrackerT>) fn;
      processFn.setStateInternalsFactory(key -> (StateInternals) stepContext.stateInternals());
      processFn.setTimerInternalsFactory(key -> stepContext.timerInternals());
      processFn.setProcessElementInvoker(
          new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
              processFn.getFn(),
              options,
              new OutputWindowedValue<OutputT>() {
                @Override
                public void outputWindowedValue(
                    OutputT output,
                    Instant timestamp,
                    Collection<? extends BoundedWindow> windows,
                    PaneInfo pane) {
                  outputManager.output(
                      mainOutputTag, WindowedValue.of(output, timestamp, windows, pane));
                }

                @Override
                public <T> void outputWindowedValue(
                    TupleTag<T> tag,
                    T output,
                    Instant timestamp,
                    Collection<? extends BoundedWindow> windows,
                    PaneInfo pane) {
                  outputManager.output(tag, WindowedValue.of(output, timestamp, windows, pane));
                }
              },
              sideInputReader,
              Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()),
              // Commit at least once every 10 seconds or 10k records.  This keeps the watermark
              // advancing smoothly, and ensures that not too much work will have to be reprocessed
              // in the event of a crash.
              10000,
              Duration.standardSeconds(10)));
      DoFnRunner<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> simpleRunner =
          new SimpleDoFnRunner<>(
              options,
              processFn,
              sideInputReader,
              outputManager,
              mainOutputTag,
              sideOutputTags,
              userStepContext,
              inputCoder,
              outputCoders,
              processFn.getInputWindowingStrategy(),
              doFnSchemaInformation,
              sideInputMapping);
      DoFnRunner<KeyedWorkItem<byte[], KV<InputT, RestrictionT>>, OutputT> fnRunner =
          new DataflowProcessFnRunner<>(simpleRunner);
      boolean hasStreamingSideInput =
          options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty();
      KeyedWorkItemCoder<byte[], KV<InputT, RestrictionT>> kwiCoder =
          (KeyedWorkItemCoder<byte[], KV<InputT, RestrictionT>>) inputCoder;
      if (hasStreamingSideInput) {
        fnRunner =
            new StreamingKeyedWorkItemSideInputDoFnRunner<>(
                fnRunner,
                ByteArrayCoder.of(),
                new StreamingSideInputFetcher<>(
                    sideInputViews,
                    kwiCoder.getElementCoder(),
                    processFn.getInputWindowingStrategy(),
                    (StreamingModeExecutionContext.StreamingModeStepContext) userStepContext),
                userStepContext);
      }
      return fnRunner;
    }
  }
}