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

import java.io.File

import org.apache.spark.sql.{CometTestBase, SaveMode}
import org.apache.spark.sql.internal.StaticSQLConf

class CometPluginsSuite extends CometTestBase {
  override protected def sparkConf: SparkConf = {
    val conf = new SparkConf()
    conf.set("spark.driver.memory", "1G")
    conf.set("spark.executor.memory", "1G")
    conf.set("spark.executor.memoryOverhead", "2G")
    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
    conf.set("spark.comet.enabled", "true")
    conf.set("spark.comet.exec.enabled", "true")
    conf.set("spark.comet.exec.onHeap.enabled", "true")
    conf.set("spark.comet.metrics.enabled", "true")
    conf
  }

  test("Register Comet extension") {
    // test common case where no extensions are previously registered
    {
      val conf = new SparkConf()
      CometDriverPlugin.registerCometSessionExtension(conf)
      assert(
        "org.apache.comet.CometSparkSessionExtensions" == conf.get(
          StaticSQLConf.SPARK_SESSION_EXTENSIONS.key))
    }
    // test case where Comet is already registered
    {
      val conf = new SparkConf()
      conf.set(
        StaticSQLConf.SPARK_SESSION_EXTENSIONS.key,
        "org.apache.comet.CometSparkSessionExtensions")
      CometDriverPlugin.registerCometSessionExtension(conf)
      assert(
        "org.apache.comet.CometSparkSessionExtensions" == conf.get(
          StaticSQLConf.SPARK_SESSION_EXTENSIONS.key))
    }
    // test case where other extensions are already registered
    {
      val conf = new SparkConf()
      conf.set(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, "foo,bar")
      CometDriverPlugin.registerCometSessionExtension(conf)
      assert(
        "foo,bar,org.apache.comet.CometSparkSessionExtensions" == conf.get(
          StaticSQLConf.SPARK_SESSION_EXTENSIONS.key))
    }
    // test case where other extensions, including Comet, are already registered
    {
      val conf = new SparkConf()
      conf.set(
        StaticSQLConf.SPARK_SESSION_EXTENSIONS.key,
        "foo,bar,org.apache.comet.CometSparkSessionExtensions")
      CometDriverPlugin.registerCometSessionExtension(conf)
      assert(
        "foo,bar,org.apache.comet.CometSparkSessionExtensions" == conf.get(
          StaticSQLConf.SPARK_SESSION_EXTENSIONS.key))
    }
  }

  test("CometSource metrics are recorded") {
    val nativeBefore = CometSource.NATIVE_OPERATORS.getCount
    val queriesBefore = CometSource.QUERIES_PLANNED.getCount

    withTempPath { dir =>
      val path = new File(dir, "test.parquet").toString
      spark.range(1000).toDF("id").write.mode(SaveMode.Overwrite).parquet(path)
      spark.read.parquet(path).filter("id > 500").collect()
    }
    spark.sparkContext.listenerBus.waitUntilEmpty()
    assert(
      CometSource.QUERIES_PLANNED.getCount > queriesBefore,
      "queries.planned should increment after query")
    assert(
      CometSource.NATIVE_OPERATORS.getCount > nativeBefore,
      "operators.native should increment for native execution")
  }

  test("metrics not double counted with AQE") {
    withSQLConf("spark.sql.adaptive.enabled" -> "true") {
      withTempPath { dir =>
        val path = new File(dir, "test.parquet").toString
        spark.range(10000).toDF("id").write.mode(SaveMode.Overwrite).parquet(path)

        spark.sparkContext.listenerBus.waitUntilEmpty()
        val queriesBefore = CometSource.QUERIES_PLANNED.getCount
        spark.read.parquet(path).filter("id > 100").collect()
        spark.read.parquet(path).filter("id > 200").collect()
        spark.sparkContext.listenerBus.waitUntilEmpty()
        val queriesAfter = CometSource.QUERIES_PLANNED.getCount
        assert(
          queriesAfter == queriesBefore + 2,
          s"Expected 2 queries, got ${queriesAfter - queriesBefore}")
      }
    }
  }

  test("Default Comet memory overhead") {
    val execMemOverhead1 = spark.conf.get("spark.executor.memoryOverhead")
    val execMemOverhead2 = spark.sessionState.conf.getConfString("spark.executor.memoryOverhead")
    val execMemOverhead3 = spark.sparkContext.getConf.get("spark.executor.memoryOverhead")
    val execMemOverhead4 = spark.sparkContext.conf.get("spark.executor.memoryOverhead")

    // 2GB + 384MB (default Comet memory overhead)
    assert(execMemOverhead1 == "3072M")
    assert(execMemOverhead2 == "3072M")
    assert(execMemOverhead3 == "3072M")
    assert(execMemOverhead4 == "3072M")
  }
}

class CometPluginsDefaultSuite extends CometTestBase {
  override protected def sparkConf: SparkConf = {
    val conf = new SparkConf()
    conf.set("spark.driver.memory", "1G")
    conf.set("spark.executor.memory", "1G")
    conf.set("spark.executor.memoryOverheadFactor", "0.5")
    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
    conf.set("spark.comet.enabled", "true")
    conf.set("spark.comet.exec.shuffle.enabled", "true")
    conf.set("spark.comet.exec.onHeap.enabled", "true")
    conf
  }

  test("Default executor memory overhead + Comet memory overhead") {
    val execMemOverhead1 = spark.conf.get("spark.executor.memoryOverhead")
    val execMemOverhead2 = spark.sessionState.conf.getConfString("spark.executor.memoryOverhead")
    val execMemOverhead3 = spark.sparkContext.getConf.get("spark.executor.memoryOverhead")
    val execMemOverhead4 = spark.sparkContext.conf.get("spark.executor.memoryOverhead")

    // Spark executor memory overhead = executor memory (1G) * memoryOverheadFactor (0.5) = 512MB
    // 512MB + 384MB (default Comet memory overhead)
    assert(execMemOverhead1 == "1536M")
    assert(execMemOverhead2 == "1536M")
    assert(execMemOverhead3 == "1536M")
    assert(execMemOverhead4 == "1536M")
  }
}

class CometPluginsNonOverrideSuite extends CometTestBase {
  override protected def sparkConf: SparkConf = {
    val conf = new SparkConf()
    conf.set("spark.driver.memory", "1G")
    conf.set("spark.executor.memory", "1G")
    conf.set("spark.executor.memoryOverhead", "2G")
    conf.set("spark.executor.memoryOverheadFactor", "0.5")
    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
    conf.set("spark.comet.enabled", "true")
    conf.set("spark.comet.exec.shuffle.enabled", "false")
    conf.set("spark.comet.exec.enabled", "false")
    conf.set("spark.comet.exec.onHeap.enabled", "true")
    conf
  }

  test("executor memory overhead is not overridden") {
    val execMemOverhead1 = spark.conf.get("spark.executor.memoryOverhead")
    val execMemOverhead2 = spark.sessionState.conf.getConfString("spark.executor.memoryOverhead")
    val execMemOverhead3 = spark.sparkContext.getConf.get("spark.executor.memoryOverhead")
    val execMemOverhead4 = spark.sparkContext.conf.get("spark.executor.memoryOverhead")

    assert(execMemOverhead1 == "2G")
    assert(execMemOverhead2 == "2G")
    assert(execMemOverhead3 == "2G")
    assert(execMemOverhead4 == "2G")
  }
}

class CometPluginsUnifiedModeOverrideSuite extends CometTestBase {
  override protected def sparkConf: SparkConf = {
    val conf = new SparkConf()
    conf.set("spark.driver.memory", "1G")
    conf.set("spark.executor.memory", "1G")
    conf.set("spark.executor.memoryOverhead", "1G")
    conf.set("spark.plugins", "org.apache.spark.CometPlugin")
    conf.set("spark.comet.enabled", "true")
    conf.set("spark.memory.offHeap.enabled", "true")
    conf.set("spark.memory.offHeap.size", "2G")
    conf.set("spark.comet.exec.shuffle.enabled", "true")
    conf.set("spark.comet.exec.enabled", "true")
    conf.set("spark.comet.memory.overhead.factor", "0.5")
    conf
  }

  /*
   * Since using unified memory executor memory should not be overridden
   */
  test("executor memory overhead is not overridden") {
    val execMemOverhead1 = spark.conf.get("spark.executor.memoryOverhead")
    val execMemOverhead2 = spark.sessionState.conf.getConfString("spark.executor.memoryOverhead")
    val execMemOverhead3 = spark.sparkContext.getConf.get("spark.executor.memoryOverhead")
    val execMemOverhead4 = spark.sparkContext.conf.get("spark.executor.memoryOverhead")

    // in unified memory mode, comet memory overhead is
    // spark.memory.offHeap.size (2G) * spark.comet.memory.overhead.factor (0.5) = 1G  and the overhead is not overridden
    assert(execMemOverhead1 == "1G")
    assert(execMemOverhead2 == "1G")
    assert(execMemOverhead3 == "1G")
    assert(execMemOverhead4 == "1G")
  }
}
