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

import static org.apache.commons.lang3.ArrayUtils.EMPTY_STRING_ARRAY;
import static org.apache.commons.lang3.StringUtils.substringBetween;
import static org.apache.commons.lang3.math.NumberUtils.toInt;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.serializers.JavaSerializer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import javax.annotation.Nullable;
import org.apache.beam.repackaged.core.org.apache.commons.lang3.ArrayUtils;
import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.SparkStructuredStreamingPipelineOptions;
import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues;
import org.apache.beam.sdk.coders.BigDecimalCoder;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.BigEndianLongCoder;
import org.apache.beam.sdk.coders.BigEndianShortCoder;
import org.apache.beam.sdk.coders.BigIntegerCoder;
import org.apache.beam.sdk.coders.BitSetCoder;
import org.apache.beam.sdk.coders.BooleanCoder;
import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.ByteCoder;
import org.apache.beam.sdk.coders.CollectionCoder;
import org.apache.beam.sdk.coders.DelegateCoder;
import org.apache.beam.sdk.coders.DequeCoder;
import org.apache.beam.sdk.coders.DoubleCoder;
import org.apache.beam.sdk.coders.DurationCoder;
import org.apache.beam.sdk.coders.FloatCoder;
import org.apache.beam.sdk.coders.InstantCoder;
import org.apache.beam.sdk.coders.IterableCoder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.LengthPrefixCoder;
import org.apache.beam.sdk.coders.ListCoder;
import org.apache.beam.sdk.coders.MapCoder;
import org.apache.beam.sdk.coders.NullableCoder;
import org.apache.beam.sdk.coders.RowCoder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.SetCoder;
import org.apache.beam.sdk.coders.ShardedKeyCoder;
import org.apache.beam.sdk.coders.SnappyCoder;
import org.apache.beam.sdk.coders.SortedMapCoder;
import org.apache.beam.sdk.coders.StringDelegateCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.TextualIntegerCoder;
import org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.coders.VarLongCoder;
import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.io.FileBasedSink;
import org.apache.beam.sdk.io.range.OffsetRange;
import org.apache.beam.sdk.transforms.join.CoGbkResult;
import org.apache.beam.sdk.transforms.join.CoGbkResultSchema;
import org.apache.beam.sdk.transforms.join.UnionCoder;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.util.construction.resources.PipelineResources;
import org.apache.beam.sdk.values.PCollectionViews;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Collections2;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.apache.spark.SparkConf;
import org.apache.spark.serializer.KryoRegistrator;
import org.apache.spark.serializer.KryoSerializer;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.execution.datasources.v2.DataWritingSparkTaskResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SparkSessionFactory {

  private static final Logger LOG = LoggerFactory.getLogger(SparkSessionFactory.class);

  // Patterns to exclude local JRE and certain artifact (groups) in Maven and Gradle cache.
  private static final Collection<String> SPARK_JAR_EXCLUDES =
      Lists.newArrayList(
          "jre/lib/ext/",
          "/org/slf4j/",
          "/org.slf4j/",
          "/log4j/",
          "/io/dropwizard/metrics/",
          "/io.dropwizard.metrics/",
          "/org/apache/spark/",
          "/org.apache.spark/",
          "/org/apache/hadoop/",
          "/org.apache.hadoop/",
          "/org/scala-lang/",
          "/org.scala-lang/",
          "/com.esotericsoftware/kryo-shaded",
          "/com/esotericsoftware/kryo-shaded");

  /**
   * Gets active {@link SparkSession} or creates one using {@link
   * SparkStructuredStreamingPipelineOptions}.
   */
  public static SparkSession getOrCreateSession(SparkStructuredStreamingPipelineOptions options) {
    if (options.getUseActiveSparkSession()) {
      return SparkSession.active();
    }
    return sessionBuilder(options.getSparkMaster(), options).getOrCreate();
  }

  /** Creates Spark session builder with some optimizations for local mode, e.g. in tests. */
  public static SparkSession.Builder sessionBuilder(String master) {
    return sessionBuilder(master, null);
  }

  private static SparkSession.Builder sessionBuilder(
      String master, @Nullable SparkStructuredStreamingPipelineOptions options) {

    SparkConf sparkConf = new SparkConf().setIfMissing("spark.master", master);
    master = sparkConf.get("spark.master"); // use effective master in the remainder of this method

    if (options != null) {
      if (options.getAppName() != null) {
        sparkConf.setAppName(options.getAppName());
      }

      if (options.getFilesToStage() != null && !options.getFilesToStage().isEmpty()) {
        // Append the files to stage provided by the user to `spark.jars`.
        PipelineResources.prepareFilesForStaging(options);
        String[] filesToStage = filterFilesToStage(options, Collections.emptyList());
        String[] jars = getSparkJars(sparkConf);
        sparkConf.setJars(jars.length > 0 ? ArrayUtils.addAll(jars, filesToStage) : filesToStage);
      } else if (!sparkConf.contains("spark.jars") && !master.startsWith("local[")) {
        // Stage classpath if `spark.jars` not set and not in local mode.
        PipelineResources.prepareFilesForStaging(options);
        // Set `spark.jars`, exclude JRE libs and jars causing conflicts using `userClassPathFirst`.
        sparkConf.setJars(filterFilesToStage(options, SPARK_JAR_EXCLUDES));
        // Enable `userClassPathFirst` to prevent issues with guava, jackson and others.
        sparkConf.setIfMissing("spark.executor.userClassPathFirst", "true");
      }
    }

    // Set to 'org.apache.spark.serializer.JavaSerializer' via system property to disable Kryo
    String serializer = sparkConf.get("spark.serializer", KryoSerializer.class.getName());
    if (serializer.equals(KryoSerializer.class.getName())) {
      // Set to 'false' via system property to disable usage of Kryo unsafe
      boolean unsafe = sparkConf.getBoolean("spark.kryo.unsafe", true);
      sparkConf.set("spark.serializer", serializer);
      sparkConf.set("spark.kryo.unsafe", Boolean.toString(unsafe));
      sparkConf.set("spark.kryo.registrator", SparkKryoRegistrator.class.getName());
      LOG.info("Configured `spark.serializer` to use KryoSerializer [unsafe={}]", unsafe);
    }

    // By default, Spark defines 200 as a number of sql partitions. This seems too much for local
    // mode, so try to align with value of "sparkMaster" option in this case.
    // We should not overwrite this value (or any user-defined spark configuration value) if the
    // user has already configured it.
    int partitions = localNumPartitions(master);
    if (partitions > 0) {
      sparkConf.setIfMissing("spark.sql.shuffle.partitions", Integer.toString(partitions));
    }

    return SparkSession.builder().config(sparkConf);
  }

  @SuppressWarnings({"return", "toarray.nullable.elements", "methodref.receiver"}) // safe to ignore
  private static String[] filterFilesToStage(
      SparkStructuredStreamingPipelineOptions opts, Collection<String> excludes) {
    Collection<String> files = opts.getFilesToStage();
    if (files == null || files.isEmpty()) {
      return EMPTY_STRING_ARRAY;
    }
    if (!excludes.isEmpty()) {
      files = Collections2.filter(files, f -> !excludes.stream().anyMatch(f::contains));
    }
    return files.toArray(EMPTY_STRING_ARRAY);
  }

  private static String[] getSparkJars(SparkConf conf) {
    return conf.contains("spark.jars") ? conf.get("spark.jars").split(",") : EMPTY_STRING_ARRAY;
  }

  private static int localNumPartitions(String master) {
    return master.startsWith("local[") ? toInt(substringBetween(master, "local[", "]")) : 0;
  }

  /**
   * {@link KryoRegistrator} for Spark to serialize broadcast variables used for side-inputs.
   *
   * <p>Note, this registrator must be public to be accessible for Kryo.
   *
   * @see SideInputValues
   */
  public static class SparkKryoRegistrator implements KryoRegistrator {
    @Override
    public void registerClasses(Kryo kryo) {
      kryo.register(InternalRow.class);
      kryo.register(InternalRow[].class);
      kryo.register(byte[][].class);
      kryo.register(HashMap.class);
      kryo.register(ArrayList.class);

      // support writing noop format
      kryo.register(DataWritingSparkTaskResult.class);

      // TODO find more efficient ways
      kryo.register(SerializablePipelineOptions.class, new JavaSerializer());

      // side input values (spark runner specific)
      kryo.register(SideInputValues.ByWindow.class);
      kryo.register(SideInputValues.Global.class);

      // avro coders
      tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroCoder");
      tryToRegister(kryo, "org.apache.beam.sdk.extensions.avro.coders.AvroGenericCoder");

      // standard coders of org.apache.beam.sdk.coders
      kryo.register(BigDecimalCoder.class);
      kryo.register(BigEndianIntegerCoder.class);
      kryo.register(BigEndianLongCoder.class);
      kryo.register(BigEndianShortCoder.class);
      kryo.register(BigIntegerCoder.class);
      kryo.register(BitSetCoder.class);
      kryo.register(BooleanCoder.class);
      kryo.register(ByteArrayCoder.class);
      kryo.register(ByteCoder.class);
      kryo.register(CollectionCoder.class);
      kryo.register(DelegateCoder.class);
      kryo.register(DequeCoder.class);
      kryo.register(DoubleCoder.class);
      kryo.register(DurationCoder.class);
      kryo.register(FloatCoder.class);
      kryo.register(InstantCoder.class);
      kryo.register(IterableCoder.class);
      kryo.register(KvCoder.class);
      kryo.register(LengthPrefixCoder.class);
      kryo.register(ListCoder.class);
      kryo.register(MapCoder.class);
      kryo.register(NullableCoder.class);
      kryo.register(RowCoder.class);
      kryo.register(SerializableCoder.class);
      kryo.register(SetCoder.class);
      kryo.register(ShardedKeyCoder.class);
      kryo.register(SnappyCoder.class);
      kryo.register(SortedMapCoder.class);
      kryo.register(StringDelegateCoder.class);
      kryo.register(StringUtf8Coder.class);
      kryo.register(TextualIntegerCoder.class);
      kryo.register(TimestampPrefixingWindowCoder.class);
      kryo.register(VarIntCoder.class);
      kryo.register(VarLongCoder.class);
      kryo.register(VoidCoder.class);

      // bounded windows and windowed value coders
      kryo.register(GlobalWindow.Coder.class);
      kryo.register(IntervalWindow.IntervalWindowCoder.class);
      kryo.register(WindowedValues.FullWindowedValueCoder.class);
      kryo.register(WindowedValues.ParamWindowedValueCoder.class);
      kryo.register(WindowedValues.ValueOnlyWindowedValueCoder.class);

      // various others
      kryo.register(OffsetRange.Coder.class);
      kryo.register(UnionCoder.class);
      kryo.register(PCollectionViews.ValueOrMetadataCoder.class);
      kryo.register(FileBasedSink.FileResultCoder.class);
      kryo.register(CoGbkResult.CoGbkResultCoder.class);
      kryo.register(CoGbkResultSchema.class);
      kryo.register(TupleTag.class);
      kryo.register(TupleTagList.class);
    }

    private void tryToRegister(Kryo kryo, String className) {
      try {
        kryo.register(Class.forName(className));
      } catch (ClassNotFoundException e) {
        LOG.info("Class {}} was not found on classpath", className);
      }
    }
  }
}
