/*
 * 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.spark.sql.comet.util

import java.io.{DataInputStream, DataOutputStream, File}
import java.nio.ByteBuffer
import java.nio.channels.Channels

import scala.jdk.CollectionConverters._

import org.apache.arrow.c.CDataDictionaryProvider
import org.apache.arrow.vector._
import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.dictionary.DictionaryProvider
import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter}
import org.apache.arrow.vector.types._
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
import org.apache.arrow.vector.util.VectorSchemaRootAppender
import org.apache.spark.{SparkEnv, SparkException}
import org.apache.spark.internal.Logging
import org.apache.spark.io.CompressionCodec
import org.apache.spark.sql.comet.execution.arrow.ArrowReaderIterator
import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream}

import org.apache.comet.Constants.COMET_CONF_DIR_ENV
import org.apache.comet.shims.CometTypeShim
import org.apache.comet.vector.CometVector

object Utils extends CometTypeShim with Logging {
  def getConfPath(confFileName: String): String = {
    sys.env
      .get(COMET_CONF_DIR_ENV)
      .map { t => new File(s"$t${File.separator}$confFileName") }
      .filter(_.isFile)
      .map(_.getAbsolutePath)
      .orNull
  }

  def stringToSeq(str: String): Seq[String] = {
    str.split(",").map(_.trim()).filter(_.nonEmpty)
  }

  /** bridges the function call to Spark's Util */
  def getSimpleName(cls: Class[_]): String = {
    org.apache.spark.util.Utils.getSimpleName(cls)
  }

  def fromArrowField(field: Field): DataType = {
    field.getType match {
      case _: ArrowType.Map =>
        val elementField = field.getChildren.get(0)
        val keyType = fromArrowField(elementField.getChildren.get(0))
        val valueType = fromArrowField(elementField.getChildren.get(1))
        MapType(keyType, valueType, elementField.getChildren.get(1).isNullable)
      case ArrowType.List.INSTANCE =>
        val elementField = field.getChildren().get(0)
        val elementType = fromArrowField(elementField)
        ArrayType(elementType, containsNull = elementField.isNullable)
      case ArrowType.Struct.INSTANCE =>
        val fields = field.getChildren().asScala.map { child =>
          val dt = fromArrowField(child)
          StructField(child.getName, dt, child.isNullable)
        }
        StructType(fields.toSeq)
      case arrowType => fromArrowType(arrowType)
    }
  }

  def fromArrowType(dt: ArrowType): DataType = dt match {
    case ArrowType.Bool.INSTANCE => BooleanType
    case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 => ByteType
    case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 2 => ShortType
    case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 4 => IntegerType
    case int: ArrowType.Int if int.getIsSigned && int.getBitWidth == 8 * 8 => LongType
    case float: ArrowType.FloatingPoint if float.getPrecision == FloatingPointPrecision.SINGLE =>
      FloatType
    case float: ArrowType.FloatingPoint if float.getPrecision == FloatingPointPrecision.DOUBLE =>
      DoubleType
    case ArrowType.Utf8.INSTANCE => StringType
    case ArrowType.Binary.INSTANCE => BinaryType
    case _: ArrowType.FixedSizeBinary => BinaryType
    case d: ArrowType.Decimal => DecimalType(d.getPrecision, d.getScale)
    case date: ArrowType.Date if date.getUnit == DateUnit.DAY => DateType
    case ts: ArrowType.Timestamp
        if ts.getUnit == TimeUnit.MICROSECOND && ts.getTimezone == null =>
      TimestampNTZType
    case ts: ArrowType.Timestamp if ts.getUnit == TimeUnit.MICROSECOND => TimestampType
    case ArrowType.Null.INSTANCE => NullType
    case yi: ArrowType.Interval if yi.getUnit == IntervalUnit.YEAR_MONTH =>
      YearMonthIntervalType()
    case di: ArrowType.Interval if di.getUnit == IntervalUnit.DAY_TIME => DayTimeIntervalType()
    case _ => throw new UnsupportedOperationException(s"Unsupported data type: ${dt.toString}")
  }

  def fromArrowSchema(schema: Schema): StructType = {
    StructType(schema.getFields.asScala.map { field =>
      val dt = fromArrowField(field)
      StructField(field.getName, dt, field.isNullable)
    }.toArray)
  }

  /** Maps data type from Spark to Arrow. NOTE: timeZoneId required for TimestampTypes */
  def toArrowType(dt: DataType, timeZoneId: String): ArrowType =
    dt match {
      case BooleanType => ArrowType.Bool.INSTANCE
      case ByteType => new ArrowType.Int(8, true)
      case ShortType => new ArrowType.Int(8 * 2, true)
      case IntegerType => new ArrowType.Int(8 * 4, true)
      case LongType => new ArrowType.Int(8 * 8, true)
      case FloatType => new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
      case DoubleType => new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)
      case _: StringType => ArrowType.Utf8.INSTANCE
      case dt if isStringCollationType(dt) => ArrowType.Utf8.INSTANCE
      case BinaryType => ArrowType.Binary.INSTANCE
      case DecimalType.Fixed(precision, scale) => new ArrowType.Decimal(precision, scale, 128)
      case DateType => new ArrowType.Date(DateUnit.DAY)
      case TimestampType =>
        if (timeZoneId == null) {
          throw new UnsupportedOperationException(
            s"${TimestampType.catalogString} must supply timeZoneId parameter")
        } else {
          new ArrowType.Timestamp(TimeUnit.MICROSECOND, timeZoneId)
        }
      case TimestampNTZType =>
        new ArrowType.Timestamp(TimeUnit.MICROSECOND, null)
      case _ =>
        throw new UnsupportedOperationException(
          s"Unsupported data type: [${dt.getClass.getName}] ${dt.catalogString}")
    }

  /** Maps field from Spark to Arrow. NOTE: timeZoneId required for TimestampType */
  def toArrowField(name: String, dt: DataType, nullable: Boolean, timeZoneId: String): Field = {
    dt match {
      case ArrayType(elementType, containsNull) =>
        val fieldType = new FieldType(nullable, ArrowType.List.INSTANCE, null)
        new Field(
          name,
          fieldType,
          Seq(toArrowField("element", elementType, containsNull, timeZoneId)).asJava)
      case StructType(fields) =>
        val fieldType = new FieldType(nullable, ArrowType.Struct.INSTANCE, null)
        new Field(
          name,
          fieldType,
          fields
            .map { field =>
              toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
            }
            .toSeq
            .asJava)
      case MapType(keyType, valueType, valueContainsNull) =>
        val mapType = new FieldType(nullable, new ArrowType.Map(false), null)
        // Note: Map Type struct can not be null, Struct Type key field can not be null
        new Field(
          name,
          mapType,
          Seq(
            toArrowField(
              MapVector.DATA_VECTOR_NAME,
              new StructType()
                .add(MapVector.KEY_NAME, keyType, nullable = false)
                .add(MapVector.VALUE_NAME, valueType, nullable = valueContainsNull),
              nullable = false,
              timeZoneId)).asJava)
      case dataType =>
        val fieldType = new FieldType(nullable, toArrowType(dataType, timeZoneId), null)
        new Field(name, fieldType, Seq.empty[Field].asJava)
    }
  }

  /**
   * Maps schema from Spark to Arrow. NOTE: timeZoneId required for TimestampType in StructType
   */
  def toArrowSchema(schema: StructType, timeZoneId: String): Schema = {
    new Schema(schema.map { field =>
      toArrowField(field.name, field.dataType, field.nullable, timeZoneId)
    }.asJava)
  }

  /**
   * Serializes a list of `ColumnarBatch` into an output stream. This method must be in `spark`
   * package because `ChunkedByteBufferOutputStream` is spark private class. As it uses Arrow
   * classes, it must be in `common` module.
   *
   * @param batches
   *   the output batches, each batch is a list of Arrow vectors wrapped in `CometVector`
   * @param out
   *   the output stream
   */
  def serializeBatches(batches: Iterator[ColumnarBatch]): Iterator[(Long, ChunkedByteBuffer)] = {
    batches.map { batch =>
      val dictionaryProvider: CDataDictionaryProvider = new CDataDictionaryProvider

      val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
      val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate)
      val out = new DataOutputStream(codec.compressedOutputStream(cbbos))

      val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
      val root = new VectorSchemaRoot(fieldVectors.asJava)
      val provider = batchProviderOpt.getOrElse(dictionaryProvider)

      val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out))
      writer.start()
      writer.writeBatch()
      root.clear()
      writer.close()

      if (out.size() > 0) {
        (batch.numRows().toLong, cbbos.toChunkedByteBuffer)
      } else {
        (batch.numRows().toLong, new ChunkedByteBuffer(Array.empty[ByteBuffer]))
      }
    }
  }

  /**
   * Decodes the byte arrays back to ColumnarBatchs and put them into buffer.
   *
   * @param bytes
   *   the serialized batches
   * @param source
   *   the class that calls this method
   * @return
   *   an iterator of ColumnarBatch
   */
  def decodeBatches(bytes: ChunkedByteBuffer, source: String): Iterator[ColumnarBatch] = {
    if (bytes.size == 0) {
      return Iterator.empty
    }

    // use Spark's compression codec (LZ4 by default) and not Comet's compression
    val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
    val cbbis = bytes.toInputStream()
    val ins = new DataInputStream(codec.compressedInputStream(cbbis))
    // batches are in Arrow IPC format
    new ArrowReaderIterator(Channels.newChannel(ins), source)
  }

  /**
   * Coalesces many small Arrow IPC batches into a single batch for broadcasting.
   *
   * Why this is necessary: The broadcast exchange collects shuffle output by calling
   * getByteArrayRdd, which serializes each ColumnarBatch independently into its own
   * ChunkedByteBuffer. The shuffle reader (CometBlockStoreShuffleReader) produces one
   * ColumnarBatch per shuffle block, and there is one block per writer task per output partition.
   * So with W writer tasks and P output partitions, the broadcast collects up to W * P tiny
   * batches. For example, with 400 writer tasks and 500 partitions, 1M rows would arrive as ~200K
   * batches of ~5 rows each.
   *
   * Without coalescing, every consumer task in the broadcast join would independently deserialize
   * all of these tiny Arrow IPC streams, paying per-stream overhead (schema parsing, buffer
   * allocation) for each one. With coalescing, we decode and append all batches into one
   * VectorSchemaRoot on the driver, then re-serialize once. Each consumer task then deserializes
   * a single Arrow IPC stream.
   */
  def coalesceBroadcastBatches(
      input: Iterator[ChunkedByteBuffer]): (Array[ChunkedByteBuffer], Long, Long) = {
    val buffers = input.filterNot(_.size == 0).toArray
    if (buffers.isEmpty) {
      return (Array.empty, 0L, 0L)
    }

    val allocator = org.apache.comet.CometArrowAllocator
      .newChildAllocator("broadcast-coalesce", 0, Long.MaxValue)
    try {
      var targetRoot: VectorSchemaRoot = null
      var totalRows = 0L
      var batchCount = 0

      val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
      try {
        for (bytes <- buffers) {
          val compressedInputStream =
            new DataInputStream(codec.compressedInputStream(bytes.toInputStream()))
          val reader =
            new ArrowStreamReader(Channels.newChannel(compressedInputStream), allocator)
          try {
            // Comet decodes dictionaries during execution, so this shouldn't happen.
            // If it does, fall back to the original uncoalesced buffers because each
            // partition can have a different dictionary, and appending index vectors
            // would silently mix indices from incompatible dictionaries.
            if (!reader.getDictionaryVectors.isEmpty) {
              logWarning(
                "Unexpected dictionary-encoded column during BroadcastExchange coalescing; " +
                  "skipping coalesce")
              reader.close()
              if (targetRoot != null) {
                targetRoot.close()
                targetRoot = null
              }
              return (buffers, 0L, 0L)
            }
            while (reader.loadNextBatch()) {
              val sourceRoot = reader.getVectorSchemaRoot
              if (targetRoot == null) {
                targetRoot = VectorSchemaRoot.create(sourceRoot.getSchema, allocator)
                targetRoot.allocateNew()
              }
              VectorSchemaRootAppender.append(targetRoot, sourceRoot)
              totalRows += sourceRoot.getRowCount
              batchCount += 1
            }
          } finally {
            reader.close()
          }
        }

        if (targetRoot == null) {
          return (Array.empty, 0L, 0L)
        }

        assert(
          targetRoot.getRowCount.toLong == totalRows,
          s"Row count mismatch after coalesce: ${targetRoot.getRowCount} != $totalRows")

        logInfo(s"Coalesced $batchCount broadcast batches into 1 ($totalRows rows)")

        val outputStream = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate)
        val compressedOutputStream =
          new DataOutputStream(codec.compressedOutputStream(outputStream))
        val writer =
          new ArrowStreamWriter(targetRoot, null, Channels.newChannel(compressedOutputStream))
        try {
          writer.start()
          writer.writeBatch()
        } finally {
          writer.close()
        }

        (Array(outputStream.toChunkedByteBuffer), batchCount.toLong, totalRows)
      } finally {
        if (targetRoot != null) {
          targetRoot.close()
        }
      }
    } finally {
      allocator.close()
    }
  }

  def getBatchFieldVectors(
      batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = {
    var provider: Option[DictionaryProvider] = None
    val fieldVectors = (0 until batch.numCols()).map { index =>
      batch.column(index) match {
        case a: CometVector =>
          val valueVector = a.getValueVector
          if (valueVector.getField.getDictionary != null) {
            if (provider.isEmpty) {
              provider = Some(a.getDictionaryProvider)
            }
          }

          getFieldVector(valueVector, "serialize")

        case c =>
          throw new SparkException(
            s"Comet execution only takes Arrow Arrays, but got ${c.getClass}. " +
              "This typically happens when a Comet scan falls back to Spark due to unsupported " +
              "data types (e.g., complex types like structs, arrays, or maps). " +
              "To resolve this, you can: " +
              "(1) enable spark.comet.scan.allowIncompatible=true to use a compatible native " +
              "scan variant, or " +
              "(2) enable spark.comet.convert.parquet.enabled=true to convert Spark Parquet " +
              "data to Arrow format automatically.")
      }
    }
    (fieldVectors, provider)
  }

  def getFieldVector(valueVector: ValueVector, reason: String): FieldVector = {
    valueVector match {
      case v @ (_: BitVector | _: TinyIntVector | _: SmallIntVector | _: IntVector |
          _: BigIntVector | _: Float4Vector | _: Float8Vector | _: VarCharVector |
          _: DecimalVector | _: DateDayVector | _: TimeStampMicroTZVector | _: VarBinaryVector |
          _: FixedSizeBinaryVector | _: TimeStampMicroVector | _: StructVector | _: ListVector |
          _: MapVector | _: NullVector) =>
        v.asInstanceOf[FieldVector]
      case _ =>
        throw new SparkException(s"Unsupported Arrow Vector for $reason: ${valueVector.getClass}")
    }
  }
}
