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

import scala.annotation.tailrec
import scala.jdk.CollectionConverters._

import org.apache.spark.sql.catalyst.expressions.{And, ArrayAppend, ArrayContains, ArrayExcept, ArrayFilter, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayPosition, ArrayRemove, ArrayRepeat, ArraysOverlap, ArraysZip, ArrayUnion, Attribute, CreateArray, ElementAt, EmptyRow, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Size, SortArray}
import org.apache.spark.sql.catalyst.util.GenericArrayData
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._

import org.apache.comet.CometConf
import org.apache.comet.CometSparkSessionExtensions.withInfo
import org.apache.comet.serde.QueryPlanSerde._
import org.apache.comet.shims.{CometExprShim, CometTypeShim}

object CometArrayRemove
    extends CometExpressionSerde[ArrayRemove]
    with CometExprShim
    with ArraysBase {

  override def convert(
      expr: ArrayRemove,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val inputTypes: Set[DataType] = expr.children.map(_.dataType).toSet
    for (dt <- inputTypes) {
      if (!isTypeSupported(dt)) {
        withInfo(expr, s"data type not supported: $dt")
        return None
      }
    }
    val arrayExprProto = exprToProto(expr.left, inputs, binding)
    val keyExprProto = exprToProto(expr.right, inputs, binding)

    scalarFunctionExprToProto("array_remove_all", arrayExprProto, keyExprProto)
  }
}

object CometArrayAppend extends CometExpressionSerde[ArrayAppend] {

  override def convert(
      expr: ArrayAppend,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val child = expr.children.head
    val elementType = child.dataType.asInstanceOf[ArrayType].elementType

    val arrayExprProto = exprToProto(expr.children.head, inputs, binding)
    val keyExprProto = exprToProto(expr.children(1), inputs, binding)

    // DataFusion's array_append always returns a list with nullable elements,
    // so we must promise ArrayType(elementType, containsNull = true) here even if
    // Spark's expr.dataType has containsNull = false (e.g. for array(1,2,3)).
    val arrayAppendScalarExpr =
      scalarFunctionExprToProtoWithReturnType(
        "array_append",
        ArrayType(elementType, containsNull = true),
        false,
        arrayExprProto,
        keyExprProto)

    val isNotNullExpr = createUnaryExpr(
      expr,
      expr.children.head,
      inputs,
      binding,
      (builder, unaryExpr) => builder.setIsNotNull(unaryExpr))

    val nullLiteralProto = exprToProto(Literal(null, elementType), Seq.empty)

    if (arrayAppendScalarExpr.isDefined && isNotNullExpr.isDefined && nullLiteralProto.isDefined) {
      val caseWhenExpr = ExprOuterClass.CaseWhen
        .newBuilder()
        .addWhen(isNotNullExpr.get)
        .addThen(arrayAppendScalarExpr.get)
        .setElseExpr(nullLiteralProto.get)
        .build()
      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setCaseWhen(caseWhenExpr)
          .build())
    } else {
      withInfo(expr, expr.children: _*)
      None
    }
  }
}

object CometArrayContains extends CometExpressionSerde[ArrayContains] {

  override def convert(
      expr: ArrayContains,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExprProto = exprToProto(expr.children.head, inputs, binding)
    val keyExprProto = exprToProto(expr.children(1), inputs, binding)

    scalarFunctionExprToProto("array_contains", arrayExprProto, keyExprProto)
  }
}

object CometSortArray extends CometExpressionSerde[SortArray] {

  override def getIncompatibleReasons(): Seq[String] = Seq(
    "When `" + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key + "=true`, sorting on" +
      " floating-point types is not 100% compatible with Spark")

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Nested arrays with `Struct` or `Null` child values are not supported natively and will" +
      " fall back to Spark.")

  private def supportedSortArrayElementType(
      dt: DataType,
      nestedInArray: Boolean = false): Boolean = {
    dt match {
      // DataFusion's array_sort compares nested arrays through Arrow's rank kernel.
      // That kernel does not support Struct or Null child values,
      // so array<array<struct<...>>> and array<array<null>> would fail at runtime.
      case _: NullType if !nestedInArray =>
        true
      case ArrayType(elementType, _) =>
        supportedSortArrayElementType(elementType, nestedInArray = true)
      case StructType(fields) if !nestedInArray =>
        fields.forall(f => supportedSortArrayElementType(f.dataType))
      case _ =>
        supportedScalarSortElementType(dt)
    }
  }

  override def getSupportLevel(expr: SortArray): SupportLevel = {
    val elementType = expr.base.dataType.asInstanceOf[ArrayType].elementType

    if (!supportedSortArrayElementType(elementType)) {
      Unsupported(Some(s"Sort on array element type $elementType is not supported"))
    } else if (CometConf.COMET_EXEC_STRICT_FLOATING_POINT.get() &&
      SupportLevel.containsFloatingPoint(elementType)) {
      Incompatible(
        Some(
          "Sorting on floating-point is not 100% compatible with Spark, and Comet is running " +
            s"with ${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true. " +
            s"${CometConf.COMPAT_GUIDE}"))
    } else {
      Compatible()
    }
  }

  override def convert(
      expr: SortArray,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExprProto = exprToProtoInternal(expr.base, inputs, binding)
    val (sortDirectionExprProto, nullOrderingExprProto) = expr.ascendingOrder match {
      case Literal(value: Boolean, BooleanType) =>
        val direction = if (value) "ASC" else "DESC"
        val nullOrdering = if (value) "NULLS FIRST" else "NULLS LAST"
        (
          exprToProtoInternal(Literal(direction), inputs, binding),
          exprToProtoInternal(Literal(nullOrdering), inputs, binding))
      case other =>
        withInfo(expr, s"ascendingOrder must be a boolean literal: $other")
        (None, None)
    }

    val sortArrayScalarExpr =
      scalarFunctionExprToProto(
        "array_sort",
        arrayExprProto,
        sortDirectionExprProto,
        nullOrderingExprProto)
    optExprWithInfo(sortArrayScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayIntersect extends CometExpressionSerde[ArrayIntersect] with CometTypeShim {

  private val incompatReason: String =
    "Result array element order may differ from Spark when the right array is longer " +
      "than the left (DataFusion probes the longer side)."

  private val unsupportedCollationReason: String =
    "array_intersect on collated strings is not supported."

  override def getIncompatibleReasons(): Seq[String] = Seq(incompatReason)

  override def getUnsupportedReasons(): Seq[String] = Seq(unsupportedCollationReason)

  override def getSupportLevel(expr: ArrayIntersect): SupportLevel = {
    if (hasNonDefaultStringCollation(expr.dataType)) {
      Unsupported(Some(unsupportedCollationReason))
    } else {
      Incompatible(Some(incompatReason))
    }
  }

  override def convert(
      expr: ArrayIntersect,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val leftArrayExprProto = exprToProto(expr.children.head, inputs, binding)
    val rightArrayExprProto = exprToProto(expr.children(1), inputs, binding)

    val arraysIntersectScalarExpr =
      scalarFunctionExprToProto("array_intersect", leftArrayExprProto, rightArrayExprProto)
    optExprWithInfo(arraysIntersectScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayMax extends CometExpressionSerde[ArrayMax] {
  override def convert(
      expr: ArrayMax,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExprProto = exprToProto(expr.children.head, inputs, binding)

    val arrayMaxScalarExpr =
      scalarFunctionExprToProto("array_max", arrayExprProto)
    optExprWithInfo(arrayMaxScalarExpr, expr)
  }
}

object CometArrayMin extends CometExpressionSerde[ArrayMin] {
  override def convert(
      expr: ArrayMin,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExprProto = exprToProto(expr.children.head, inputs, binding)

    val arrayMinScalarExpr = scalarFunctionExprToProto("array_min", arrayExprProto)
    optExprWithInfo(arrayMinScalarExpr, expr)
  }
}

object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] {
  override def convert(
      expr: ArraysOverlap,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val leftArrayExprProto = exprToProto(expr.left, inputs, binding)
    val rightArrayExprProto = exprToProto(expr.right, inputs, binding)

    val arraysOverlapScalarExpr = scalarFunctionExprToProtoWithReturnType(
      "spark_arrays_overlap",
      BooleanType,
      false,
      leftArrayExprProto,
      rightArrayExprProto)
    optExprWithInfo(arraysOverlapScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayRepeat extends CometExpressionSerde[ArrayRepeat] {

  override def convert(
      expr: ArrayRepeat,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val elementProto = exprToProto(expr.left, inputs, binding)
    val countProto = exprToProto(expr.right, inputs, binding)
    val returnType = ArrayType(elementType = expr.left.dataType)
    for {
      countIsNotNullExpr <- countIsNotNullExpr(expr, inputs, binding)
      arrayRepeatExprProto <- scalarFunctionExprToProto("array_repeat", elementProto, countProto)
      nullLiteralExprProto <- exprToProtoInternal(Literal(null, returnType), inputs, binding)
    } yield {
      val caseWhenProto = ExprOuterClass.CaseWhen
        .newBuilder()
        .addWhen(countIsNotNullExpr)
        .addThen(arrayRepeatExprProto)
        .setElseExpr(nullLiteralExprProto)
        .build()
      ExprOuterClass.Expr
        .newBuilder()
        .setCaseWhen(caseWhenProto)
        .build()
    }
  }

  private def countIsNotNullExpr(
      expr: ArrayRepeat,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    createUnaryExpr(
      expr,
      expr.right,
      inputs,
      binding,
      (builder, countExpr) => builder.setIsNotNull(countExpr))
  }
}

object CometArrayCompact extends CometExpressionSerde[Expression] {

  override def convert(
      expr: Expression,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val child = expr.children.head
    val elementType = child.dataType.asInstanceOf[ArrayType].elementType

    val arrayExprProto = exprToProto(child, inputs, binding)

    // Use Comet's SparkArrayCompact UDF instead of DataFusion's array_remove_all.
    // DF 53 changed array_remove_all to return NULL when the element arg is NULL,
    // which breaks the array_compact use case.
    // TODO: upstream to datafusion-spark crate
    val arrayCompactScalarExpr = scalarFunctionExprToProtoWithReturnType(
      "spark_array_compact",
      ArrayType(elementType = elementType),
      false,
      arrayExprProto)
    optExprWithInfo(arrayCompactScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayExcept extends CometExpressionSerde[ArrayExcept] with CometExprShim {

  override def getIncompatibleReasons(): Seq[String] = Seq(
    "Null handling and ordering may differ from Spark")

  override def getSupportLevel(expr: ArrayExcept): SupportLevel = Incompatible(None)

  @tailrec
  def isTypeSupported(dt: DataType): Boolean = {
    import DataTypes._
    dt match {
      case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
          _: DecimalType | DateType | TimestampType | TimestampNTZType | StringType =>
        true
      case BinaryType => false
      case ArrayType(elementType, _) => isTypeSupported(elementType)
      case _: StructType =>
        false
      case _ => false
    }
  }

  override def convert(
      expr: ArrayExcept,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val inputTypes = expr.children.map(_.dataType).toSet
    for (dt <- inputTypes) {
      if (!isTypeSupported(dt)) {
        withInfo(expr, s"data type not supported: $dt")
        return None
      }
    }
    val leftArrayExprProto = exprToProto(expr.left, inputs, binding)
    val rightArrayExprProto = exprToProto(expr.right, inputs, binding)

    val arrayExceptScalarExpr =
      scalarFunctionExprToProto("array_except", leftArrayExprProto, rightArrayExprProto)
    optExprWithInfo(arrayExceptScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayJoin extends CometExpressionSerde[ArrayJoin] {

  override def getIncompatibleReasons(): Seq[String] = Seq("Null handling may differ from Spark")

  override def getSupportLevel(expr: ArrayJoin): SupportLevel = Incompatible(None)

  override def convert(
      expr: ArrayJoin,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExpr = expr.asInstanceOf[ArrayJoin]
    val arrayExprProto = exprToProto(arrayExpr.array, inputs, binding)
    val delimiterExprProto = exprToProto(arrayExpr.delimiter, inputs, binding)

    arrayExpr.nullReplacement match {
      case Some(nullReplacementExpr) =>
        val nullReplacementExprProto = exprToProto(nullReplacementExpr, inputs, binding)

        val arrayJoinScalarExpr = scalarFunctionExprToProto(
          "array_to_string",
          arrayExprProto,
          delimiterExprProto,
          nullReplacementExprProto)

        optExprWithInfo(
          arrayJoinScalarExpr,
          expr,
          arrayExpr,
          arrayExpr.delimiter,
          nullReplacementExpr)
      case None =>
        val arrayJoinScalarExpr =
          scalarFunctionExprToProto("array_to_string", arrayExprProto, delimiterExprProto)

        optExprWithInfo(arrayJoinScalarExpr, expr, arrayExpr, arrayExpr.delimiter)
    }
  }
}

object CometArrayInsert extends CometExpressionSerde[ArrayInsert] {

  override def getSupportLevel(expr: ArrayInsert): SupportLevel = Compatible()

  override def convert(
      expr: ArrayInsert,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val srcExprProto = exprToProtoInternal(expr.children.head, inputs, binding)
    val posExprProto = exprToProtoInternal(expr.children(1), inputs, binding)
    val itemExprProto = exprToProtoInternal(expr.children(2), inputs, binding)
    val legacyNegativeIndex =
      SQLConf.get.getConfString("spark.sql.legacy.negativeIndexInArrayInsert").toBoolean
    if (srcExprProto.isDefined && posExprProto.isDefined && itemExprProto.isDefined) {
      val arrayInsertBuilder = ExprOuterClass.ArrayInsert
        .newBuilder()
        .setSrcArrayExpr(srcExprProto.get)
        .setPosExpr(posExprProto.get)
        .setItemExpr(itemExprProto.get)
        .setLegacyNegativeIndex(legacyNegativeIndex)

      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setArrayInsert(arrayInsertBuilder)
          .build())
    } else {
      withInfo(
        expr,
        "unsupported arguments for ArrayInsert",
        expr.children.head,
        expr.children(1),
        expr.children(2))
      None
    }
  }
}

object CometArrayUnion extends CometExpressionSerde[ArrayUnion] {
  override def convert(
      expr: ArrayUnion,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val leftArrayExprProto = exprToProto(expr.children.head, inputs, binding)
    val rightArrayExprProto = exprToProto(expr.children(1), inputs, binding)

    val arraysUnionScalarExpr =
      scalarFunctionExprToProto("array_union", leftArrayExprProto, rightArrayExprProto)
    optExprWithInfo(arraysUnionScalarExpr, expr, expr.children: _*)
  }
}

object CometCreateArray extends CometExpressionSerde[CreateArray] {
  override def convert(
      expr: CreateArray,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val children = expr.children

    // Handle empty array: return literal directly to avoid DataFusion coerce_types bug
    // when make_array is called with 0 arguments (issue #3338)
    if (children.isEmpty) {
      val emptyArrayLiteral =
        Literal.create(new GenericArrayData(Array.empty[Any]), expr.dataType)
      return exprToProtoInternal(emptyArrayLiteral, inputs, binding)
    }

    val childExprs = children.map(exprToProtoInternal(_, inputs, binding))

    if (childExprs.forall(_.isDefined)) {
      scalarFunctionExprToProto("make_array", childExprs: _*)
    } else {
      withInfo(expr, "unsupported arguments for CreateArray", children: _*)
      None
    }
  }
}

object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] {

  override def convert(
      expr: GetArrayItem,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val childExpr = exprToProtoInternal(expr.child, inputs, binding)
    val ordinalExpr = exprToProtoInternal(expr.ordinal, inputs, binding)

    if (childExpr.isDefined && ordinalExpr.isDefined) {
      val listExtractBuilder = ExprOuterClass.ListExtract
        .newBuilder()
        .setChild(childExpr.get)
        .setOrdinal(ordinalExpr.get)
        .setOneBased(false)
        .setFailOnError(expr.failOnError)

      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setListExtract(listExtractBuilder)
          .build())
    } else {
      withInfo(expr, "unsupported arguments for GetArrayItem", expr.child, expr.ordinal)
      None
    }
  }
}

object CometArrayReverse extends CometExpressionSerde[Reverse] with ArraysBase {
  val unsupportedReason = "reverse on array containing binary is not supported"

  override def getIncompatibleReasons(): Seq[String] = Seq(unsupportedReason)

  @tailrec
  private def containsBinary(dt: DataType): Boolean = {
    dt match {
      case BinaryType => true
      case ArrayType(elementType, _) => containsBinary(elementType)
      case _ => false
    }
  }

  override def getSupportLevel(expr: Reverse): SupportLevel = {
    if (containsBinary(expr.child.dataType)) {
      Incompatible(Some(unsupportedReason))
    } else {
      Compatible(None)
    }
  }

  override def convert(
      expr: Reverse,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    if (!isTypeSupported(expr.child.dataType)) {
      withInfo(expr, s"child data type not supported: ${expr.child.dataType}")
      return None
    }
    val reverseExprProto = exprToProto(expr.child, inputs, binding)
    val reverseScalarExpr = scalarFunctionExprToProto("array_reverse", reverseExprProto)
    optExprWithInfo(reverseScalarExpr, expr, expr.children: _*)
  }

}

object CometElementAt extends CometExpressionSerde[ElementAt] {

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Input must be an array. `Map` inputs are not supported.")

  override def convert(
      expr: ElementAt,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val childExpr = exprToProtoInternal(expr.left, inputs, binding)
    val ordinalExpr = exprToProtoInternal(expr.right, inputs, binding)
    val defaultExpr = expr.defaultValueOutOfBound.flatMap(exprToProtoInternal(_, inputs, binding))

    if (!expr.left.dataType.isInstanceOf[ArrayType]) {
      withInfo(expr, "Input is not an array")
      return None
    }

    if (childExpr.isDefined && ordinalExpr.isDefined &&
      defaultExpr.isDefined == expr.defaultValueOutOfBound.isDefined) {
      val arrayExtractBuilder = ExprOuterClass.ListExtract
        .newBuilder()
        .setChild(childExpr.get)
        .setOrdinal(ordinalExpr.get)
        .setOneBased(true)
        .setFailOnError(expr.failOnError)

      defaultExpr.foreach(arrayExtractBuilder.setDefaultValue)

      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setListExtract(arrayExtractBuilder)
          .build())
    } else {
      withInfo(expr, "unsupported arguments for ElementAt", expr.left, expr.right)
      None
    }
  }
}

object CometFlatten extends CometExpressionSerde[Flatten] with ArraysBase {

  override def convert(
      expr: Flatten,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val inputTypes = expr.children.map(_.dataType).toSet
    for (dt <- inputTypes) {
      if (!isTypeSupported(dt)) {
        withInfo(expr, s"data type not supported: $dt")
        return None
      }
    }
    val flattenExprProto = exprToProto(expr.child, inputs, binding)
    val flattenScalarExpr = scalarFunctionExprToProto("flatten", flattenExprProto)
    optExprWithInfo(flattenScalarExpr, expr, expr.children: _*)
  }
}

object CometArrayFilter extends CometExpressionSerde[ArrayFilter] {

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Only supports `array_filter` when the function is `IsNotNull` (used by `array_compact`)")

  override def getSupportLevel(expr: ArrayFilter): SupportLevel = {
    expr.function.children.headOption match {
      case Some(_: IsNotNull) => Compatible()
      case _ => Unsupported()
    }
  }

  override def convert(
      expr: ArrayFilter,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    CometArrayCompact.convert(expr, inputs, binding)
  }
}

object CometSize extends CometExpressionSerde[Size] {

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Only supports `ArrayType` input; `MapType` input is not supported")

  override def getSupportLevel(expr: Size): SupportLevel = {
    expr.child.dataType match {
      case _: ArrayType => Compatible()
      case _: MapType => Unsupported(Some("size does not support map inputs"))
      case other =>
        // this should be unreachable because Spark only supports map and array inputs
        Unsupported(Some(s"Unsupported child data type: $other"))
    }
  }

  override def convert(
      expr: Size,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val arrayExprProto = exprToProto(expr.child, inputs, binding)
    for {
      isNotNullExprProto <- createIsNotNullExprProto(expr, inputs, binding)
      sizeScalarExprProto <- scalarFunctionExprToProto("size", arrayExprProto)
      emptyLiteralExprProto <- createLiteralExprProto(SQLConf.get.legacySizeOfNull)
    } yield {
      val caseWhenExpr = ExprOuterClass.CaseWhen
        .newBuilder()
        .addWhen(isNotNullExprProto)
        .addThen(sizeScalarExprProto)
        .setElseExpr(emptyLiteralExprProto)
        .build()
      ExprOuterClass.Expr
        .newBuilder()
        .setCaseWhen(caseWhenExpr)
        .build()
    }
  }

  private def createIsNotNullExprProto(
      expr: Size,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    createUnaryExpr(
      expr,
      expr.child,
      inputs,
      binding,
      (builder, unaryExpr) => builder.setIsNotNull(unaryExpr))
  }

  private def createLiteralExprProto(legacySizeOfNull: Boolean): Option[ExprOuterClass.Expr] = {
    val value = if (legacySizeOfNull) -1 else null
    exprToProto(Literal(value, IntegerType), Seq.empty)
  }

}

object CometArrayPosition extends CometExpressionSerde[ArrayPosition] with ArraysBase {

  override def getSupportLevel(expr: ArrayPosition): SupportLevel = Compatible()

  override def convert(
      expr: ArrayPosition,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    if (expr.children.forall(_.foldable)) {
      withInfo(expr, "all arguments are literals, falling back to Spark")
      return None
    }
    // Check if input types are supported
    val inputTypes: Set[DataType] = expr.children.map(_.dataType).toSet
    for (dt <- inputTypes) {
      if (!isTypeSupported(dt)) {
        withInfo(expr, s"data type not supported: $dt")
        return None
      }
    }

    val arrayExprProto = exprToProto(expr.left, inputs, binding)
    val elementExprProto = exprToProto(expr.right, inputs, binding)

    // Use spark_array_position which returns Int64 and 0 when not found
    // (matching Spark's behavior)
    val optExpr =
      scalarFunctionExprToProto("spark_array_position", arrayExprProto, elementExprProto)
    optExprWithInfo(optExpr, expr, expr.left, expr.right)
  }
}

object CometArraysZip extends CometExpressionSerde[ArraysZip] {

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Not all input data types are supported; falls back to Spark for unsupported types")

  private def isTypeSupported(dt: DataType): Boolean = {
    import DataTypes._
    dt match {
      case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
          _: DecimalType | DateType | TimestampType | TimestampNTZType | StringType | NullType |
          BinaryType =>
        true
      case ArrayType(elementType, _) => isTypeSupported(elementType)
      case StructType(fields) => fields.forall(f => isTypeSupported(f.dataType))
      case _ => false
    }
  }

  override def getSupportLevel(expr: ArraysZip): SupportLevel = {
    val inputTypes = expr.children.map(_.dataType).toSet
    for (dt <- inputTypes) {
      if (!isTypeSupported(dt)) {
        return Unsupported(Some(s"Unsupported child data type: $dt"))
      }
    }
    Compatible()
  }

  override def convert(
      expr: ArraysZip,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {

    val exprChildren: Seq[Option[ExprOuterClass.Expr]] =
      expr.children.map(exprToProtoInternal(_, inputs, binding))
    val names: Seq[Any] = expr.names.map(_.eval(EmptyRow))

    // mimic Spark's ArraysZip behavior: returns NULL if any argument is NULL
    val combinedNullCheck = expr.children.map(child => IsNotNull(child)).reduce(And)
    val isNotNullExpr = exprToProtoInternal(combinedNullCheck, inputs, binding)
    val nullLiteralProto = exprToProto(Literal(null, expr.dataType), Seq.empty)

    if (exprChildren.forall(
        _.isDefined) && isNotNullExpr.isDefined && nullLiteralProto.isDefined) {
      val arraysZip: ExprOuterClass.ArraysZip = ExprOuterClass.ArraysZip
        .newBuilder()
        .addAllValues(exprChildren.map(_.get).asJava)
        .addAllNames(names.map(_.toString).asJava)
        .build()

      val caseWhenExpr = ExprOuterClass.CaseWhen
        .newBuilder()
        .addWhen(isNotNullExpr.get)
        .addThen(ExprOuterClass.Expr.newBuilder().setArraysZip(arraysZip).build())
        .setElseExpr(nullLiteralProto.get)
        .build()
      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setCaseWhen(caseWhenExpr)
          .build())

    } else {
      withInfo(expr, "unsupported arguments for ArraysZip", expr.children ++ expr.names: _*)
      None
    }
  }
}

trait ArraysBase {

  def isTypeSupported(dt: DataType): Boolean = {
    import DataTypes._
    dt match {
      case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
          _: DecimalType | DateType | TimestampType | TimestampNTZType | StringType =>
        true
      case BinaryType => false
      case ArrayType(elementType, _) => isTypeSupported(elementType)
      case _: StructType =>
        // https://github.com/apache/datafusion-comet/issues/1307
        false
      case _ => false
    }
  }
}
