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

import java.util.UUID
import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, Executors, ScheduledExecutorService, TimeUnit}
import java.util.concurrent.atomic.{AtomicLong, AtomicReference}

import scala.concurrent.duration.FiniteDuration
import scala.jdk.CollectionConverters._
import scala.util.control.NonFatal

import com.google.common.cache.CacheBuilder
import io.grpc.stub.StreamObserver

import org.apache.spark.{SparkEnv, SparkSQLException}
import org.apache.spark.connect.proto
import org.apache.spark.internal.{Logging, LogKeys}
import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_MILLIS
import org.apache.spark.sql.connect.config.Connect.{CONNECT_EXECUTE_MANAGER_ABANDONED_TOMBSTONES_SIZE, CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT, CONNECT_EXECUTE_MANAGER_MAINTENANCE_INTERVAL}
import org.apache.spark.sql.connect.execution.ExecuteGrpcResponseSender
import org.apache.spark.util.ThreadUtils

// Unique key identifying execution by combination of user, session and operation id
case class ExecuteKey(userId: String, sessionId: String, operationId: String)

object ExecuteKey {
  def apply(request: proto.ExecutePlanRequest, sessionHolder: SessionHolder): ExecuteKey = {
    val operationId = if (request.hasOperationId) {
      try {
        UUID.fromString(request.getOperationId).toString
      } catch {
        case _: IllegalArgumentException =>
          throw new SparkSQLException(
            errorClass = "INVALID_HANDLE.FORMAT",
            messageParameters = Map("handle" -> request.getOperationId))
      }
    } else {
      UUID.randomUUID().toString
    }
    ExecuteKey(sessionHolder.userId, sessionHolder.sessionId, operationId)
  }
}

/**
 * Global tracker of all ExecuteHolder executions.
 *
 * All ExecuteHolders are created, and removed through it. It keeps track of all the executions,
 * and removes executions that have been abandoned.
 */
private[connect] class SparkConnectExecutionManager() extends Logging {

  /** Concurrent hash table containing all the current executions. */
  private val executions: ConcurrentMap[ExecuteKey, ExecuteHolder] =
    new ConcurrentHashMap[ExecuteKey, ExecuteHolder]()

  /** Graveyard of tombstones of executions that were abandoned and removed. */
  private val abandonedTombstones = CacheBuilder
    .newBuilder()
    .maximumSize(SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_ABANDONED_TOMBSTONES_SIZE))
    .build[ExecuteKey, ExecuteInfo]()

  /** The time when the last execution was removed. */
  private val lastExecutionTimeNs: AtomicLong = new AtomicLong(System.nanoTime())

  /** Executor for the periodic maintenance */
  private val scheduledExecutor: AtomicReference[ScheduledExecutorService] =
    new AtomicReference[ScheduledExecutorService]()

  /**
   * Create a new ExecuteHolder and register it with this global manager and with its session.
   */
  private[connect] def createExecuteHolder(
      executeKey: ExecuteKey,
      request: proto.ExecutePlanRequest,
      sessionHolder: SessionHolder): ExecuteHolder = {
    val opId = executeKey.operationId
    val executeHolder = executions.compute(
      executeKey,
      (executeKey, oldExecuteHolder) => {

        // Check if the operation already exists, either in the active execution map
        if (oldExecuteHolder != null) {
          throw new SparkSQLException(
            errorClass = "INVALID_HANDLE.OPERATION_ALREADY_EXISTS",
            messageParameters = Map("handle" -> opId))
        }
        // Check if the operation is already in the graveyard of abandoned executions, or was
        // recently completed. Prevents double execution when client retries on a lost response.
        if (getAbandonedTombstone(executeKey).isDefined ||
          sessionHolder.getOperationStatus(opId).isDefined) {

          logInfo(
            log"Operation ${MDC(LogKeys.EXECUTE_KEY, executeKey)}: Already tombstoned: " +
              log"${MDC(LogKeys.STATUS, getAbandonedTombstone(executeKey).isDefined)}.")
          logInfo(
            log"Operation ${MDC(LogKeys.EXECUTE_KEY, executeKey)}: Seen previously: " +
              log"${MDC(LogKeys.STATUS, sessionHolder.getOperationStatus(opId).isDefined)}.")

          throw new SparkSQLException(
            errorClass = "INVALID_HANDLE.OPERATION_ABANDONED",
            messageParameters = Map("handle" -> opId))
        }
        new ExecuteHolder(executeKey, request, sessionHolder)
      })

    sessionHolder.addOperationId(opId)

    logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, executeHolder.key)} is created.")

    schedulePeriodicChecks() // Starts the maintenance thread if it hasn't started.

    executeHolder
  }

  /**
   * Create a new ExecuteHolder and register it with this global manager and with its session.
   */
  private[connect] def createExecuteHolder(v: proto.ExecutePlanRequest): ExecuteHolder = {
    val previousSessionId = v.hasClientObservedServerSideSessionId match {
      case true => Some(v.getClientObservedServerSideSessionId)
      case false => None
    }
    val sessionHolder = SparkConnectService
      .getOrCreateIsolatedSession(v.getUserContext.getUserId, v.getSessionId, previousSessionId)
    val executeKey = ExecuteKey(v, sessionHolder)
    createExecuteHolder(executeKey, v, sessionHolder)
  }

  /**
   * Remove an ExecuteHolder from this global manager and from its session. Interrupt the
   * execution if still running, free all resources.
   */
  private[connect] def removeExecuteHolder(key: ExecuteKey, abandoned: Boolean = false): Unit = {
    val executeHolder = executions.get(key)

    if (executeHolder == null) {
      return
    }

    // Put into abandonedTombstones before removing it from executions, so that the client ends up
    // getting an INVALID_HANDLE.OPERATION_ABANDONED error on a retry.
    if (abandoned) {
      abandonedTombstones.put(key, executeHolder.getExecuteInfo)
      executeHolder.sessionHolder.closeOperation(executeHolder.operationId)
    }

    // Remove the execution from the map *after* putting it in abandonedTombstones.
    executions.remove(key)
    executeHolder.sessionHolder.closeOperation(executeHolder.operationId)

    updateLastExecutionTime()

    logInfo(log"ExecuteHolder ${MDC(LogKeys.EXECUTE_KEY, key)} is removed.")

    executeHolder.close()
    if (abandoned) {
      // Update in abandonedTombstones: above it wasn't yet updated with closedTime etc.
      abandonedTombstones.put(key, executeHolder.getExecuteInfo)
    }
  }

  private[connect] def getExecuteHolder(key: ExecuteKey): Option[ExecuteHolder] = {
    Option(executions.get(key))
  }

  /**
   * Create a new ExecuteHolder, register it with this global manager and with its session, and
   * attach the given response observer to it.
   */
  private[connect] def createExecuteHolderAndAttach(
      executeKey: ExecuteKey,
      request: proto.ExecutePlanRequest,
      sessionHolder: SessionHolder,
      responseObserver: StreamObserver[proto.ExecutePlanResponse]): ExecuteHolder = {
    val executeHolder = createExecuteHolder(executeKey, request, sessionHolder)
    try {
      executeHolder.eventsManager.postStarted()
      executeHolder.start()
    } catch {
      // Errors raised before the execution holder has finished spawning a thread are considered
      // plan execution failure, and the client should not try reattaching it afterwards.
      case t: Throwable =>
        removeExecuteHolder(executeHolder.key)
        throw t
    }

    try {
      val responseSender =
        new ExecuteGrpcResponseSender[proto.ExecutePlanResponse](executeHolder, responseObserver)
      executeHolder.runGrpcResponseSender(responseSender)
    } finally {
      executeHolder.afterInitialRPC()
    }
    executeHolder
  }

  /**
   * Reattach the given response observer to the given ExecuteHolder.
   */
  private[connect] def reattachExecuteHolder(
      executeHolder: ExecuteHolder,
      responseObserver: StreamObserver[proto.ExecutePlanResponse],
      lastConsumedResponseId: Option[String]): Unit = {
    if (!executeHolder.reattachable) {
      logWarning(log"Reattach to not reattachable operation.")
      throw new SparkSQLException(
        errorClass = "INVALID_CURSOR.NOT_REATTACHABLE",
        messageParameters = Map.empty)
    } else if (executeHolder.isOrphan()) {
      logWarning(log"Reattach to an orphan operation.")
      removeExecuteHolder(executeHolder.key)
      throw new IllegalStateException("Operation was orphaned because of an internal error.")
    }

    val responseSender =
      new ExecuteGrpcResponseSender[proto.ExecutePlanResponse](executeHolder, responseObserver)
    lastConsumedResponseId match {
      case Some(lastResponseId) =>
        // start from response after lastResponseId
        executeHolder.runGrpcResponseSender(responseSender, lastResponseId)
      case None =>
        // start from the start of the stream.
        executeHolder.runGrpcResponseSender(responseSender)
    }
  }

  private[connect] def removeAllExecutionsForSession(key: SessionKey): Unit = {
    executions.forEach((_, executeHolder) => {
      if (executeHolder.sessionHolder.key == key) {
        val info = executeHolder.getExecuteInfo
        logInfo(
          log"Execution ${MDC(LogKeys.EXECUTE_INFO, info)} removed in removeSessionExecutions.")
        removeExecuteHolder(executeHolder.key, abandoned = true)
      }
    })
  }

  /** Get info about abandoned execution, if there is one. */
  private[connect] def getAbandonedTombstone(key: ExecuteKey): Option[ExecuteInfo] = {
    Option(abandonedTombstones.getIfPresent(key))
  }

  /**
   * If there are no executions, return Left with System.nanoTime of last active execution.
   * Otherwise return Right with list of ExecuteInfo of all executions.
   */
  def listActiveExecutions: Either[Long, Seq[ExecuteInfo]] = {
    if (executions.isEmpty) {
      Left(lastExecutionTimeNs.getAcquire())
    } else {
      Right(executions.values().asScala.map(_.getExecuteInfo).toBuffer.toSeq)
    }
  }

  /**
   * Return list of executions that got abandoned and removed by periodic maintenance. This is a
   * cache, and the tombstones will be eventually removed.
   */
  def listAbandonedExecutions: Seq[ExecuteInfo] = {
    abandonedTombstones.asMap.asScala.values.toSeq
  }

  private[connect] def shutdown(): Unit = {
    val executor = scheduledExecutor.getAndSet(null)
    if (executor != null) {
      ThreadUtils.shutdown(executor, FiniteDuration(1, TimeUnit.MINUTES))
    }

    // note: this does not cleanly shut down the executions, but the server is shutting down.
    executions.clear()
    abandonedTombstones.invalidateAll()

    updateLastExecutionTime()
  }

  /**
   * Updates the last execution time after the last execution has been removed.
   */
  private def updateLastExecutionTime(): Unit = {
    lastExecutionTimeNs.getAndUpdate(prev => prev.max(System.nanoTime()))
  }

  /**
   * Schedules periodic maintenance checks if it is not already scheduled. The checks are looking
   * for executions that have not been closed, but are left with no RPC attached to them, and
   * removes them after a timeout.
   */
  private def schedulePeriodicChecks(): Unit = {
    var executor = scheduledExecutor.getAcquire()
    if (executor == null) {
      executor = Executors.newSingleThreadScheduledExecutor()
      if (scheduledExecutor.compareAndExchangeRelease(null, executor) == null) {
        val interval = SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_MAINTENANCE_INTERVAL)
        logInfo(
          log"Starting thread for cleanup of abandoned executions every " +
            log"${MDC(LogKeys.INTERVAL, interval)} ms")
        executor.scheduleAtFixedRate(
          () => {
            try {
              val timeoutNs =
                SparkEnv.get.conf.get(CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT) * NANOS_PER_MILLIS
              periodicMaintenance(timeoutNs)
            } catch {
              case NonFatal(ex) => logWarning("Unexpected exception in periodic task", ex)
            }
          },
          interval,
          interval,
          TimeUnit.MILLISECONDS)
      }
    }
  }

  // Visible for testing.
  private[connect] def periodicMaintenance(timeoutNs: Long): Unit = {
    // Find any detached executions that expired and should be removed.
    logDebug("Started periodic run of SparkConnectExecutionManager maintenance.")

    val nowNs = System.nanoTime()
    executions.forEach((_, executeHolder) => {
      executeHolder.lastAttachedRpcTimeNs match {
        case Some(detachedNs) =>
          if (detachedNs + timeoutNs <= nowNs) {
            val info = executeHolder.getExecuteInfo
            logInfo(
              log"Found execution ${MDC(LogKeys.EXECUTE_INFO, info)} that was abandoned " +
                log"and expired and will be removed.")
            removeExecuteHolder(executeHolder.key, abandoned = true)
          }
        case _ => // execution is active
      }
    })

    logDebug("Finished periodic run of SparkConnectExecutionManager maintenance.")
  }

  // For testing.
  private[connect] def setAllRPCsDeadline(deadlineNs: Long) = {
    executions.values().asScala.foreach(_.setGrpcResponseSendersDeadline(deadlineNs))
  }

  // For testing.
  private[connect] def interruptAllRPCs() = {
    executions.values().asScala.foreach(_.interruptGrpcResponseSenders())
  }

  private[connect] def listExecuteHolders: Seq[ExecuteHolder] = {
    executions.values().asScala.toSeq
  }
}
