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

import com.hazelcast.jet.core.Processor;
import com.hazelcast.jet.core.ProcessorMetaSupplier;
import com.hazelcast.jet.core.Vertex;
import com.hazelcast.jet.function.SupplierEx;
import com.hazelcast.jet.impl.util.ExceptionUtil;
import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.beam.runners.core.construction.CreatePCollectionViewTranslation;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.ParDoTranslation;
import org.apache.beam.runners.core.construction.ReadTranslation;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.jet.processors.AssignWindowP;
import org.apache.beam.runners.jet.processors.BoundedSourceP;
import org.apache.beam.runners.jet.processors.FlattenP;
import org.apache.beam.runners.jet.processors.ImpulseP;
import org.apache.beam.runners.jet.processors.ParDoP;
import org.apache.beam.runners.jet.processors.StatefulParDoP;
import org.apache.beam.runners.jet.processors.TestStreamP;
import org.apache.beam.runners.jet.processors.ViewP;
import org.apache.beam.runners.jet.processors.WindowGroupP;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.TransformHierarchy.Node;
import org.apache.beam.sdk.testing.TestStream;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
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.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;

@SuppressWarnings("unchecked")
class JetTransformTranslators {

  /** A map from a Transform URN to the translator. */
  private static final Map<String, JetTransformTranslator> TRANSLATORS = new HashMap<>();

  static {
    TRANSLATORS.put(PTransformTranslation.READ_TRANSFORM_URN, new ReadSourceTranslator());
    TRANSLATORS.put(PTransformTranslation.CREATE_VIEW_TRANSFORM_URN, new CreateViewTranslator());
    TRANSLATORS.put(PTransformTranslation.PAR_DO_TRANSFORM_URN, new ParDoTranslator());
    TRANSLATORS.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator());
    TRANSLATORS.put(PTransformTranslation.FLATTEN_TRANSFORM_URN, new FlattenTranslator());
    TRANSLATORS.put(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, new WindowTranslator());
    TRANSLATORS.put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator());
    TRANSLATORS.put(PTransformTranslation.TEST_STREAM_TRANSFORM_URN, new TestStreamTranslator());
  }

  static JetTransformTranslator<?> getTranslator(PTransform<?, ?> transform) {
    String urn = PTransformTranslation.urnForTransformOrNull(transform);
    return urn == null ? null : TRANSLATORS.get(urn);
  }

  private static class ReadSourceTranslator<T>
      implements JetTransformTranslator<PTransform<PBegin, PCollection<T>>> {

    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      if (!Utils.isBounded(appliedTransform)) {
        throw new UnsupportedOperationException(); // todo
      }

      BoundedSource<T> source;
      try {
        source =
            ReadTranslation.boundedSourceFromTransform(
                (AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>>)
                    appliedTransform);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder =
          Utils.getCoder((PCollection) Utils.getOutput(appliedTransform).getValue());

      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);
      SerializablePipelineOptions pipelineOptions = context.getOptions();
      ProcessorMetaSupplier processorSupplier =
          BoundedSourceP.supplier(source, pipelineOptions, outputCoder, vertexId);

      Vertex vertex = dagBuilder.addVertex(vertexId, processorSupplier);

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class ParDoTranslator
      implements JetTransformTranslator<PTransform<PCollection, PCollectionTuple>> {

    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      boolean usesStateOrTimers = Utils.usesStateOrTimers(appliedTransform);
      DoFn<?, ?> doFn = Utils.getDoFn(appliedTransform);

      Map<TupleTag<?>, PValue> outputs = Utils.getOutputs(appliedTransform);

      TupleTag<?> mainOutputTag;
      try {
        mainOutputTag = ParDoTranslation.getMainOutputTag(appliedTransform);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
      Map<TupleTag<?>, Integer> outputMap = new HashMap<>();
      int count = 1;
      for (TupleTag<?> tag : outputs.keySet()) {
        if (!outputMap.containsKey(tag)) {
          outputMap.put(tag, count++);
        }
      }
      final WindowingStrategy<?, ?> windowingStrategy =
          Utils.getWindowingStrategy(appliedTransform);

      Map<TupleTag<?>, Coder<?>> outputValueCoders = Utils.getOutputValueCoders(appliedTransform);
      Map<TupleTag<?>, Coder> outputCoders =
          Utils.getCoders(Utils.getOutputs(appliedTransform), Map.Entry::getKey);

      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String stepId =
          transformName.contains("/")
              ? transformName.substring(0, transformName.indexOf('/'))
              : transformName;
      String vertexId =
          dagBuilder.newVertexId(transformName) + (usesStateOrTimers ? " - STATEFUL" : "");
      SerializablePipelineOptions pipelineOptions = context.getOptions();
      Coder inputValueCoder = Utils.getInput(appliedTransform).getCoder();
      Coder inputCoder = Utils.getCoder(Utils.getInput(appliedTransform));
      List<PCollectionView<?>> sideInputs = Utils.getSideInputs(appliedTransform);
      Map<? extends PCollectionView<?>, Coder> sideInputCoders =
          sideInputs.stream()
              .collect(Collectors.toMap(si -> si, si -> Utils.getCoder(si.getPCollection())));
      DoFnSchemaInformation doFnSchemaInformation =
          ParDoTranslation.getSchemaInformation(appliedTransform);
      SupplierEx<Processor> processorSupplier =
          usesStateOrTimers
              ? new StatefulParDoP.Supplier(
                  stepId,
                  vertexId,
                  doFn,
                  windowingStrategy,
                  doFnSchemaInformation,
                  pipelineOptions,
                  mainOutputTag,
                  outputMap.keySet(),
                  inputCoder,
                  sideInputCoders,
                  outputCoders,
                  inputValueCoder,
                  outputValueCoders,
                  sideInputs)
              : new ParDoP.Supplier(
                  stepId,
                  vertexId,
                  doFn,
                  windowingStrategy,
                  doFnSchemaInformation,
                  pipelineOptions,
                  mainOutputTag,
                  outputMap.keySet(),
                  inputCoder,
                  sideInputCoders,
                  outputCoders,
                  inputValueCoder,
                  outputValueCoders,
                  sideInputs);

      Vertex vertex = dagBuilder.addVertex(vertexId, processorSupplier);
      dagBuilder.registerConstructionListeners((DAGBuilder.WiringListener) processorSupplier);

      PValue mainInput = Utils.getMainInput(pipeline, node);
      dagBuilder.registerEdgeEndPoint(Utils.getTupleTagId(mainInput), vertex);

      Map<TupleTag<?>, PValue> additionalInputs = Utils.getAdditionalInputs(node);
      if (additionalInputs != null && !additionalInputs.isEmpty()) {
        for (TupleTag<?> tupleTag : additionalInputs.keySet()) {
          dagBuilder.registerEdgeEndPoint(tupleTag.getId(), vertex);
        }
      }

      for (Map.Entry<TupleTag<?>, PValue> entry : outputs.entrySet()) {
        TupleTag<?> pCollId = entry.getKey();
        String edgeId = Utils.getTupleTagId(entry.getValue());
        dagBuilder.registerCollectionOfEdge(edgeId, pCollId.getId());
        dagBuilder.registerEdgeStartPoint(edgeId, vertex, outputCoders.get(pCollId));
      }

      return vertex;
    }
  }

  private static class GroupByKeyTranslator<K, InputT>
      implements JetTransformTranslator<
          PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, Iterable<InputT>>>>> {

    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      String transformName = appliedTransform.getFullName();

      PCollection<KV<K, InputT>> input = Utils.getInput(appliedTransform);
      Coder inputCoder = Utils.getCoder(input);
      Coder inputValueCoder = Utils.getInput(appliedTransform).getCoder();
      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder = Utils.getCoder((PCollection) output.getValue());

      WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();

      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);
      Vertex vertex =
          dagBuilder.addVertex(
              vertexId,
              WindowGroupP.supplier(
                  context.getOptions(),
                  inputValueCoder,
                  inputCoder,
                  outputCoder,
                  windowingStrategy,
                  vertexId));

      dagBuilder.registerEdgeEndPoint(Utils.getTupleTagId(input), vertex);

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class CreateViewTranslator<T>
      implements JetTransformTranslator<PTransform<PCollection<T>, PCollection<T>>> {

    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      PCollectionView<T> view;
      try {
        view =
            CreatePCollectionViewTranslation.getView(
                (AppliedPTransform<
                        PCollection<T>, PCollection<T>, PTransform<PCollection<T>, PCollection<T>>>)
                    appliedTransform);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }

      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);
      PCollection<T> input = Utils.getInput(appliedTransform);
      Coder inputCoder = Utils.getCoder(input);
      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder = Utils.getCoder((PCollection) output.getValue());

      Vertex vertex =
          dagBuilder.addVertex(
              vertexId,
              ViewP.supplier(inputCoder, outputCoder, input.getWindowingStrategy(), vertexId));

      dagBuilder.registerEdgeEndPoint(Utils.getTupleTagId(input), vertex);

      String viewTag = Utils.getTupleTagId(view);
      dagBuilder.registerCollectionOfEdge(viewTag, view.getTagInternal().getId());
      dagBuilder.registerEdgeStartPoint(viewTag, vertex, outputCoder);

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class FlattenTranslator<T>
      implements JetTransformTranslator<PTransform<PCollectionList<T>, PCollection<T>>> {

    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      Collection<PValue> mainInputs = Utils.getMainInputs(pipeline, node);
      Map<String, Coder> inputCoders =
          Utils.getCoders(
              Utils.getInputs(appliedTransform), e -> Utils.getTupleTagId(e.getValue()));
      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder = Utils.getCoder((PCollection) output.getValue());

      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(appliedTransform.getFullName());
      FlattenP.Supplier processorSupplier =
          new FlattenP.Supplier(inputCoders, outputCoder, vertexId);
      Vertex vertex = dagBuilder.addVertex(vertexId, processorSupplier);
      dagBuilder.registerConstructionListeners(processorSupplier);

      for (PValue value : mainInputs) {
        PCollection<T> input = (PCollection<T>) value;
        dagBuilder.registerEdgeEndPoint(Utils.getTupleTagId(input), vertex);
      }

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class WindowTranslator<T>
      implements JetTransformTranslator<PTransform<PCollection<T>, PCollection<T>>> {
    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      WindowingStrategy<T, BoundedWindow> windowingStrategy =
          (WindowingStrategy<T, BoundedWindow>)
              ((PCollection) Utils.getOutput(appliedTransform).getValue()).getWindowingStrategy();

      PCollection<WindowedValue> input = Utils.getInput(appliedTransform);
      Coder inputCoder = Utils.getCoder(input);
      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder =
          Utils.getCoder((PCollection) Utils.getOutput(appliedTransform).getValue());

      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);

      Vertex vertex =
          dagBuilder.addVertex(
              vertexId,
              AssignWindowP.supplier(inputCoder, outputCoder, windowingStrategy, vertexId));
      dagBuilder.registerEdgeEndPoint(Utils.getTupleTagId(input), vertex);

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class ImpulseTranslator
      implements JetTransformTranslator<PTransform<PBegin, PCollection<byte[]>>> {
    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);

      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder = Utils.getCoder((PCollection) output.getValue());
      Vertex vertex = dagBuilder.addVertex(vertexId, ImpulseP.supplier(vertexId));

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }
  }

  private static class TestStreamTranslator<T>
      implements JetTransformTranslator<PTransform<PBegin, PCollection<T>>> {
    @Override
    public Vertex translate(
        Pipeline pipeline,
        AppliedPTransform<?, ?, ?> appliedTransform,
        Node node,
        JetTranslationContext context) {
      String transformName = appliedTransform.getFullName();
      DAGBuilder dagBuilder = context.getDagBuilder();
      String vertexId = dagBuilder.newVertexId(transformName);

      TestStream<T> testStream = (TestStream<T>) appliedTransform.getTransform();

      // events in the transform are not serializable, we have to translate them. We'll also flatten
      // the collection.
      Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
      Coder outputCoder = Utils.getCoder((PCollection) output.getValue());
      TestStream.TestStreamCoder<T> payloadCoder =
          TestStream.TestStreamCoder.of(testStream.getValueCoder());
      byte[] encodedPayload = getEncodedPayload(testStream, payloadCoder);
      Vertex vertex =
          dagBuilder.addVertex(
              vertexId, TestStreamP.supplier(encodedPayload, payloadCoder, outputCoder));

      String outputEdgeId = Utils.getTupleTagId(output.getValue());
      dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
      dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
      return vertex;
    }

    private static <T> byte[] getEncodedPayload(
        TestStream<T> testStream, TestStream.TestStreamCoder<T> coder) {
      try {
        return CoderUtils.encodeToByteArray(coder, testStream);
      } catch (CoderException e) {
        throw ExceptionUtil.rethrow(e);
      }
    }
  }
}