/*
 * 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.comet.fuzz

import scala.util.Random

import org.rogach.scallop.{ScallopConf, Subcommand}
import org.rogach.scallop.ScallopOption

import org.apache.spark.sql.SparkSession

import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}

class Conf(arguments: Seq[String]) extends ScallopConf(arguments) {
  object generateData extends Subcommand("data") {
    val numFiles: ScallopOption[Int] =
      opt[Int](required = true, descr = "Number of files to generate")
    val numRows: ScallopOption[Int] = opt[Int](required = true, descr = "Number of rows per file")
    val randomSeed: ScallopOption[Long] =
      opt[Long](required = false, descr = "Random seed to use")
    val generateArrays: ScallopOption[Boolean] =
      opt[Boolean](required = false, descr = "Whether to generate arrays")
    val generateStructs: ScallopOption[Boolean] =
      opt[Boolean](required = false, descr = "Whether to generate structs")
    val generateMaps: ScallopOption[Boolean] =
      opt[Boolean](required = false, descr = "Whether to generate maps")
    val excludeNegativeZero: ScallopOption[Boolean] =
      opt[Boolean](required = false, descr = "Whether to exclude negative zero")
  }
  addSubcommand(generateData)
  object generateQueries extends Subcommand("queries") {
    val numFiles: ScallopOption[Int] =
      opt[Int](required = true, descr = "Number of input files to use")
    val numQueries: ScallopOption[Int] =
      opt[Int](required = true, descr = "Number of queries to generate")
    val randomSeed: ScallopOption[Long] =
      opt[Long](required = false, descr = "Random seed to use")
  }
  addSubcommand(generateQueries)
  object runQueries extends Subcommand("run") {
    val filename: ScallopOption[String] =
      opt[String](required = true, descr = "File to write queries to")
    val numFiles: ScallopOption[Int] =
      opt[Int](required = true, descr = "Number of input files to use")
    val showFailedSparkQueries: ScallopOption[Boolean] =
      opt[Boolean](required = false, descr = "Whether to show failed Spark queries")
  }
  addSubcommand(runQueries)
  verify()
}

object Main {

  lazy val spark: SparkSession = SparkSession
    .builder()
    .getOrCreate()

  def main(args: Array[String]): Unit = {
    val conf = new Conf(args.toIndexedSeq)
    conf.subcommand match {
      case Some(conf.generateData) =>
        val r = conf.generateData.randomSeed.toOption match {
          case Some(seed) => new Random(seed)
          case None => new Random()
        }
        for (i <- 0 until conf.generateData.numFiles()) {
          ParquetGenerator.makeParquetFile(
            r,
            spark,
            s"test$i.parquet",
            numRows = conf.generateData.numRows(),
            SchemaGenOptions(
              generateArray = conf.generateData.generateArrays(),
              generateStruct = conf.generateData.generateStructs(),
              generateMap = conf.generateData.generateMaps(),
              // create two columns of each primitive type so that they can be used in binary
              // expressions such as `a + b` and `a <  b`
              primitiveTypes = SchemaGenOptions.defaultPrimitiveTypes ++
                SchemaGenOptions.defaultPrimitiveTypes),
            DataGenOptions(
              allowNull = true,
              generateNegativeZero = !conf.generateData.excludeNegativeZero()))
        }
      case Some(conf.generateQueries) =>
        val r = conf.generateQueries.randomSeed.toOption match {
          case Some(seed) => new Random(seed)
          case None => new Random()
        }
        QueryGen.generateRandomQueries(
          r,
          spark,
          numFiles = conf.generateQueries.numFiles(),
          conf.generateQueries.numQueries())
      case Some(conf.runQueries) =>
        QueryRunner.runQueries(
          spark,
          conf.runQueries.numFiles(),
          conf.runQueries.filename(),
          conf.runQueries.showFailedSparkQueries())
      case _ =>
        // scalastyle:off println
        println("Invalid subcommand")
        // scalastyle:on println
        sys.exit(-1)
    }
  }
}
