import com.databricks.spark.sql.perf.tpcds.{TPCDS, TPCDSTables}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._

// ── Configure these before running ──────────────────────────────────────────
// These must match the values used in generate-data.scala
val useHive    = true
val useIceberg = false

val rootDir        = "file:///path/to/tpcds/data"     // same as generate-data.scala
val databaseName   = "tpcds"                           // same as generate-data.scala
val scaleFactor    = "50"                              // same as generate-data.scala
val dsdgenDir      = "/path/to/tpcds-kit/tools"
val resultLocation = "file:///path/to/tpcds/results"  // where per-query timings are written
// ────────────────────────────────────────────────────────────────────────────

val sqlContext = if (useHive) {
  SparkSession.builder()
    .appName("tpcds-benchmark")
    .enableHiveSupport()
    .getOrCreate()
    .sqlContext
} else if (useIceberg) {
  SparkSession.builder()
    .appName("tpcds-benchmark")
    .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
    .config("spark.sql.catalog.spark_catalog.type", "hadoop")
    .config("spark.sql.catalog.spark_catalog.warehouse", "file:///tmp/iceberg_warehouse")
    .getOrCreate()
    .sqlContext
} else {
  throw new RuntimeException("Set either useHive or useIceberg to true")
}

sqlContext.sql(s"use $databaseName")

val tables = new TPCDSTables(sqlContext,
  dsdgenDir           = dsdgenDir,
  scaleFactor         = scaleFactor,
  useDoubleForDecimal = true,
  useStringForDate    = false)

tables.analyzeTables(databaseName, analyzeColumns = false)

val tpcds      = new TPCDS(sqlContext = sqlContext)
val iterations = 1
val queries    = tpcds.tpcds2_4Queries
val timeout    = 24 * 60 * 60

val experiment = tpcds.runExperiment(
  queries,
  iterations     = iterations,
  resultLocation = resultLocation,
  forkThread     = true)

experiment.waitForFinish(timeout)

val result1 = spark.read.json(resultLocation)
  .filter(s"timestamp=${experiment.timestamp}")
  .select(explode($"results").as("r"))
result1.createOrReplaceTempView("result1")

spark.sql(
  "select r.name, r.numRows, " +
  "bround((r.parsingTime+r.analysisTime+r.optimizationTime+r.planningTime+r.executionTime)/1000.0,1) " +
  "as runtime_sec from result1").show(1000)

spark.sql(
  "select sum(bround((r.parsingTime+r.analysisTime+r.optimizationTime+r.planningTime+r.executionTime)/1000.0,1)) " +
  "as total_runtime_sec from result1").show(1000)