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

import java.util.HashMap;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.FlinkStreamingPortablePipelineTranslator;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.Preconditions;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.Collector;
import org.joda.time.Instant;

/**
 * An adapter class that allows one to apply Apache Beam PTransforms directly to Flink DataStreams.
 */
public class BeamFlinkDataStreamAdapter {
  private final PipelineOptions pipelineOptions;
  private final CoderRegistry coderRegistry = CoderRegistry.createDefault();

  public BeamFlinkDataStreamAdapter() {
    this(PipelineOptionsFactory.create());
  }

  public BeamFlinkDataStreamAdapter(PipelineOptions pipelineOptions) {
    this.pipelineOptions = pipelineOptions;
  }

  public <InputT, OutputT, CollectionT extends PCollection<? extends InputT>>
      DataStream<OutputT> applyBeamPTransform(
          DataStream<InputT> input, PTransform<CollectionT, PCollection<OutputT>> transform) {
    return (DataStream)
        getNonNull(
            applyBeamPTransformInternal(
                ImmutableMap.of("input", input),
                (pipeline, map) -> (CollectionT) getNonNull(map, "input"),
                (output) -> ImmutableMap.of("output", output),
                transform,
                input.getExecutionEnvironment()),
            "output");
  }

  public <OutputT> DataStream<OutputT> applyBeamPTransform(
      Map<String, ? extends DataStream<?>> inputs,
      PTransform<PCollectionTuple, PCollection<OutputT>> transform) {
    return (DataStream)
        getNonNull(
            applyBeamPTransformInternal(
                inputs,
                BeamAdapterUtils::mapToTuple,
                (output) -> ImmutableMap.of("output", output),
                transform,
                inputs.values().stream().findAny().get().getExecutionEnvironment()),
            "output");
  }

  public <OutputT> DataStream<OutputT> applyBeamPTransform(
      StreamExecutionEnvironment executionEnvironment,
      PTransform<PBegin, PCollection<OutputT>> transform) {
    return (DataStream)
        getNonNull(
            applyBeamPTransformInternal(
                ImmutableMap.of(),
                (pipeline, map) -> PBegin.in(pipeline),
                (output) -> ImmutableMap.of("output", output),
                transform,
                executionEnvironment),
            "output");
  }

  public <InputT, CollectionT extends PCollection<? extends InputT>>
      Map<String, DataStream<?>> applyMultiOutputBeamPTransform(
          DataStream<InputT> input, PTransform<CollectionT, PCollectionTuple> transform) {
    return applyBeamPTransformInternal(
        ImmutableMap.of("input", input),
        (pipeline, map) -> (CollectionT) getNonNull(map, "input"),
        BeamAdapterUtils::tupleToMap,
        transform,
        input.getExecutionEnvironment());
  }

  public Map<String, DataStream<?>> applyMultiOutputBeamPTransform(
      Map<String, ? extends DataStream<?>> inputs,
      PTransform<PCollectionTuple, PCollectionTuple> transform) {
    return applyBeamPTransformInternal(
        inputs,
        BeamAdapterUtils::mapToTuple,
        BeamAdapterUtils::tupleToMap,
        transform,
        inputs.values().stream().findAny().get().getExecutionEnvironment());
  }

  public Map<String, DataStream<?>> applyMultiOutputBeamPTransform(
      StreamExecutionEnvironment executionEnvironment,
      PTransform<PBegin, PCollectionTuple> transform) {
    return applyBeamPTransformInternal(
        ImmutableMap.of(),
        (pipeline, map) -> PBegin.in(pipeline),
        BeamAdapterUtils::tupleToMap,
        transform,
        executionEnvironment);
  }

  public <InputT, CollectionT extends PCollection<? extends InputT>>
      void applyNoOutputBeamPTransform(
          DataStream<InputT> input, PTransform<CollectionT, PDone> transform) {
    applyBeamPTransformInternal(
        ImmutableMap.of("input", input),
        (pipeline, map) -> (CollectionT) getNonNull(map, "input"),
        pDone -> ImmutableMap.of(),
        transform,
        input.getExecutionEnvironment());
  }

  public void applyNoOutputBeamPTransform(
      Map<String, ? extends DataStream<?>> inputs, PTransform<PCollectionTuple, PDone> transform) {
    applyBeamPTransformInternal(
        inputs,
        BeamAdapterUtils::mapToTuple,
        pDone -> ImmutableMap.of(),
        transform,
        inputs.values().stream().findAny().get().getExecutionEnvironment());
  }

  public void applyNoOutputBeamPTransform(
      StreamExecutionEnvironment executionEnvironment, PTransform<PBegin, PDone> transform) {
    applyBeamPTransformInternal(
        ImmutableMap.of(),
        (pipeline, map) -> PBegin.in(pipeline),
        pDone -> ImmutableMap.of(),
        transform,
        executionEnvironment);
  }

  private <BeamInputT extends PInput, BeamOutputT extends POutput>
      Map<String, DataStream<?>> applyBeamPTransformInternal(
          Map<String, ? extends DataStream<?>> inputs,
          BiFunction<Pipeline, Map<String, PCollection<?>>, BeamInputT> toBeamInput,
          Function<BeamOutputT, Map<String, PCollection<?>>> fromBeamOutput,
          PTransform<? super BeamInputT, BeamOutputT> transform,
          StreamExecutionEnvironment executionEnvironment) {
    return BeamAdapterUtils.applyBeamPTransformInternal(
        inputs,
        toBeamInput,
        fromBeamOutput,
        transform,
        executionEnvironment,
        false,
        dataStream -> dataStream.getType(),
        pipelineOptions,
        coderRegistry,
        (flinkInputs, pipelineProto, env) -> {
          Map<String, DataStream<?>> flinkOutputs = new HashMap<>();
          FlinkStreamingPortablePipelineTranslator translator =
              new FlinkStreamingPortablePipelineTranslator(
                  ImmutableMap.of(
                      FlinkInput.URN, flinkInputTranslator(flinkInputs),
                      FlinkOutput.URN, flinkOutputTranslator(flinkOutputs)));
          FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext context =
              translator.createTranslationContext(
                  JobInfo.create(
                      "unusedJobId",
                      "unusedJobName",
                      "unusedRetrievalToken",
                      PipelineOptionsTranslation.toProto(pipelineOptions)),
                  pipelineOptions.as(FlinkPipelineOptions.class),
                  env);
          translator.translate(context, translator.prepareForTranslation(pipelineProto));
          return flinkOutputs;
        });
  }

  private <InputT>
      FlinkStreamingPortablePipelineTranslator.PTransformTranslator<
              FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext>
          flinkInputTranslator(Map<String, ? extends DataStream<?>> inputMap) {
    return (String id,
        RunnerApi.Pipeline p,
        FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext context) -> {
      // When we run into a FlinkInput operator, it "produces" the corresponding input as its
      // "computed result."
      RunnerApi.PTransform transform = p.getComponents().getTransformsOrThrow(id);
      String inputId = transform.getSpec().getPayload().toStringUtf8();
      DataStream<InputT> flinkInput =
          Preconditions.checkStateNotNull(
              (DataStream<InputT>) inputMap.get(inputId),
              "missing input referenced in proto: %s",
              inputId);
      context.addDataStream(
          Iterables.getOnlyElement(transform.getOutputsMap().values()),
          flinkInput.process(
              new ProcessFunction<InputT, WindowedValue<InputT>>() {
                @Override
                public void processElement(
                    InputT value,
                    ProcessFunction<InputT, WindowedValue<InputT>>.Context ctx,
                    Collector<WindowedValue<InputT>> out)
                    throws Exception {
                  out.collect(
                      WindowedValues.timestampedValueInGlobalWindow(
                          value,
                          ctx.timestamp() == null
                              ? BoundedWindow.TIMESTAMP_MIN_VALUE
                              : Instant.ofEpochMilli(ctx.timestamp())));
                }
              },
              BeamAdapterCoderUtils.coderToTypeInformation(
                  WindowedValues.getFullCoder(
                      BeamAdapterCoderUtils.typeInformationToCoder(
                          flinkInput.getType(), coderRegistry),
                      GlobalWindow.Coder.INSTANCE),
                  pipelineOptions)));
    };
  }

  private <InputT>
      FlinkStreamingPortablePipelineTranslator.PTransformTranslator<
              FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext>
          flinkOutputTranslator(Map<String, DataStream<?>> outputMap) {
    return (String id,
        RunnerApi.Pipeline p,
        FlinkStreamingPortablePipelineTranslator.StreamingTranslationContext context) -> {
      // When we run into a FlinkOutput operator, we cache the computed PCollection to return to the
      // user.
      RunnerApi.PTransform transform = p.getComponents().getTransformsOrThrow(id);
      DataStream<WindowedValue<InputT>> inputDataStream =
          context.getDataStreamOrThrow(Iterables.getOnlyElement(transform.getInputsMap().values()));
      String outputId = transform.getSpec().getPayload().toStringUtf8();
      Coder<InputT> outputCoder =
          BeamAdapterCoderUtils.lookupCoder(
              p, Iterables.getOnlyElement(transform.getInputsMap().values()));
      outputMap.put(
          outputId,
          inputDataStream.transform(
              "StripWindows",
              BeamAdapterCoderUtils.coderToTypeInformation(outputCoder, pipelineOptions),
              new UnwrapWindowOperator<InputT>()));
    };
  }

  /**
   * Forwards the Beam timestamps to the underlying Flink timestamps, but unlike {@link
   * DataStream#assignTimestampsAndWatermarks(WatermarkStrategy)} does not discard the underlying
   * watermark signals.
   *
   * @param <T> user element type
   */
  private static class UnwrapWindowOperator<T> extends AbstractStreamOperator<T>
      implements OneInputStreamOperator<WindowedValue<T>, T> {
    @Override
    public void processElement(StreamRecord<WindowedValue<T>> element) {
      output.collect(
          element.replace(
              element.getValue().getValue(), element.getValue().getTimestamp().getMillis()));
    }
  }

  private <K, V> V getNonNull(Map<K, V> map, K key) {
    return Preconditions.checkStateNotNull(map.get(Preconditions.checkArgumentNotNull(key)));
  }
}
