/*
 * 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.execution.streaming.sources

import java.util.concurrent.atomic.AtomicInteger
import javax.annotation.concurrent.GuardedBy

import scala.collection.mutable.ListBuffer

import org.json4s.{Formats, NoTypeHints}
import org.json4s.jackson.Serialization

import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.rpc.{
  RpcAddress,
  RpcEndpointRef,
  RpcEnv,
  ThreadSafeRpcEndpoint
}
import org.apache.spark.sql.{Encoder, SparkSession, SQLContext}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
import org.apache.spark.sql.connector.read.InputPartition
import org.apache.spark.sql.connector.read.PartitionReader
import org.apache.spark.sql.connector.read.PartitionReaderFactory
import org.apache.spark.sql.connector.read.streaming.{
  Offset => OffsetV2,
  PartitionOffset,
  ReadLimit,
  SupportsRealTimeMode,
  SupportsRealTimeRead
}
import org.apache.spark.sql.connector.read.streaming.SupportsRealTimeRead.RecordStatus
import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock
import org.apache.spark.sql.execution.streaming.Offset
import org.apache.spark.sql.execution.streaming.runtime._
import org.apache.spark.util.{Clock, RpcUtils}

/**
 * A low latency memory source from memory, only for unit test purpose.
 * This class is very similar to ContinuousMemoryStream, except that it implements the
 * interface of SupportsRealTimeMode, rather than ContinuousStream
 * The overall strategy here is:
 *  * LowLatencyMemoryStream maintains a list of records for each partition. addData() will
 *    distribute records evenly-ish across partitions.
 *  * RecordEndpoint is set up as an endpoint for executor-side
 *    LowLatencyMemoryStreamInputPartitionReader instances to poll. It returns the record at
 *    the specified offset within the list, or null if that offset doesn't yet have a record.
 * This differs from the existing memory source implementation as data is sent once to
 * tasks as part of the Partition/Split metadata at the beginning of a batch.
 */
class LowLatencyMemoryStream[A: Encoder](
    id: Int,
    sparkSession: SparkSession,
    numPartitions: Int = 2,
    clock: Clock = LowLatencyClock.getClock)
    extends MemoryStreamBaseClass[A](0, sparkSession)
    with SupportsRealTimeMode {
  private implicit val formats: Formats = Serialization.formats(NoTypeHints)

  @GuardedBy("this")
  private val records =
    Seq.fill(numPartitions)(new ListBuffer[(UnsafeRow, Long)]())

  @GuardedBy("this")
  private var numRecords: Long = 0L

  private val recordEndpoint = new LowLatencyMemoryStreamEndpoint(records, this)
  @volatile private var endpointRef: RpcEndpointRef = _

  override def addData(data: IterableOnce[A]): Offset = synchronized {
    // Distribute data evenly among partition lists.
    val timestamp = clock.getTimeMillis()
    data.iterator.foreach { item =>
      val partitionId: Int = (numRecords % numPartitions).toInt
      records(partitionId) += ((toRow(item).copy().asInstanceOf[UnsafeRow], timestamp))
      numRecords += 1
    }

    // The new target offset is the offset where all records in all partitions have been processed.
    LowLatencyMemoryStreamOffset((0 until numPartitions).map(i => (i, records(i).size)).toMap)
  }

  def addData(partitionId: Int, data: IterableOnce[A]): Offset = synchronized {
    require(
      partitionId >= 0 && partitionId < numPartitions,
      s"Partition ID $partitionId is out of bounds for $numPartitions partitions."
    )

    // Add data to the specified partition.
    val timestamp = clock.getTimeMillis()
    data.iterator.foreach { item =>
      records(partitionId) += ((toRow(item).copy().asInstanceOf[UnsafeRow], timestamp))
      numRecords += 1
    }

    // The new target offset is the offset where all records in all partitions have been processed.
    LowLatencyMemoryStreamOffset((0 until numPartitions).map(i => (i, records(i).size)).toMap)
  }

  override def initialOffset(): OffsetV2 = {
    LowLatencyMemoryStreamOffset((0 until numPartitions).map(i => (i, 0)).toMap)
  }

  override def latestOffset(startOffset: OffsetV2, limit: ReadLimit): OffsetV2 = synchronized {
    LowLatencyMemoryStreamOffset((0 until numPartitions).map(i => (i, records(i).size)).toMap)
  }

  override def deserializeOffset(json: String): LowLatencyMemoryStreamOffset = {
    LowLatencyMemoryStreamOffset(Serialization.read[Map[Int, Int]](json))
  }

  override def mergeOffsets(offsets: Array[PartitionOffset]): LowLatencyMemoryStreamOffset = {
    LowLatencyMemoryStreamOffset(
      offsets.map {
        case ContinuousRecordPartitionOffset(part, num) => (part, num)
      }.toMap
    )
  }

  override def planInputPartitions(start: OffsetV2): Array[InputPartition] = {
    val startOffset = start.asInstanceOf[LowLatencyMemoryStreamOffset]
    synchronized {
      val endpointName = s"LowLatencyRecordEndpoint-${java.util.UUID.randomUUID()}-$id"
      endpointRef = recordEndpoint.rpcEnv.setupEndpoint(endpointName, recordEndpoint)

      startOffset.partitionNums.map {
        case (part, index) =>
          LowLatencyMemoryStreamInputPartition(
            endpointName,
            endpointRef.address,
            part,
            index,
            Int.MaxValue
          )
      }.toArray
    }
  }

  override def planInputPartitions(start: OffsetV2, end: OffsetV2): Array[InputPartition] = {
    val startOffset = start.asInstanceOf[LowLatencyMemoryStreamOffset]
    val endOffset = end.asInstanceOf[LowLatencyMemoryStreamOffset]
    synchronized {
      val endpointName = s"LowLatencyRecordEndpoint-${java.util.UUID.randomUUID()}-$id"
      endpointRef = recordEndpoint.rpcEnv.setupEndpoint(endpointName, recordEndpoint)

      startOffset.partitionNums.map {
        case (part, index) =>
          LowLatencyMemoryStreamInputPartition(
            endpointName,
            endpointRef.address,
            part,
            index,
            endOffset.partitionNums(part)
          )
      }.toArray
    }
  }

  override def createReaderFactory(): PartitionReaderFactory = {
    new LowLatencyMemoryStreamReaderFactory(clock)
  }

  override def stop(): Unit = {
    if (endpointRef != null) recordEndpoint.rpcEnv.stop(endpointRef)
  }

  override def commit(end: OffsetV2): Unit = {}

  override def reset(): Unit = synchronized {
    super.reset()
    records.foreach(_.clear())
  }
}

object LowLatencyMemoryStream {
  protected val memoryStreamId = new AtomicInteger(0)

  /** Creates a LowLatencyMemoryStream with an implicit SQLContext (backward compatible). */
  def apply[A: Encoder](implicit sqlContext: SQLContext): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](memoryStreamId.getAndIncrement(), sqlContext.sparkSession)

  /** Creates a LowLatencyMemoryStream with specified partitions (SQLContext). */
  def apply[A: Encoder](numPartitions: Int)(
      implicit sqlContext: SQLContext): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](
      memoryStreamId.getAndIncrement(),
      sqlContext.sparkSession,
      numPartitions = numPartitions)

  /** Creates a LowLatencyMemoryStream with explicit SparkSession. */
  def apply[A: Encoder](sparkSession: SparkSession): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](memoryStreamId.getAndIncrement(), sparkSession)

  /** Creates a LowLatencyMemoryStream with specified partitions (SparkSession). */
  def apply[A: Encoder](sparkSession: SparkSession, numPartitions: Int): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](
      memoryStreamId.getAndIncrement(),
      sparkSession,
      numPartitions = numPartitions)

  /** Creates a single partition LowLatencyMemoryStream (SQLContext). */
  def singlePartition[A: Encoder](implicit sqlContext: SQLContext): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](memoryStreamId.getAndIncrement(), sqlContext.sparkSession, 1)

  /** Creates a single partition LowLatencyMemoryStream (SparkSession). */
  def singlePartition[A: Encoder](sparkSession: SparkSession): LowLatencyMemoryStream[A] =
    new LowLatencyMemoryStream[A](memoryStreamId.getAndIncrement(), sparkSession, 1)
}

/**
 * An input partition for LowLatency memory stream.
 */
case class LowLatencyMemoryStreamInputPartition(
    driverEndpointName: String,
    driverEndpointAddress: RpcAddress,
    partition: Int,
    startOffset: Int,
    endOffset: Int)
    extends InputPartition

class LowLatencyMemoryStreamReaderFactory(clock: Clock) extends PartitionReaderFactory {
  override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
    val p = partition.asInstanceOf[LowLatencyMemoryStreamInputPartition]
    new LowLatencyMemoryStreamPartitionReader(
      p.driverEndpointName,
      p.driverEndpointAddress,
      p.partition,
      p.startOffset,
      p.endOffset,
      clock
    )
  }
}

/**
 * An input partition reader for LowLatency memory stream.
 *
 * Polls the driver endpoint for new records.
 */
class LowLatencyMemoryStreamPartitionReader(
    driverEndpointName: String,
    driverEndpointAddress: RpcAddress,
    partition: Int,
    startOffset: Int,
    endOffset: Int,
    clock: Clock)
    extends SupportsRealTimeRead[InternalRow] {
  // Avoid tracking the ref, given that we create a new one for each partition reader
  // because a new driver endpoint is created for each LowLatencyMemoryStream. If we track the ref,
  // we can end up with a lot of refs (1000s) if a test suite has so many test cases and can lead to
  // issues with the tracking array. Causing the test suite to be flaky.
  private val endpoint = RpcUtils.makeDriverRef(
    driverEndpointName,
    driverEndpointAddress.host,
    driverEndpointAddress.port,
    SparkEnv.get.rpcEnv
  )

  private var currentOffset = startOffset
  private var current: Option[(InternalRow, Long)] = None

  // Defense-in-depth against failing to propagate the task context. Since it's not inheritable,
  // we have to do a bit of error prone work to get it into every thread used by LowLatency
  // processing. We hope that some unit test will end up instantiating a LowLatency memory stream
  // in such cases.
  if (TaskContext.get() == null) {
    throw new IllegalStateException("Task context was not set!")
  }
  override def nextWithTimeout(
      startTimeMs: java.lang.Long, timeoutMs: java.lang.Long): RecordStatus = {
    // SPARK-55699: Use the reference time passed in by the caller instead of getting the latest
    // time from LowLatencyClock, to avoid inconsistent reading when LowLatencyClock is a
    // manual clock.
    var elapsedTimeMs = 0L
    current = getRecordWithTimestamp
    while (current.isEmpty) {
      val POLL_TIME = 10L
      if (elapsedTimeMs >= timeoutMs) {
        return RecordStatus.newStatusWithoutArrivalTime(false)
      }
      Thread.sleep(POLL_TIME)
      current = getRecordWithTimestamp
      elapsedTimeMs = clock.getTimeMillis() - startTimeMs
    }
    currentOffset += 1
    RecordStatus.newStatusWithArrivalTimeMs(current.get._2)
  }

  override def next(): Boolean = {
    current = getRecordWithTimestamp
    if (current.isDefined) {
      currentOffset += 1
      true
    } else {
      false
    }
  }

  override def get(): InternalRow = current.get._1
  override def close(): Unit = {}

  override def getOffset: ContinuousRecordPartitionOffset =
    ContinuousRecordPartitionOffset(partition, currentOffset)

  private def getRecordWithTimestamp: Option[(InternalRow, Long)] = {
    if (currentOffset >= endOffset) {
      return None
    }
    endpoint.askSync[Option[(InternalRow, Long)]](
      GetRecord(ContinuousRecordPartitionOffset(partition, currentOffset))
    )
  }
}

case class LowLatencyMemoryStreamOffset(partitionNums: Map[Int, Int]) extends Offset {
  private implicit val formats: Formats = Serialization.formats(NoTypeHints)
  override def json(): String = Serialization.write(partitionNums)
}

/**
 * Custom endpoint for LowLatencyMemoryStream to support (row, timestamp) retrieval via RPC.
 */
class LowLatencyMemoryStreamEndpoint(
    records: Seq[scala.collection.mutable.Seq[(UnsafeRow, Long)]],
    lock: Object
) extends ThreadSafeRpcEndpoint {
  override val rpcEnv: RpcEnv = SparkEnv.get.rpcEnv

  override def receiveAndReply(
      context: org.apache.spark.rpc.RpcCallContext
  ): PartialFunction[Any, Unit] = {
    case GetRecord(ContinuousRecordPartitionOffset(partitionId, offset)) =>
      lock.synchronized {
        val buf = records(partitionId)
        val record = if (buf.size <= offset) None else Some(buf(offset))

        context.reply(record)
      }
  }
}
