/*
 * 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.serde.operator

import scala.collection.mutable.ListBuffer
import scala.jdk.CollectionConverters._

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions.{Expression, Literal}
import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefaultValues
import org.apache.spark.sql.comet.{CometNativeExec, CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SubqueryAdaptiveBroadcastExec}
import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
import org.apache.spark.sql.internal.SQLConf

import org.apache.comet.{CometConf, ConfigEntry}
import org.apache.comet.CometConf.COMET_EXEC_ENABLED
import org.apache.comet.CometSparkSessionExtensions.{hasExplainInfo, isSpark35Plus, isSpark41Plus, withInfo}
import org.apache.comet.objectstore.NativeConfig
import org.apache.comet.parquet.CometParquetUtils
import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel}
import org.apache.comet.serde.ExprOuterClass.Expr
import org.apache.comet.serde.OperatorOuterClass.Operator
import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType}

/**
 * Validation and serde logic for `native_datafusion` scans.
 */
object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging {

  /** Determine whether the scan is supported and tag the Spark plan with any fallback reasons */
  def isSupported(scanExec: FileSourceScanExec): Boolean = {

    if (hasExplainInfo(scanExec)) {
      // this node has already been tagged with fallback reasons
      return false
    }

    if (!COMET_EXEC_ENABLED.get()) {
      withInfo(scanExec, s"Full native scan disabled because ${COMET_EXEC_ENABLED.key} disabled")
    }

    // AQE DPP (SubqueryAdaptiveBroadcastExec) is converted to CometSubqueryBroadcastExec
    // by CometPlanAdaptiveDynamicPruningFilters (queryStageOptimizerRule, Spark 3.5+).
    // Non-AQE DPP (SubqueryBroadcastExec/SubqueryExec) is converted by
    // CometExecRule.convertSubqueryBroadcasts. Both are resolved through the lazy
    // partition serialization path in CometNativeScanExec.
    //
    // On Spark 3.4, injectQueryStageOptimizerRule is unavailable, so the AQE DPP conversion
    // rule can't run. CometScanRule.transformV1Scan rejects AQE DPP on 3.4, so this check
    // is a safety net: if the scan somehow reached here with AQE DPP on 3.4, reject it.
    if (!isSpark35Plus && scanExec.partitionFilters.exists(isAqeDynamicPruningFilter)) {
      withInfo(scanExec, "Native DataFusion scan does not support AQE DPP on Spark 3.4")
    }

    if (SQLConf.get.ignoreCorruptFiles ||
      scanExec.relation.options
        .get("ignorecorruptfiles") // Spark sets this to lowercase.
        .contains("true")) {
      withInfo(scanExec, "Full native scan disabled because ignoreCorruptFiles enabled")
    }

    if (SQLConf.get.ignoreMissingFiles ||
      scanExec.relation.options
        .get("ignoremissingfiles") // Spark sets this to lowercase.
        .contains("true")) {

      withInfo(scanExec, "Full native scan disabled because ignoreMissingFiles enabled")
    }

    // the scan is supported if no fallback reasons were added to the node
    !hasExplainInfo(scanExec)
  }

  /** Detects AQE DPP (SubqueryAdaptiveBroadcastExec), as opposed to non-AQE DPP. */
  private def isAqeDynamicPruningFilter(e: Expression): Boolean =
    e.exists {
      case sub: InSubqueryExec => sub.plan.isInstanceOf[SubqueryAdaptiveBroadcastExec]
      case _ => false
    }

  override def enabledConfig: Option[ConfigEntry[Boolean]] = None

  override def getSupportLevel(operator: CometScanExec): SupportLevel = {
    // all checks happen in CometScanRule before ScanExec is converted to CometScanExec, so
    // we always report compatible here because this serde object is for the converted CometScanExec
    Compatible()
  }

  override def convert(
      scan: CometScanExec,
      builder: Operator.Builder,
      childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
    val nativeScanBuilder = OperatorOuterClass.NativeScan.newBuilder()
    val commonBuilder = OperatorOuterClass.NativeScanCommon.newBuilder()

    // Set source in common (used as part of injection key)
    commonBuilder.setSource(scan.simpleStringWithNodeId())

    val scanTypes = scan.output.flatten { attr =>
      serializeDataType(attr.dataType)
    }

    if (scanTypes.length == scan.output.length) {
      commonBuilder.addAllFields(scanTypes.asJava)

      // Sink operators don't have children
      builder.clearChildren()

      if (scan.conf.getConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED) &&
        CometConf.COMET_RESPECT_PARQUET_FILTER_PUSHDOWN.get(scan.conf)) {

        val dataFilters = new ListBuffer[Expr]()
        for (filter <- scan.supportedDataFilters) {
          exprToProto(filter, scan.output) match {
            case Some(proto) => dataFilters += proto
            case _ =>
              logWarning(s"Unsupported data filter $filter")
          }
        }
        commonBuilder.addAllDataFilters(dataFilters.asJava)
      }

      val possibleDefaultValues = getExistenceDefaultValues(scan.requiredSchema)
      if (possibleDefaultValues.exists(_ != null)) {
        // Our schema has default values. Serialize two lists, one with the default values
        // and another with the indexes in the schema so the native side can map missing
        // columns to these default values.
        val (defaultValues, indexes) = possibleDefaultValues.zipWithIndex
          .filter { case (expr, _) => expr != null }
          .map { case (expr, index) =>
            // ResolveDefaultColumnsUtil.getExistenceDefaultValues has evaluated these
            // expressions and they should now just be literals.
            (Literal(expr), index.toLong.asInstanceOf[java.lang.Long])
          }
          .unzip
        commonBuilder.addAllDefaultValues(
          defaultValues.flatMap(exprToProto(_, scan.output)).toIterable.asJava)
        commonBuilder.addAllDefaultValuesIndexes(indexes.toIterable.asJava)
      }

      // Extract object store options from first file (S3 configs apply to all files in scan).
      // Use selectedPartitions (static) instead of getFilePartitions() because at planning time
      // DPP subqueries haven't been resolved yet. Object store options don't depend on DPP.
      val firstFileUri = scan.selectedPartitions
        .flatMap(_.files.headOption)
        .headOption
        .map(_.getPath.toUri)

      val partitionSchema = schema2Proto(scan.relation.partitionSchema.fields)
      val requiredSchema = schema2Proto(scan.requiredSchema.fields)
      val dataSchema = schema2Proto(scan.relation.dataSchema.fields)

      val dataSchemaIndexes = scan.requiredSchema.fields.map(field => {
        scan.relation.dataSchema.fieldIndex(field.name)
      })
      val partitionSchemaIndexes = Array
        .range(
          scan.relation.dataSchema.fields.length,
          scan.relation.dataSchema.length + scan.relation.partitionSchema.fields.length)

      val projectionVector = (dataSchemaIndexes ++ partitionSchemaIndexes).map(idx =>
        idx.toLong.asInstanceOf[java.lang.Long])

      commonBuilder.addAllProjectionVector(projectionVector.toIterable.asJava)

      // In `CometScanRule`, we ensure partitionSchema is supported.
      assert(partitionSchema.length == scan.relation.partitionSchema.fields.length)

      commonBuilder.addAllDataSchema(dataSchema.toIterable.asJava)
      commonBuilder.addAllRequiredSchema(requiredSchema.toIterable.asJava)
      commonBuilder.addAllPartitionSchema(partitionSchema.toIterable.asJava)
      commonBuilder.setSessionTimezone(scan.conf.getConfString("spark.sql.session.timeZone"))
      commonBuilder.setCaseSensitive(scan.conf.getConf[Boolean](SQLConf.CASE_SENSITIVE))

      // SPARK-53535 (Spark 4.1+): when reading a struct whose requested fields are all
      // missing in the Parquet file, the new default preserves the parent struct's
      // nullness from the file (so non-null parents materialize as a struct of all-null
      // fields). Pre-4.1 Spark hardcodes the legacy behavior (whole struct null), which
      // matches the Comet default we use as fallback.
      val returnNullStructConfKey =
        "spark.sql.legacy.parquet.returnNullStructIfAllFieldsMissing"
      val returnNullStructDefault = if (isSpark41Plus) "false" else "true"
      commonBuilder.setReturnNullStructIfAllFieldsMissing(
        scan.conf.getConfString(returnNullStructConfKey, returnNullStructDefault).toBoolean)

      // Field-ID matching: only ask the native side to do extra work when the conf is on AND
      // the requested schema actually carries IDs. Spark's ParquetReadSupport applies the same
      // gate before invoking matchIdField.
      val useFieldId =
        scan.conf.getConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED) &&
          ParquetUtils.hasFieldIds(scan.requiredSchema)
      commonBuilder.setUseFieldId(useFieldId)
      commonBuilder.setIgnoreMissingFieldId(
        scan.conf.getConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID))

      // Collect S3/cloud storage configurations
      val hadoopConf = scan.relation.sparkSession.sessionState
        .newHadoopConfWithOptions(scan.relation.options)

      commonBuilder.setEncryptionEnabled(CometParquetUtils.encryptionEnabled(hadoopConf))

      firstFileUri.foreach { uri =>
        val objectStoreOptions =
          NativeConfig.extractObjectStoreOptions(hadoopConf, uri)
        objectStoreOptions.foreach { case (key, value) =>
          commonBuilder.putObjectStoreOptions(key, value)
        }
      }

      // Set common data in NativeScan (file_partition will be populated at execution time)
      nativeScanBuilder.setCommon(commonBuilder.build())

      Some(builder.setNativeScan(nativeScanBuilder).build())

    } else {
      // There are unsupported scan type
      withInfo(
        scan,
        s"unsupported Comet operator: ${scan.nodeName}, due to unsupported data types above")
      None
    }

  }

  override def createExec(nativeOp: Operator, op: CometScanExec): CometNativeExec = {
    CometNativeScanExec(nativeOp, op.wrapped, op.session, op)
  }
}
