/*
 * 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.spark.structuredstreaming.translation.batch;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import org.apache.beam.runners.core.construction.PipelineOptionsSerializationUtils;
import org.apache.beam.runners.core.construction.ReadTranslation;
import org.apache.beam.runners.core.serialization.Base64Serializer;
import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
import org.apache.beam.runners.spark.structuredstreaming.translation.TranslationContext;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

class ReadSourceTranslatorBatch<T>
    implements TransformTranslator<PTransform<PBegin, PCollection<T>>> {

  private static String sourceProviderClass = DatasetSourceBatch.class.getCanonicalName();

  @SuppressWarnings("unchecked")
  @Override
  public void translateTransform(
      PTransform<PBegin, PCollection<T>> transform, TranslationContext context) {
    AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> rootTransform =
        (AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>>)
            context.getCurrentTransform();

        BoundedSource<T> source;
    try {
      source = ReadTranslation.boundedSourceFromTransform(rootTransform);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    SparkSession sparkSession = context.getSparkSession();

    String serializedSource = Base64Serializer.serializeUnchecked(source);
    Map<String, String> datasetSourceOptions = new HashMap<>();
    datasetSourceOptions.put(DatasetSourceBatch.BEAM_SOURCE_OPTION, serializedSource);
    datasetSourceOptions.put(DatasetSourceBatch.DEFAULT_PARALLELISM,
        String.valueOf(context.getSparkSession().sparkContext().defaultParallelism()));
    datasetSourceOptions.put(DatasetSourceBatch.PIPELINE_OPTIONS,
        PipelineOptionsSerializationUtils.serializeToJson(context.getOptions()));
    Dataset<Row> rowDataset = sparkSession.read().format(sourceProviderClass).options(datasetSourceOptions)
        .load();

    MapFunction<Row, WindowedValue> func = new MapFunction<Row, WindowedValue>() {
      @Override public WindowedValue call(Row value) throws Exception {
        //there is only one value put in each Row by the InputPartitionReader
        return value.<WindowedValue>getAs(0);
      }
    };
    //TODO: is there a better way than using the raw WindowedValue? Can an Encoder<WindowedValue<T>>
    // be created ?
    Dataset<WindowedValue> dataset = rowDataset.map(func, Encoders.kryo(WindowedValue.class));

    PCollection<T> output = (PCollection<T>) context.getOutput();
    context.putDatasetRaw(output, dataset);
  }
}

@@ -909,56 +909,8 @@
 on;%0A
-import java.util.HashMap;%0Aimport java.util.Map;%0A
 impo
@@ -2718,96 +2718,97 @@
     
-Map%3CString, String%3E datasetSourceOptions = new HashMap%3C%3E();%0A    datasetSourceOptions.put
+Dataset%3CRow%3E rowDataset = sparkSession.read().format(sourceProviderClass)%0A        .option
 (Dat
@@ -2864,38 +2864,24 @@
 rce)
-;
 %0A    
-datasetSourceOptions.put
+    .option
 (Dat
@@ -2917,16 +2917,20 @@
 LELISM,%0A
+    
         
@@ -3011,38 +3011,24 @@
 ()))
-;%0A    datasetSourceOptions.put
+%0A        .option
 (Dat
@@ -3061,16 +3061,20 @@
 PTIONS,%0A
+    
         
@@ -3149,126 +3149,8 @@
 ()))
-;%0A    Dataset%3CRow%3E rowDataset = sparkSession.read().format(sourceProviderClass).options(datasetSourceOptions)%0A        
 .loa

improve readability of options passing to the source