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

import java.io.{BufferedOutputStream, BufferedReader, FileOutputStream, FileReader}

import scala.collection.mutable
import scala.collection.mutable.ListBuffer

import org.apache.spark.sql.catalyst.expressions.Cast

import org.apache.comet.CometConf.COMET_ONHEAP_MEMORY_OVERHEAD
import org.apache.comet.expressions.{CometCast, CometEvalMode}
import org.apache.comet.serde.{CodegenDispatchFallback, CometAggregateExpressionSerde, CometExpressionSerde, Compatible, Incompatible, QueryPlanSerde, Unsupported}

/**
 * Utility for generating markdown documentation from the configs.
 *
 * This is invoked when running `mvn clean package -DskipTests`.
 */
object GenerateDocs {

  private val publicConfigs: Set[ConfigEntry[_]] = CometConf.allConfs.filter(_.isPublic).toSet

  /**
   * Documentation notes for a single expression.
   *
   * @param name
   *   expression class simple name
   * @param compatibleNotes
   *   differences from Spark that are always present
   * @param incompatibleReasons
   *   reasons the native implementation is incompatible with Spark
   * @param unsupportedReasons
   *   cases that Comet does not support
   * @param codegenDispatch
   *   whether the serde opts into codegen-dispatch fallback for its incompatible cases, so that
   *   the expression stays native and Spark-compatible by default instead of falling back to
   *   Spark
   */
  private case class ExprNotes(
      name: String,
      compatibleNotes: Seq[String],
      incompatibleReasons: Seq[String],
      unsupportedReasons: Seq[String],
      codegenDispatch: Boolean)

  private type CategoryNotes = Seq[ExprNotes]

  /** Build the documentation notes for a single expression serde. */
  private def exprNotes(cls: Class[_], serde: CometExpressionSerde[_]): ExprNotes =
    ExprNotes(
      cls.getSimpleName,
      serde.getCompatibleNotes(),
      serde.getIncompatibleReasons(),
      serde.getUnsupportedReasons(),
      serde.isInstanceOf[CodegenDispatchFallback])

  /** Build the documentation notes for a single aggregate expression serde. */
  private def aggExprNotes(cls: Class[_], serde: CometAggregateExpressionSerde[_]): ExprNotes =
    ExprNotes(
      cls.getSimpleName,
      serde.getCompatibleNotes(),
      serde.getIncompatibleReasons(),
      serde.getUnsupportedReasons(),
      // Aggregate serdes do not participate in codegen-dispatch fallback.
      codegenDispatch = false)

  /**
   * Mapping from expression category to the compatibility guide filename where that category's
   * auto-generated notes should be written, along with a function that produces the notes for
   * that category from the serde maps in `QueryPlanSerde`. Filenames are resolved relative to the
   * per-Spark-version compatibility/expressions directory.
   */
  private def categoryPages: Map[String, (String, () => CategoryNotes)] = Map(
    "array" -> ("array.md",
    () =>
      QueryPlanSerde.arrayExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "datetime" -> ("datetime.md",
    () =>
      QueryPlanSerde.temporalExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "math" -> ("math.md",
    () =>
      QueryPlanSerde.mathExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "struct" -> ("struct.md",
    () =>
      QueryPlanSerde.structExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "aggregate" -> ("aggregate.md",
    () =>
      QueryPlanSerde.aggrSerdeMap.toSeq.map { case (cls, serde) =>
        aggExprNotes(cls, serde)
      }),
    "string" -> ("string.md",
    () =>
      QueryPlanSerde.stringExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "map" -> ("map.md",
    () =>
      QueryPlanSerde.mapExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "misc" -> ("misc.md",
    () =>
      QueryPlanSerde.miscExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }),
    "url" -> ("url.md",
    () =>
      QueryPlanSerde.urlExpressions.toSeq.map { case (cls, serde) =>
        exprNotes(cls, serde)
      }))

  /**
   * Args:
   *   - args(0): user guide root directory (e.g. `docs/source/user-guide/latest/`).
   *   - args(1) (optional): per-Spark-version subdirectory for compatibility pages (e.g.
   *     `spark-3.4`). When omitted, compat pages are written to `compatibility/expressions/`
   *     directly, preserving the legacy flat layout used by released-version doc trees.
   */
  def main(args: Array[String]): Unit = {
    val userGuideLocation = args(0)
    val compatPagesDir = if (args.length > 1) {
      s"$userGuideLocation/compatibility/expressions/${args(1)}"
    } else {
      s"$userGuideLocation/compatibility/expressions"
    }
    generateConfigReference(s"$userGuideLocation/configs.md")
    generateCompatibilityGuide(s"$compatPagesDir/cast.md")
    for ((category, (filename, notesFn)) <- categoryPages) {
      generateExpressionCompatNotes(s"$compatPagesDir/$filename", category, notesFn())
    }
  }

  private def generateConfigReference(filename: String): Unit = {
    val pattern = "<!--BEGIN:CONFIG_TABLE\\[(.*)]-->".r
    val lines = readFile(filename)
    val w = new BufferedOutputStream(new FileOutputStream(filename))
    for (line <- lines) {
      w.write(s"${line.stripTrailing()}\n".getBytes)
      line match {
        case pattern(category) =>
          w.write("<!-- prettier-ignore-start -->\n".getBytes)
          w.write("| Config | Description | Default Value |\n".getBytes)
          w.write("|--------|-------------|---------------|\n".getBytes)
          category match {
            case "enable_expr" =>
              for (expr <- QueryPlanSerde.exprSerdeMap.keys.map(_.getSimpleName).toList.sorted) {
                val config = s"spark.comet.expression.$expr.enabled"
                w.write(
                  s"| `$config` | Enable Comet acceleration for `$expr` | true |\n".getBytes)
              }
              w.write("<!-- prettier-ignore-end -->\n".getBytes)
            case "enable_agg_expr" =>
              for (expr <- QueryPlanSerde.aggrSerdeMap.keys.map(_.getSimpleName).toList.sorted) {
                val config = s"spark.comet.expression.$expr.enabled"
                w.write(
                  s"| `$config` | Enable Comet acceleration for `$expr` | true |\n".getBytes)
              }
              w.write("<!-- prettier-ignore-end -->\n".getBytes)
            case _ =>
              val urlPattern = """Comet\s+(Compatibility|Tuning|Tracing)\s+Guide\s+\(""".r
              val confs = publicConfigs.filter(_.category == category).toList.sortBy(_.key)
              for (conf <- confs) {
                // convert links to Markdown
                val doc =
                  urlPattern.replaceAllIn(conf.doc.trim, m => s"[Comet ${m.group(1)} Guide](")
                // append env var info if present
                val docWithEnvVar = conf.envVar match {
                  case Some(envVarName) =>
                    s"$doc It can be overridden by the environment variable `$envVarName`."
                  case None => doc
                }
                if (conf.defaultValue.isEmpty) {
                  w.write(s"| `${conf.key}` | $docWithEnvVar | |\n".getBytes)
                } else {
                  val isBytesConf = conf.key == COMET_ONHEAP_MEMORY_OVERHEAD.key
                  if (isBytesConf) {
                    val bytes = conf.defaultValue.get.asInstanceOf[Long]
                    w.write(s"| `${conf.key}` | $docWithEnvVar | $bytes MiB |\n".getBytes)
                  } else {
                    val defaultVal = conf.defaultValueString
                    w.write(s"| `${conf.key}` | $docWithEnvVar | $defaultVal |\n".getBytes)
                  }
                }
              }
              w.write("<!-- prettier-ignore-end -->\n".getBytes)
          }
        case _ =>
      }
    }
    w.close()
  }

  private def generateCompatibilityGuide(filename: String): Unit = {
    val lines = readFile(filename)
    val w = new BufferedOutputStream(new FileOutputStream(filename))
    for (line <- lines) {
      w.write(s"${line.stripTrailing()}\n".getBytes)
      if (line.trim == "<!--BEGIN:CAST_LEGACY_TABLE-->") {
        writeCastMatrixForMode(w, CometEvalMode.LEGACY)
      } else if (line.trim == "<!--BEGIN:CAST_TRY_TABLE-->") {
        writeCastMatrixForMode(w, CometEvalMode.TRY)
      } else if (line.trim == "<!--BEGIN:CAST_ANSI_TABLE-->") {
        writeCastMatrixForMode(w, CometEvalMode.ANSI)
      }
    }
    w.close()
  }

  private def generateExpressionCompatNotes(
      filename: String,
      category: String,
      notes: CategoryNotes): Unit = {
    val beginTag = s"<!--BEGIN:EXPR_COMPAT[$category]-->"
    val lines = readFile(filename)
    val w = new BufferedOutputStream(new FileOutputStream(filename))
    for (line <- lines) {
      w.write(s"${line.stripTrailing()}\n".getBytes)
      if (line.trim == beginTag) {
        writeExpressionCompatNotes(w, notes)
      }
    }
    w.close()
  }

  private def writeExpressionCompatNotes(w: BufferedOutputStream, notes: CategoryNotes): Unit = {
    val sorted = notes.sortBy(_.name).filter { n =>
      n.compatibleNotes.nonEmpty || n.incompatibleReasons.nonEmpty || n.unsupportedReasons.nonEmpty
    }
    for (n <- sorted) {
      val name = n.name
      w.write(s"\n## $name\n".getBytes)
      if (n.compatibleNotes.nonEmpty) {
        w.write(
          ("\nThe following differences from Spark are always present and do not require" +
            " any additional configuration:\n\n").getBytes)
        for (note <- n.compatibleNotes) {
          w.write(s"- $note\n".getBytes)
        }
      }
      if (n.incompatibleReasons.nonEmpty) {
        val header = if (n.codegenDispatch) {
          s"\nBy default, Comet accelerates `$name` using JVM codegen dispatch, which runs" +
            " Spark's generated code inside Comet's native pipeline and matches Spark exactly." +
            s" Set `spark.comet.expression.$name.allowIncompatible=true` to use Comet's faster" +
            " native implementation instead, which has the following differences from Spark:\n\n"
        } else {
          s"\nThe following incompatibilities cause `$name` to fall back to Spark by default." +
            s" Set `spark.comet.expression.$name.allowIncompatible=true` to enable Comet" +
            " acceleration despite these differences.\n\n"
        }
        w.write(header.getBytes)
        for (reason <- n.incompatibleReasons) {
          w.write(s"- $reason\n".getBytes)
        }
      }
      if (n.unsupportedReasons.nonEmpty) {
        w.write("\nThe following cases are not supported by Comet:\n\n".getBytes)
        for (reason <- n.unsupportedReasons) {
          w.write(s"- $reason\n".getBytes)
        }
      }
    }
  }

  private def writeCastMatrixForMode(w: BufferedOutputStream, mode: CometEvalMode.Value): Unit = {
    val sortedTypes = CometCast.supportedTypes.sortBy(_.typeName)
    val typeNames = sortedTypes.map(_.typeName.replace("(10,2)", ""))

    // Collect annotations for meaningful notes
    val annotations = mutable.ListBuffer[(String, String, String)]()

    w.write("<!-- prettier-ignore-start -->\n".getBytes)

    // Write header row
    w.write("| |".getBytes)
    for (toTypeName <- typeNames) {
      w.write(s" $toTypeName |".getBytes)
    }
    w.write("\n".getBytes)

    // Write separator row
    w.write("|---|".getBytes)
    for (_ <- typeNames) {
      w.write("---|".getBytes)
    }
    w.write("\n".getBytes)

    // Write data rows
    for ((fromType, fromTypeName) <- sortedTypes.zip(typeNames)) {
      w.write(s"| $fromTypeName |".getBytes)
      for ((toType, toTypeName) <- sortedTypes.zip(typeNames)) {
        val cell = if (fromType == toType) {
          "-"
        } else if (!Cast.canCast(fromType, toType)) {
          "N/A"
        } else {
          val supportLevel = CometCast.isSupported(fromType, toType, None, mode)
          supportLevel match {
            case Compatible(notes) =>
              notes.filter(_.trim.nonEmpty).foreach { note =>
                annotations += ((fromTypeName, toTypeName, note.trim.replace("(10,2)", "")))
              }
              "C"
            case Incompatible(notes) =>
              notes.filter(_.trim.nonEmpty).foreach { note =>
                annotations += ((fromTypeName, toTypeName, note.trim.replace("(10,2)", "")))
              }
              "I"
            case Unsupported(_) =>
              "U"
          }
        }
        w.write(s" $cell |".getBytes)
      }
      w.write("\n".getBytes)
    }

    // Write annotations if any
    if (annotations.nonEmpty) {
      w.write("\n**Notes:**\n".getBytes)
      for ((from, to, note) <- annotations.distinct) {
        w.write(s"- **$from -> $to**: $note\n".getBytes)
      }
    }

    w.write("<!-- prettier-ignore-end -->\n".getBytes)
  }

  /** Read file into memory */
  private def readFile(filename: String): Seq[String] = {
    val r = new BufferedReader(new FileReader(filename))
    val buffer = new ListBuffer[String]()
    var line = r.readLine()
    var skipping = false
    while (line != null) {
      if (line.startsWith("<!--BEGIN:")) {
        buffer += line
        skipping = true
      } else if (line.startsWith("<!--END:")) {
        buffer += line
        skipping = false
      } else if (!skipping) {
        buffer += line
      }
      line = r.readLine()
    }
    r.close()
    buffer.toSeq
  }
}
