/*
 * 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.spark.sql.auron

import org.apache.spark.SparkEnv
import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.SparkSessionExtensions
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.ColumnarRule
import org.apache.spark.sql.execution.LocalTableScanExec
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.internal.SQLConf

import org.apache.auron.spark.configuration.SparkAuronConfiguration

class AuronSparkSessionExtension extends (SparkSessionExtensions => Unit) with Logging {
  Shims.get.initExtension()

  override def apply(extensions: SparkSessionExtensions): Unit = {
    SparkEnv.get.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true")
    SparkEnv.get.conf.set(SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY.key, "true")
    logInfo(s"${classOf[AuronSparkSessionExtension].getName} enabled")

    Shims.get.onApplyingExtension()

    extensions.injectColumnar(sparkSession => {
      AuronColumnarOverrides(sparkSession)
    })
  }
}

object AuronSparkSessionExtension extends Logging {
  def dumpSimpleSparkPlanTreeNode(exec: SparkPlan, depth: Int = 0): Unit = {
    val nodeName = exec.nodeName
    val convertible = exec
      .getTagValue(AuronConvertStrategy.convertibleTag)
      .getOrElse(false)
    val strategy =
      exec.getTagValue(AuronConvertStrategy.convertStrategyTag).getOrElse(Default)
    logInfo(s" +${"-" * depth} $nodeName (convertible=$convertible, strategy=$strategy)")
    exec.children.foreach(dumpSimpleSparkPlanTreeNode(_, depth + 1))
  }
}

case class AuronColumnarOverrides(sparkSession: SparkSession) extends ColumnarRule with Logging {
  import AuronSparkSessionExtension._

  override def preColumnarTransitions: Rule[SparkPlan] = {
    new Rule[SparkPlan] {
      override def apply(sparkPlan: SparkPlan): SparkPlan = {
        if (!SparkAuronConfiguration.AURON_ENABLED.get()) {
          return sparkPlan // performs no conversion if auron is not enabled
        }

        if (sparkPlan.isInstanceOf[LocalTableScanExec]) {
          return sparkPlan // skip useless local table scan (generated by set, addjar, etc)
        }

        if (AuronConverters.enableShuffleExechange) {
          assert(
            AuronConverters.supportedShuffleManager,
            "spark.auron.enable.shuffleExchange=true requires an Auron shuffle manager. " +
              "Set spark.shuffle.manager to one of: " +
              "org.apache.spark.sql.execution.auron.shuffle.AuronShuffleManager, " +
              "org.apache.spark.sql.execution.auron.shuffle.AuronUniffleShuffleManager, " +
              "org.apache.spark.sql.execution.auron.shuffle.AuronCelebornShuffleManager.")
        }

        // generate convert strategy
        AuronConvertStrategy.apply(sparkPlan)
        logInfo("Auron convert strategy for current stage:")
        dumpSimpleSparkPlanTreeNode(sparkPlan)

        val sparkPlanTransformed = AuronConverters.convertSparkPlanRecursively(sparkPlan)
        logInfo("Auron convert result for current stage:")
        dumpSimpleSparkPlanTreeNode(sparkPlanTransformed)

        logInfo(s"Transformed spark plan after preColumnarTransitions:\n${sparkPlanTransformed
          .treeString(verbose = true, addSuffix = true)}")

        // post-transform
        Shims.get.postTransform(sparkPlanTransformed, sparkSession.sparkContext)
        sparkPlanTransformed
      }
    }
  }
}
