/*
 * 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.connect.client

import java.io.SequenceInputStream
import java.lang.ref.Cleaner
import java.util.Objects

import scala.collection.mutable
import scala.jdk.CollectionConverters._
import scala.util.{Failure, Success, Try}

import com.google.protobuf.ByteString
import org.apache.arrow.memory.BufferAllocator
import org.apache.arrow.vector.ipc.message.{ArrowMessage, ArrowRecordBatch}
import org.apache.arrow.vector.types.pojo

import org.apache.spark.connect.proto
import org.apache.spark.connect.proto.ExecutePlanResponse.ObservedMetrics
import org.apache.spark.internal.Logging
import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, RowEncoder}
import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{ProductEncoder, UnboundRowEncoder}
import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
import org.apache.spark.sql.connect.client.arrow.ArrowDeserializingIterator
import org.apache.spark.sql.connect.common.{DataTypeProtoConverter, LiteralValueProtoConverter}
import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.sql.util.{AbstractMessageIterator, ArrowUtils, CloseableIterator, ConcatenatingArrowStreamReader, MessageIterator}

private[sql] class SparkResult[T](
    responses: CloseableIterator[proto.ExecutePlanResponse],
    allocator: BufferAllocator,
    encoder: AgnosticEncoder[T],
    timeZoneId: String)
    extends AutoCloseable
    with Logging { self =>

  case class StageInfo(
      stageId: Long,
      numTasks: Long,
      completedTasks: Long = 0,
      inputBytesRead: Long = 0,
      completed: Boolean = false)

  object StageInfo {
    def apply(stageInfo: proto.ExecutePlanResponse.ExecutionProgress.StageInfo): StageInfo = {
      StageInfo(
        stageInfo.getStageId,
        stageInfo.getNumTasks,
        stageInfo.getNumCompletedTasks,
        stageInfo.getInputBytesRead,
        stageInfo.getDone)
    }
  }

  object Progress {
    def apply(progress: proto.ExecutePlanResponse.ExecutionProgress): Progress = {
      Progress(
        progress.getStagesList.asScala.map(StageInfo(_)).toSeq,
        progress.getNumInflightTasks)
    }
  }

  /**
   * Progress of the query execution. This information can be accessed from the iterator.
   */
  case class Progress(stages: Seq[StageInfo], inflight: Long)

  var progress: Progress = new Progress(Seq.empty, 0)
  private[this] var opId: String = _
  private[this] var numRecords: Int = 0
  private[this] var structType: StructType = _
  private[this] var arrowSchema: pojo.Schema = _
  private[this] var nextResultIndex: Int = 0
  private val resultMap = mutable.Map.empty[Int, (Long, Seq[ArrowMessage])]
  private val observedMetrics = mutable.Map.empty[String, Try[Row]]
  private val cleanable =
    SparkResult.cleaner.register(this, new SparkResultCloseable(resultMap, responses))

  /**
   * Update RowEncoder and recursively update the fields of the ProductEncoder if found.
   */
  private def createEncoder[E](
      enc: AgnosticEncoder[E],
      dataType: DataType): AgnosticEncoder[E] = {
    enc match {
      case UnboundRowEncoder =>
        // Replace the row encoder with the encoder inferred from the schema.
        RowEncoder
          .encoderFor(dataType.asInstanceOf[StructType])
          .asInstanceOf[AgnosticEncoder[E]]
      case ProductEncoder(clsTag, fields, outer) if ProductEncoder.isTuple(clsTag) =>
        // Recursively continue updating the tuple product encoder
        val schema = dataType.asInstanceOf[StructType]
        assert(fields.length <= schema.fields.length)
        val updatedFields = fields.zipWithIndex.map { case (f, id) =>
          f.copy(enc = createEncoder(f.enc, schema.fields(id).dataType))
        }
        ProductEncoder(clsTag, updatedFields, outer)
      case _ =>
        enc
    }
  }

  private def processResponses(
      stopOnOperationId: Boolean = false,
      stopOnSchema: Boolean = false,
      stopOnArrowSchema: Boolean = false,
      stopOnFirstNonEmptyResponse: Boolean = false): Boolean = {
    var nonEmpty = false
    var stop = false
    val arrowBatchChunksToAssemble = mutable.Buffer.empty[ByteString]
    while (!stop && responses.hasNext) {
      val response = responses.next()

      // Collect **all** metrics for this response, whether or not registered to an Observation
      // object.
      observedMetrics ++= processObservedMetrics(response.getObservedMetricsList)

      // Save and validate operationId
      if (opId == null) {
        opId = response.getOperationId
      }
      if (opId != response.getOperationId) {
        // backwards compatibility:
        // response from an old server without operationId field would have getOperationId == "".
        throw new IllegalStateException(
          "Received response with wrong operationId. " +
            s"Expected '$opId' but received '${response.getOperationId}'.")
      }
      stop |= stopOnOperationId

      // Update the execution status. This information can now be accessed directly from
      // the iterator.
      if (response.hasExecutionProgress) {
        progress = Progress(response.getExecutionProgress)
      }

      if (response.hasSchema) {
        // The original schema should arrive before ArrowBatches.
        structType =
          DataTypeProtoConverter.toCatalystType(response.getSchema).asInstanceOf[StructType]
        stop |= stopOnSchema
      }
      if (response.hasArrowBatch) {
        val arrowBatch = response.getArrowBatch
        logDebug(
          s"Received arrow batch rows=${arrowBatch.getRowCount} " +
            s"Number of chunks in batch=${arrowBatch.getNumChunksInBatch} " +
            s"Chunk index=${arrowBatch.getChunkIndex} " +
            s"size=${arrowBatch.getData.size()}")

        if (arrowBatchChunksToAssemble.nonEmpty) {
          // Expect next chunk of the same batch
          if (arrowBatch.getChunkIndex != arrowBatchChunksToAssemble.size) {
            throw new IllegalStateException(
              s"Expected chunk index ${arrowBatchChunksToAssemble.size} of the " +
                s"arrow batch but got ${arrowBatch.getChunkIndex}.")
          }
        } else {
          // Expect next batch
          if (arrowBatch.hasStartOffset) {
            val expectedStartOffset = arrowBatch.getStartOffset
            if (numRecords != expectedStartOffset) {
              throw new IllegalStateException(
                s"Expected arrow batch to start at row offset $numRecords in results, " +
                  s"but received arrow batch starting at offset $expectedStartOffset.")
            }
          }
          if (arrowBatch.getChunkIndex != 0) {
            throw new IllegalStateException(
              s"Expected chunk index 0 of the next arrow batch " +
                s"but got ${arrowBatch.getChunkIndex}.")
          }
        }

        arrowBatchChunksToAssemble += arrowBatch.getData

        // Assemble the chunks to an arrow batch to process if
        // (a) chunking is not enabled (numChunksInBatch is not set or is 0,
        //     in this case, it is the single chunk in the batch)
        // (b) or the client has received all chunks of the batch.
        if (!arrowBatch.hasNumChunksInBatch ||
          arrowBatch.getNumChunksInBatch == 0 ||
          arrowBatchChunksToAssemble.size == arrowBatch.getNumChunksInBatch) {

          val numChunks = arrowBatchChunksToAssemble.size
          val inputStreams =
            arrowBatchChunksToAssemble.map(_.newInput()).iterator.asJavaEnumeration
          val input = new SequenceInputStream(inputStreams)
          arrowBatchChunksToAssemble.clear()
          logDebug(s"Assembling arrow batch from $numChunks chunks.")

          val expectedNumRows = arrowBatch.getRowCount
          val reader = new MessageIterator(input, allocator)
          if (arrowSchema == null) {
            arrowSchema = reader.schema
            stop |= stopOnArrowSchema
          } else if (arrowSchema != reader.schema) {
            throw new IllegalStateException(
              s"""Schema Mismatch between expected and received schema:
                 |=== Expected Schema ===
                 |$arrowSchema
                 |=== Received Schema ===
                 |${reader.schema}
                 |""".stripMargin)
          }
          if (structType == null) {
            // If the schema is not available yet, fallback to the arrow schema.
            structType = ArrowUtils.fromArrowSchema(reader.schema)
          }

          var numRecordsInBatch = 0
          val messages = Seq.newBuilder[ArrowMessage]
          while (reader.hasNext) {
            val message = reader.next()
            message match {
              case batch: ArrowRecordBatch =>
                numRecordsInBatch += batch.getLength
              case _ =>
            }
            messages += message
          }
          if (numRecordsInBatch != expectedNumRows) {
            throw new IllegalStateException(
              s"Expected $expectedNumRows rows in arrow batch but got $numRecordsInBatch.")
          }
          val messagesInBatch = messages.result()
          // Skip the entire result if it is empty.
          if (numRecordsInBatch > 0) {
            numRecords += numRecordsInBatch
            resultMap.put(nextResultIndex, (reader.bytesRead, messagesInBatch))
            nextResultIndex += 1
            nonEmpty |= true
            stop |= stopOnFirstNonEmptyResponse
          } else {
            messagesInBatch.foreach(_.close())
          }
        }
      }
    }
    nonEmpty
  }

  private def processObservedMetrics(
      metrics: java.util.List[ObservedMetrics]): Iterable[(String, Try[Row])] = {
    metrics.asScala.map { metric =>
      metric.getName -> SparkResult.transformObservedMetrics(metric)
    }
  }

  /**
   * Returns the number of elements in the result.
   */
  def length: Int = {
    // We need to process all responses to make sure numRecords is correct.
    processResponses()
    numRecords
  }

  /**
   * @return
   *   the schema of the result.
   */
  def schema: StructType = {
    if (structType == null) {
      processResponses(stopOnSchema = true)
    }
    structType
  }

  /**
   * @return
   *   the operationId of the result.
   */
  def operationId: String = {
    if (opId == null) {
      processResponses(stopOnOperationId = true)
    }
    opId
  }

  /**
   * Create an Array with the contents of the result.
   */
  def toArray: Array[T] = {
    val result = encoder.clsTag.newArray(length)
    val rows = iterator
    try {
      var i = 0
      while (rows.hasNext) {
        result(i) = rows.next()
        assert(i < numRecords)
        i += 1
      }
    } finally {
      rows.close()
    }
    result
  }

  /**
   * Returns all observed metrics in the result.
   */
  def getObservedMetrics: Map[String, Row] = {
    // We need to process all responses to get all metrics.
    processResponses()
    observedMetrics.view.mapValues(_.get).toMap
  }

  /**
   * Returns an iterator over the contents of the result.
   */
  def iterator: CloseableIterator[T] =
    buildIterator(destructive = false)

  /**
   * Returns an destructive iterator over the contents of the result.
   */
  def destructiveIterator: CloseableIterator[T] =
    buildIterator(destructive = true)

  private def buildIterator(destructive: Boolean): CloseableIterator[T] = {
    new CloseableIterator[T] {
      private[this] var iter: CloseableIterator[T] = _

      private def initialize(): Unit = {
        if (iter == null) {
          iter = new ArrowDeserializingIterator(
            createEncoder(encoder, schema),
            new ConcatenatingArrowStreamReader(
              allocator,
              Iterator.single(new ResultMessageIterator(destructive)),
              destructive),
            timeZoneId)
        }
      }

      override def hasNext: Boolean = {
        initialize()
        iter.hasNext
      }

      override def next(): T = {
        initialize()
        iter.next()
      }

      override def close(): Unit = {
        if (iter != null) {
          iter.close()
        }
      }
    }
  }

  /**
   * Close this result, freeing any underlying resources.
   */
  override def close(): Unit = cleanable.clean()

  private class ResultMessageIterator(destructive: Boolean) extends AbstractMessageIterator {
    private[this] var totalBytesRead = 0L
    private[this] var nextResultIndex = 0
    private[this] var current: Iterator[ArrowMessage] = Iterator.empty

    override def bytesRead: Long = totalBytesRead

    override def schema: pojo.Schema = {
      if (arrowSchema == null) {
        // We need a schema to proceed. Spark Connect will always
        // return a result (with a schema) even if the result is empty.
        processResponses(stopOnArrowSchema = true)
        Objects.requireNonNull(arrowSchema)
      }
      arrowSchema
    }

    override def hasNext: Boolean = {
      if (current.hasNext) {
        return true
      }
      val hasNextResult = if (!resultMap.contains(nextResultIndex)) {
        self.processResponses(stopOnFirstNonEmptyResponse = true)
      } else {
        true
      }
      if (hasNextResult) {
        val Some((sizeInBytes, messages)) = if (destructive) {
          resultMap.remove(nextResultIndex)
        } else {
          resultMap.get(nextResultIndex)
        }
        totalBytesRead += sizeInBytes
        current = messages.iterator
        nextResultIndex += 1
      }
      hasNextResult
    }

    override def next(): ArrowMessage = {
      if (!hasNext) {
        throw new NoSuchElementException()
      }
      current.next()
    }
  }
}

private[sql] object SparkResult {
  private val cleaner: Cleaner = Cleaner.create()

  /** Return value is a Seq of pairs, to preserve the order of values. */
  private[sql] def transformObservedMetrics(metric: ObservedMetrics): Try[Row] = {
    // Check if the metric contains errors
    if (metric.hasRootErrorIdx) {
      Failure(
        GrpcExceptionConverter
          .errorsToThrowable(metric.getRootErrorIdx, metric.getErrorsList.asScala.toSeq))
    } else {
      assert(metric.getKeysCount == metric.getValuesCount)
      var schema = new StructType()
      val values = mutable.ArrayBuilder.make[Any]
      values.sizeHint(metric.getKeysCount)
      (0 until metric.getKeysCount).foreach { i =>
        val key = metric.getKeys(i)
        val value = LiteralValueProtoConverter.toScalaValue(metric.getValues(i))
        schema = schema.add(key, LiteralValueProtoConverter.getDataType(metric.getValues(i)))
        values += value
      }
      Success(new GenericRowWithSchema(values.result(), schema))
    }
  }
}

private[client] class SparkResultCloseable(
    resultMap: mutable.Map[Int, (Long, Seq[ArrowMessage])],
    responses: CloseableIterator[proto.ExecutePlanResponse])
    extends AutoCloseable
    with Runnable {
  override def close(): Unit = {
    resultMap.values.foreach(_._2.foreach(_.close()))
    responses.close()
  }

  override def run(): Unit = {
    close()
  }
}
