/*
 * 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.net.URI
import java.util.{Locale, UUID}
import java.util.concurrent.Executor

import scala.collection.mutable
import scala.jdk.CollectionConverters._
import scala.util.Properties
import scala.util.control.NonFatal

import com.google.protobuf
import com.google.protobuf.ByteString
import io.grpc._

import org.apache.spark.SparkBuildInfo.{spark_version => SPARK_VERSION}
import org.apache.spark.SparkThrowable
import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.connect.proto
import org.apache.spark.connect.proto.UserContext
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{ERROR, RATIO, SIZE, TIME}
import org.apache.spark.sql.connect.RuntimeConfig
import org.apache.spark.sql.connect.common.ProtoUtils
import org.apache.spark.sql.connect.common.config.ConnectCommon
import org.apache.spark.sql.util.CloseableIterator
import org.apache.spark.util.SparkSystemUtils

/**
 * Conceptually the remote spark session that communicates with the server.
 */
private[sql] class SparkConnectClient(
    private[sql] val configuration: SparkConnectClient.Configuration,
    private[sql] val channel: ManagedChannel)
    extends Logging {

  private val userContext: UserContext = configuration.userContext

  private[this] val stubState = new SparkConnectStubState(channel, configuration.retryPolicies)
  private[this] val bstub =
    new CustomSparkConnectBlockingStub(channel, stubState)
  private[this] val stub =
    new CustomSparkConnectStub(channel, stubState)

  private[client] def userAgent: String = configuration.userAgent

  /**
   * Placeholder method.
   * @return
   *   User ID.
   */
  private[sql] def userId: String = userContext.getUserId

  // Generate a unique session ID for this client. This UUID must be unique to allow
  // concurrent Spark sessions of the same user. If the channel is closed, creating
  // a new client will create a new session ID.
  private[sql] val sessionId: String = configuration.sessionId.getOrElse(UUID.randomUUID.toString)

  private val conf: RuntimeConfig = new RuntimeConfig(this)

  // Cached plan compression options.
  private var _planCompressionOptions: Option[Option[PlanCompressionOptions]] = None

  // Get the plan compression options. The options are cached after the first call.
  private[sql] def getPlanCompressionOptions: Option[PlanCompressionOptions] = {
    _planCompressionOptions match {
      case Some(options) => options
      case None =>
        val options =
          try {
            Some(
              PlanCompressionOptions(
                thresholdBytes =
                  conf.get("spark.connect.session.planCompression.threshold").toInt,
                algorithm = conf.get("spark.connect.session.planCompression.defaultAlgorithm")))
          } catch {
            // Disable plan compression if the server does not support it. Other exceptions are not
            // swallowed.
            case e: NoSuchElementException =>
              logWarning(
                log"Plan compression is disabled because the server does not support it",
                e)
              None
            case e: SparkThrowable
                if e.getCondition == "INVALID_CONF_VALUE"
                  || e.getCondition == "SQL_CONF_NOT_FOUND"
                  || e.getCondition == "CONFIG_NOT_AVAILABLE" =>
              logWarning(
                log"Plan compression is disabled because the server does not support it",
                e)
              None
          }
        _planCompressionOptions = Some(options)
        options
    }
  }

  // For testing and internal use only.
  private[sql] def setPlanCompressionOptions(
      planCompressionOptions: Option[PlanCompressionOptions]): Unit = {
    _planCompressionOptions = Some(planCompressionOptions)
  }

  /**
   * Handle plan compression errors.
   */
  private def handlePlanCompressionErrors[E](fn: => E): E = {
    try {
      fn
    } catch {
      // If the server cannot parse the compressed plan, disable plan compression for subsequent
      // requests on the session.
      case e: SparkThrowable if e.getCondition == "CONNECT_INVALID_PLAN.CANNOT_PARSE" =>
        logWarning(
          log"Disabling plan compression for the session due to " +
            log"CONNECT_INVALID_PLAN.CANNOT_PARSE error.")
        setPlanCompressionOptions(None)
        // Retry the code block without plan compression.
        fn
    }
  }

  /**
   * Hijacks the stored server side session ID with the given suffix. Used for testing to make
   * sure that server is validating the session ID.
   */
  private[sql] def hijackServerSideSessionIdForTesting(suffix: String) = {
    stubState.responseValidator.hijackServerSideSessionIdForTesting(suffix)
  }

  /**
   * Returns true if the session is valid on both the client and the server. A session becomes
   * invalid if the server side information about the client, e.g., session ID, does not
   * correspond to the actual client state.
   */
  private[sql] def isSessionValid: Boolean = {
    // The last known state of the session is store in `responseValidator`, because it is where the
    // client gets responses from the server.
    stubState.responseValidator.isSessionValid
  }

  private[sql] val artifactManager: ArtifactManager = {
    new ArtifactManager(configuration, sessionId, bstub, stub)
  }

  /**
   * Manually triggers upload of all classfile artifacts to the Spark Connect Server
   */
  private[sql] def uploadAllClassFileArtifacts(): Unit =
    artifactManager.uploadAllClassFileArtifacts()

  /**
   * Returns the server-side session id obtained from the first request, if there was a request
   * already.
   */
  private def serverSideSessionId: Option[String] = {
    stubState.responseValidator.getServerSideSessionId
  }

  /**
   * Dispatch the [[proto.AnalyzePlanRequest]] to the Spark Connect server.
   * @return
   *   A [[proto.AnalyzePlanResponse]] from the Spark Connect server.
   */
  def analyze(request: proto.AnalyzePlanRequest): proto.AnalyzePlanResponse = {
    artifactManager.uploadAllClassFileArtifacts()
    bstub.analyzePlan(request)
  }

  private def isValidUUID(uuid: String): Boolean = {
    try {
      UUID.fromString(uuid)
      true
    } catch {
      case _: IllegalArgumentException => false
    }
  }

  /**
   * Try to compress the plan if it exceeds the threshold defined in the planCompressionOptions.
   * Return the original plan if compression is disabled, not needed, or not effective.
   */
  private def tryCompressPlan(plan: proto.Plan): proto.Plan = {
    def tryCompressMessage(
        message: protobuf.Message,
        opType: proto.Plan.CompressedOperation.OpType,
        options: PlanCompressionOptions): Option[proto.Plan.CompressedOperation] = {
      val serialized = message.toByteArray
      if (serialized.length > options.thresholdBytes) {
        try {
          import com.github.luben.zstd.Zstd

          val startTime = System.nanoTime()
          val compressed = Zstd.compress(serialized)
          val duration = (System.nanoTime() - startTime) / 1e9
          val savingRatio = 1 - compressed.length.toDouble / serialized.length
          logDebug(
            log"Plan compression: original_size=${MDC(SIZE, serialized.length)}, " +
              log"compressed_size=${MDC(SIZE, compressed.length)}, " +
              log"saving_ratio=${MDC(RATIO, savingRatio)}, " +
              log"duration_s=${MDC(TIME, duration)}")
          if (compressed.length < serialized.length) {
            return Some(
              proto.Plan.CompressedOperation
                .newBuilder()
                .setData(ByteString.copyFrom(compressed))
                .setOpType(opType)
                .setCompressionCodec(proto.CompressionCodec.COMPRESSION_CODEC_ZSTD)
                .build())
          } else {
            logDebug(log"Plan compression not effective. Using original plan.")
          }
        } catch {
          case _: NoClassDefFoundError | _: ClassNotFoundException =>
            logInfo(log"Zstd library not available. Disabling plan compression.")
            setPlanCompressionOptions(None)
          case NonFatal(e) =>
            logWarning(
              log"Failed to compress plan: ${MDC(ERROR, e.getMessage)}. Using original " +
                log"plan and disabling plan compression.")
            setPlanCompressionOptions(None)
        }
      }
      None
    }

    def maybeCompressPlan(
        plan: proto.Plan,
        message: protobuf.Message,
        opType: proto.Plan.CompressedOperation.OpType,
        clearFn: proto.Plan.Builder => proto.Plan.Builder,
        options: PlanCompressionOptions): proto.Plan = {
      tryCompressMessage(message, opType, options) match {
        case Some(compressedOperation) =>
          clearFn(proto.Plan.newBuilder(plan)).setCompressedOperation(compressedOperation).build()
        case None => plan
      }
    }

    getPlanCompressionOptions match {
      case Some(options) if options.algorithm == "ZSTD" && options.thresholdBytes >= 0 =>
        plan.getOpTypeCase match {
          case proto.Plan.OpTypeCase.ROOT =>
            maybeCompressPlan(
              plan,
              plan.getRoot,
              proto.Plan.CompressedOperation.OpType.OP_TYPE_RELATION,
              _.clearRoot(),
              options)
          case proto.Plan.OpTypeCase.COMMAND =>
            maybeCompressPlan(
              plan,
              plan.getCommand,
              proto.Plan.CompressedOperation.OpType.OP_TYPE_COMMAND,
              _.clearCommand(),
              options)
          case _ => plan
        }
      case _ => plan
    }
  }

  /**
   * Execute the plan and return response iterator.
   *
   * It returns CloseableIterator. For resource management it is better to close it once you are
   * done. If you don't close it, it and the underlying data will be cleaned up once the iterator
   * is garbage collected.
   */
  def execute(
      plan: proto.Plan,
      operationId: Option[String] = None): CloseableIterator[proto.ExecutePlanResponse] = {
    artifactManager.uploadAllClassFileArtifacts()
    handlePlanCompressionErrors {
      // Compress the plan if needed.
      val maybeCompressedPlan = tryCompressPlan(plan)
      val request = proto.ExecutePlanRequest
        .newBuilder()
        .setPlan(maybeCompressedPlan)
        .setUserContext(userContext)
        .setSessionId(sessionId)
        .setClientType(userAgent)
        .addAllTags(tags.get.toSeq.asJava)

      // Add request option to allow result chunking.
      if (configuration.allowArrowBatchChunking) {
        val chunkingOptionsBuilder = proto.ResultChunkingOptions
          .newBuilder()
          .setAllowArrowBatchChunking(true)
        configuration.preferredArrowChunkSize.foreach { size =>
          chunkingOptionsBuilder.setPreferredArrowChunkSize(size)
        }
        request.addRequestOptions(
          proto.ExecutePlanRequest.RequestOption
            .newBuilder()
            .setResultChunkingOptions(chunkingOptionsBuilder.build())
            .build())
      }

      serverSideSessionId.foreach(session =>
        request.setClientObservedServerSideSessionId(session))
      operationId.foreach { opId =>
        require(
          isValidUUID(opId),
          s"Invalid operationId: $opId. The id must be an UUID string of " +
            "the format `00112233-4455-6677-8899-aabbccddeeff`")
        request.setOperationId(opId)
      }
      if (configuration.useReattachableExecute) {
        bstub.executePlanReattachable(request.build())
      } else {
        bstub.executePlan(request.build())
      }
    }
  }

  /**
   * Dispatch the [[proto.ConfigRequest]] to the Spark Connect server.
   * @return
   *   A [[proto.ConfigResponse]] from the Spark Connect server.
   */
  def config(operation: proto.ConfigRequest.Operation): proto.ConfigResponse = {
    val request = proto.ConfigRequest
      .newBuilder()
      .setOperation(operation)
      .setSessionId(sessionId)
      .setClientType(userAgent)
      .setUserContext(userContext)
    serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
    bstub.config(request.build())
  }

  /**
   * Builds a [[proto.AnalyzePlanRequest]] from `plan` and dispatched it to the Spark Connect
   * server.
   * @return
   *   A [[proto.AnalyzePlanResponse]] from the Spark Connect server.
   */
  def analyze(
      method: proto.AnalyzePlanRequest.AnalyzeCase,
      plan: Option[proto.Plan] = None,
      explainMode: Option[proto.AnalyzePlanRequest.Explain.ExplainMode] = None)
      : proto.AnalyzePlanResponse = {
    handlePlanCompressionErrors {
      val builder = proto.AnalyzePlanRequest.newBuilder()
      // Compress the plan if needed.
      val maybeCompressedPlan = plan match {
        case Some(p) => Some(tryCompressPlan(p))
        case None => None
      }
      method match {
        case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
          assert(maybeCompressedPlan.isDefined)
          builder.setSchema(
            proto.AnalyzePlanRequest.Schema
              .newBuilder()
              .setPlan(maybeCompressedPlan.get)
              .build())
        case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
          if (explainMode.isEmpty) {
            throw new IllegalArgumentException(s"ExplainMode is required in Explain request")
          }
          assert(maybeCompressedPlan.isDefined)
          builder.setExplain(
            proto.AnalyzePlanRequest.Explain
              .newBuilder()
              .setPlan(maybeCompressedPlan.get)
              .setExplainMode(explainMode.get)
              .build())
        case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
          assert(maybeCompressedPlan.isDefined)
          builder.setIsLocal(
            proto.AnalyzePlanRequest.IsLocal
              .newBuilder()
              .setPlan(maybeCompressedPlan.get)
              .build())
        case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
          assert(maybeCompressedPlan.isDefined)
          builder.setIsStreaming(
            proto.AnalyzePlanRequest.IsStreaming
              .newBuilder()
              .setPlan(maybeCompressedPlan.get)
              .build())
        case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
          assert(maybeCompressedPlan.isDefined)
          builder.setInputFiles(
            proto.AnalyzePlanRequest.InputFiles
              .newBuilder()
              .setPlan(maybeCompressedPlan.get)
              .build())
        case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION =>
          builder.setSparkVersion(proto.AnalyzePlanRequest.SparkVersion.newBuilder().build())
        case other => throw new IllegalArgumentException(s"Unknown Analyze request $other")
      }
      analyze(builder)
    }
  }

  def sameSemantics(plan: proto.Plan, otherPlan: proto.Plan): proto.AnalyzePlanResponse = {
    handlePlanCompressionErrors {
      val builder = proto.AnalyzePlanRequest.newBuilder()
      // Compress the plan if needed.
      val maybeCompressedPlan = tryCompressPlan(plan)
      val otherMaybeCompressedPlan = tryCompressPlan(otherPlan)
      builder.setSameSemantics(
        proto.AnalyzePlanRequest.SameSemantics
          .newBuilder()
          .setTargetPlan(maybeCompressedPlan)
          .setOtherPlan(otherMaybeCompressedPlan))
      analyze(builder)
    }
  }

  def semanticHash(plan: proto.Plan): proto.AnalyzePlanResponse = {
    handlePlanCompressionErrors {
      val builder = proto.AnalyzePlanRequest.newBuilder()
      // Compress the plan if needed.
      val maybeCompressedPlan = tryCompressPlan(plan)
      builder.setSemanticHash(
        proto.AnalyzePlanRequest.SemanticHash
          .newBuilder()
          .setPlan(maybeCompressedPlan))
      analyze(builder)
    }
  }

  private[sql] def analyze(
      builder: proto.AnalyzePlanRequest.Builder): proto.AnalyzePlanResponse = {
    val request = builder
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType(userAgent)
    serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
    analyze(request.build())
  }

  private[sql] def interruptAll(): proto.InterruptResponse = {
    val builder = proto.InterruptRequest.newBuilder()
    val request = builder
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType(userAgent)
      .setInterruptType(proto.InterruptRequest.InterruptType.INTERRUPT_TYPE_ALL)
    serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
    bstub.interrupt(request.build())
  }

  private[sql] def interruptTag(tag: String): proto.InterruptResponse = {
    val builder = proto.InterruptRequest.newBuilder()
    val request = builder
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType(userAgent)
      .setInterruptType(proto.InterruptRequest.InterruptType.INTERRUPT_TYPE_TAG)
      .setOperationTag(tag)
    serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
    bstub.interrupt(request.build())
  }

  private[sql] def interruptOperation(id: String): proto.InterruptResponse = {
    val builder = proto.InterruptRequest.newBuilder()
    val request = builder
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType(userAgent)
      .setInterruptType(proto.InterruptRequest.InterruptType.INTERRUPT_TYPE_OPERATION_ID)
      .setOperationId(id)
    serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session))
    bstub.interrupt(request.build())
  }

  private[sql] def releaseSession(): proto.ReleaseSessionResponse = {
    val builder = proto.ReleaseSessionRequest.newBuilder()
    val request = builder
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType(userAgent)
    bstub.releaseSession(request.build())
  }

  private[this] val tags = new InheritableThreadLocal[mutable.Set[String]] {
    override def childValue(parent: mutable.Set[String]): mutable.Set[String] = {
      // Note: make a clone such that changes in the parent tags aren't reflected in
      // those of the children threads.
      parent.clone()
    }
    override protected def initialValue(): mutable.Set[String] = new mutable.HashSet[String]()
  }

  private[sql] def addTag(tag: String): Unit = {
    // validation is also done server side, but this will give error earlier.
    ProtoUtils.throwIfInvalidTag(tag)
    tags.get += tag
  }

  private[sql] def removeTag(tag: String): Unit = {
    // validation is also done server side, but this will give error earlier.
    ProtoUtils.throwIfInvalidTag(tag)
    tags.get.remove(tag)
  }

  private[sql] def getTags(): Set[String] = {
    tags.get.toSet
  }

  private[sql] def clearTags(): Unit = {
    tags.get.clear()
  }

  def copy(): SparkConnectClient = configuration.toSparkConnectClient

  /**
   * Returns whether arrow batch chunking is allowed.
   */
  def allowArrowBatchChunking: Boolean = configuration.allowArrowBatchChunking

  /**
   * Returns the preferred arrow chunk size in bytes.
   */
  def preferredArrowChunkSize: Option[Int] = configuration.preferredArrowChunkSize

  /**
   * Add a single artifact to the client session.
   *
   * Currently only local files with extensions .jar and .class are supported.
   */
  def addArtifact(path: String): Unit = artifactManager.addArtifact(path)

  /**
   * Add a single artifact to the client session.
   *
   * Currently only local files with extensions .jar and .class are supported.
   */
  def addArtifact(uri: URI): Unit = artifactManager.addArtifact(uri)

  /**
   * Add a single in-memory artifact to the session while preserving the directory structure
   * specified by `target` under the session's working directory of that particular file
   * extension.
   *
   * Supported target file extensions are .jar and .class.
   *
   * ==Example==
   * {{{
   *  addArtifact(bytesBar, "foo/bar.class")
   *  addArtifact(bytesFlat, "flat.class")
   *  // Directory structure of the session's working directory for class files would look like:
   *  // ${WORKING_DIR_FOR_CLASS_FILES}/flat.class
   *  // ${WORKING_DIR_FOR_CLASS_FILES}/foo/bar.class
   * }}}
   */
  def addArtifact(bytes: Array[Byte], target: String): Unit =
    artifactManager.addArtifact(bytes, target)

  /**
   * Add a single artifact to the session while preserving the directory structure specified by
   * `target` under the session's working directory of that particular file extension.
   *
   * Supported target file extensions are .jar and .class.
   *
   * ==Example==
   * {{{
   *  addArtifact("/Users/dummyUser/files/foo/bar.class", "foo/bar.class")
   *  addArtifact("/Users/dummyUser/files/flat.class", "flat.class")
   *  // Directory structure of the session's working directory for class files would look like:
   *  // ${WORKING_DIR_FOR_CLASS_FILES}/flat.class
   *  // ${WORKING_DIR_FOR_CLASS_FILES}/foo/bar.class
   * }}}
   */
  def addArtifact(source: String, target: String): Unit =
    artifactManager.addArtifact(source, target)

  /**
   * Add multiple artifacts to the session.
   *
   * Currently only local files with extensions .jar and .class are supported.
   */
  def addArtifacts(uri: Seq[URI]): Unit = artifactManager.addArtifacts(uri)

  /**
   * Register a [[ClassFinder]] for dynamically generated classes.
   */
  def registerClassFinder(finder: ClassFinder): Unit = artifactManager.registerClassFinder(finder)

  /**
   * Shutdown the client's connection to the server.
   */
  def shutdown(): Unit = {
    channel.shutdownNow()
  }

  /**
   * Clone this client session, creating a new session with the same configuration and shared
   * state as the current session but with independent runtime state.
   *
   * @return
   *   A new SparkConnectClient instance with the cloned session.
   */
  @DeveloperApi
  def cloneSession(): SparkConnectClient = {
    clone(None)
  }

  /**
   * Clone this client session with a custom session ID, creating a new session with the same
   * configuration and shared state as the current session but with independent runtime state.
   *
   * @param newSessionId
   *   Custom session ID to use for the cloned session (must be a valid UUID).
   * @return
   *   A new SparkConnectClient instance with the cloned session.
   */
  @DeveloperApi
  def clone(newSessionId: String): SparkConnectClient = {
    clone(Some(newSessionId))
  }

  private def clone(newSessionId: Option[String]): SparkConnectClient = {
    val requestBuilder = proto.CloneSessionRequest
      .newBuilder()
      .setUserContext(userContext)
      .setSessionId(sessionId)
      .setClientType("scala")

    newSessionId.foreach(requestBuilder.setNewSessionId)

    val response: proto.CloneSessionResponse = bstub.cloneSession(requestBuilder.build())

    // Assert that the returned session ID matches the requested ID if one was provided
    newSessionId.foreach { expectedId =>
      require(
        response.getNewSessionId == expectedId,
        s"Returned session ID '${response.getNewSessionId}' does not match " +
          s"requested ID '$expectedId'")
    }

    // Create a new client with the cloned session ID
    val newConfiguration = configuration.copy(sessionId = Some(response.getNewSessionId))
    new SparkConnectClient(newConfiguration, configuration.createChannel())
  }
}

// Options for plan compression
case class PlanCompressionOptions(thresholdBytes: Int, algorithm: String)

object SparkConnectClient {

  private[sql] val SPARK_REMOTE: String = "SPARK_REMOTE"

  private val DEFAULT_USER_AGENT: String = "_SPARK_CONNECT_SCALA"

  private val AUTH_TOKEN_META_DATA_KEY: Metadata.Key[String] =
    Metadata.Key.of("Authentication", Metadata.ASCII_STRING_MARSHALLER)

  // for internal tests
  private[sql] def apply(channel: ManagedChannel): SparkConnectClient = {
    new SparkConnectClient(Configuration(), channel)
  }

  def builder(): Builder = new Builder()

  /**
   * This is a helper class that is used to create a GRPC channel based on either a set host and
   * port or a NameResolver-compliant URI connection string.
   */
  class Builder(private var _configuration: Configuration) {
    def this() = this(Configuration())

    def configuration: Configuration = _configuration

    def userId(id: String): Builder = {
      require(id != null && id.nonEmpty)
      _configuration = _configuration.copy(userId = id)
      this
    }

    def userId: Option[String] = Option(_configuration.userId)

    def userName(name: String): Builder = {
      require(name != null && name.nonEmpty)
      _configuration = _configuration.copy(userName = name)
      this
    }

    def userName: Option[String] = Option(_configuration.userName)

    def host(inputHost: String): Builder = {
      require(inputHost != null)
      _configuration = _configuration.copy(host = inputHost)
      this
    }

    def host: String = _configuration.host

    def port(inputPort: Int): Builder = {
      _configuration = _configuration.copy(port = inputPort)
      this
    }

    def port: Int = _configuration.port

    /**
     * Setting the token implicitly sets the use_ssl=true. All the following examples yield the
     * same results:
     *
     * {{{
     * sc://localhost/;token=aaa
     * sc://localhost/;use_ssl=true;token=aaa
     * sc://localhost/;token=aaa;use_ssl=true
     * }}}
     *
     * @param inputToken
     *   the user token.
     * @return
     *   this builder.
     */
    def token(inputToken: String): Builder = {
      require(inputToken != null && inputToken.nonEmpty)
      _configuration = _configuration.copy(token = Option(inputToken))
      this
    }

    def token: Option[String] = _configuration.token

    def enableSsl(): Builder = {
      _configuration = _configuration.copy(isSslEnabled = Option(true))
      this
    }

    /**
     * Disables the SSL. Throws exception if the token has been set.
     *
     * @return
     *   this builder.
     */
    def disableSsl(): Builder = {
      _configuration = _configuration.copy(isSslEnabled = Option(false))
      this
    }

    def sslEnabled: Boolean = _configuration.isSslEnabled.contains(true)

    def retryPolicy(policies: Seq[RetryPolicy]): Builder = {
      _configuration = _configuration.copy(retryPolicies = policies)
      this
    }

    def retryPolicy(policy: RetryPolicy): Builder = {
      retryPolicy(List(policy))
    }

    private object URIParams {
      val PARAM_USER_ID = "user_id"
      val PARAM_USE_SSL = "use_ssl"
      val PARAM_TOKEN = "token"
      val PARAM_USER_AGENT = "user_agent"
      val PARAM_SESSION_ID = "session_id"
      val PARAM_GRPC_MAX_MESSAGE_SIZE = "grpc_max_message_size"
    }

    private def verifyURI(uri: URI): Unit = {
      if (uri.getScheme != "sc") {
        throw new IllegalArgumentException("Scheme for connection URI must be 'sc'.")
      }
      if (uri.getHost == null) {
        throw new IllegalArgumentException(s"Host for connection URI must be defined.")
      }
      // Java URI considers everything after the authority segment as "path" until the
      // ? (query)/# (fragment) components as shown in the regex
      // [scheme:][//authority][path][?query][#fragment].
      // However, with the Spark Connect definition, configuration parameter are passed in the
      // style of the HTTP URL Path Parameter Syntax (e.g
      // sc://hostname:port/;param1=value;param2=value).
      // Thus, we manually parse the "java path" to get the "correct path" and configuration
      // parameters.
      val pathAndParams = uri.getPath.split(';')
      if (pathAndParams.nonEmpty && (pathAndParams(0) != "/" && pathAndParams(0) != "")) {
        throw new IllegalArgumentException(
          s"Path component for connection URI must be empty: " +
            s"${pathAndParams(0)}")
      }
    }

    def userAgent(value: String): Builder = {
      require(value != null)
      _configuration = _configuration.copy(userAgent = genUserAgent(value))
      this
    }

    def sessionId(value: String): Builder = {
      try {
        UUID.fromString(value).toString
      } catch {
        case e: IllegalArgumentException =>
          throw new IllegalArgumentException(
            "Parameter value 'session_id' must be a valid UUID format.",
            e)
      }
      _configuration = _configuration.copy(sessionId = Some(value))
      this
    }

    def sessionId: Option[String] = _configuration.sessionId

    def userAgent: String = _configuration.userAgent

    def grpcMaxMessageSize(messageSize: Int): Builder = {
      _configuration = _configuration.copy(grpcMaxMessageSize = messageSize)
      this
    }

    def grpcMaxMessageSize: Int = _configuration.grpcMaxMessageSize

    def grpcMaxRecursionLimit(recursionLimit: Int): Builder = {
      _configuration = _configuration.copy(grpcMaxRecursionLimit = recursionLimit)
      this
    }

    def grpcMaxRecursionLimit: Int = _configuration.grpcMaxRecursionLimit

    def option(key: String, value: String): Builder = {
      _configuration = _configuration.copy(metadata = _configuration.metadata + ((key, value)))
      this
    }

    def options: Map[String, String] = _configuration.metadata

    private def parseURIParams(uri: URI): Unit = {
      val params = uri.getPath.split(';').drop(1).filter(_ != "")
      params.foreach { kv =>
        val (key, value) = {
          val arr = kv.split('=')
          if (arr.length != 2) {
            throw new IllegalArgumentException(
              s"Parameter $kv is not a valid parameter" +
                s" key-value pair")
          }
          (arr(0), arr(1))
        }
        key match {
          case URIParams.PARAM_USER_ID => userId(value)
          case URIParams.PARAM_USER_AGENT => userAgent(value)
          case URIParams.PARAM_TOKEN => token(value)
          case URIParams.PARAM_USE_SSL =>
            if (java.lang.Boolean.valueOf(value)) enableSsl() else disableSsl()
          case URIParams.PARAM_SESSION_ID => sessionId(value)
          case URIParams.PARAM_GRPC_MAX_MESSAGE_SIZE => grpcMaxMessageSize(value.toInt)
          case _ => option(key, value)
        }
      }
    }

    /**
     * Configure the builder using the env SPARK_REMOTE environment variable.
     */
    def loadFromEnvironment(): Builder = {
      lazy val isAPIModeConnect =
        Option(System.getProperty(org.apache.spark.sql.SparkSessionBuilder.API_MODE_KEY))
          .getOrElse("classic")
          .toLowerCase(Locale.ROOT) == "connect"
      Option(System.getProperty("spark.remote")) // Set from Spark Submit
        .orElse(sys.env.get(SparkConnectClient.SPARK_REMOTE))
        .orElse {
          if (isAPIModeConnect) {
            Option(System.getProperty("spark.master")).orElse(sys.env.get("MASTER"))
          } else {
            None
          }
        }
        .foreach(connectionString)
      this
    }

    /**
     * Creates the channel with a target connection string, per the documentation of Spark
     * Connect.
     *
     * Note: The connection string, if used, will override any previous host/port settings.
     */
    def connectionString(connectionString: String): Builder = {
      val uri = new URI(connectionString)
      verifyURI(uri)
      parseURIParams(uri)
      host(uri.getHost)
      val inputPort = uri.getPort
      if (inputPort != -1) {
        port(uri.getPort)
      }
      this
    }

    /**
     * Configure the builder with the given CLI arguments.
     */
    def parse(args: Array[String]): Builder = {
      SparkConnectClientParser.parse(args.toList, this)
      this
    }

    /**
     * Add an interceptor to be used during channel creation.
     *
     * Note that interceptors added last are executed first by gRPC.
     */
    def interceptor(interceptor: ClientInterceptor): Builder = {
      val interceptors = _configuration.interceptors ++ List(interceptor)
      _configuration = _configuration.copy(interceptors = interceptors)
      this
    }

    /**
     * Disable reattachable execute.
     */
    def disableReattachableExecute(): Builder = {
      _configuration = _configuration.copy(useReattachableExecute = false)
      this
    }

    /**
     * Enable reattachable execute.
     *
     * It makes client more robust, enabling reattaching to an ExecutePlanResponse stream in case
     * of intermittent connection errors.
     */
    def enableReattachableExecute(): Builder = {
      _configuration = _configuration.copy(useReattachableExecute = true)
      this
    }

    def allowArrowBatchChunking(allow: Boolean): Builder = {
      _configuration = _configuration.copy(allowArrowBatchChunking = allow)
      this
    }

    def allowArrowBatchChunking: Boolean = _configuration.allowArrowBatchChunking

    def preferredArrowChunkSize(size: Option[Int]): Builder = {
      size.foreach(s => require(s > 0, "preferredArrowChunkSize must be positive"))
      _configuration = _configuration.copy(preferredArrowChunkSize = size)
      this
    }

    def preferredArrowChunkSize: Option[Int] = _configuration.preferredArrowChunkSize

    def build(): SparkConnectClient = _configuration.toSparkConnectClient
  }

  /**
   * Appends the Spark, Scala & JVM version, and the used OS to the user-provided user agent.
   */
  private def genUserAgent(value: String): String = {
    val scalaVersion = Properties.versionNumberString
    val jvmVersion = SparkSystemUtils.javaVersion.split("_")(0)
    val osName = {
      if (SparkSystemUtils.isMac) "darwin"
      else if (SparkSystemUtils.isLinux) "linux"
      else if (SparkSystemUtils.isWindows) "windows"
      else "unknown"
    }
    List(
      value,
      s"spark/$SPARK_VERSION",
      s"scala/$scalaVersion",
      s"jvm/$jvmVersion",
      s"os/$osName").mkString(" ")
  }

  private lazy val sparkUser =
    sys.env.getOrElse("SPARK_USER", System.getProperty("user.name", null))

  /**
   * Helper class that fully captures the configuration for a [[SparkConnectClient]].
   */
  private[sql] case class Configuration(
      userId: String = sparkUser,
      userName: String = sparkUser,
      host: String = "localhost",
      port: Int = ConnectCommon.CONNECT_GRPC_BINDING_PORT,
      token: Option[String] = None,
      isSslEnabled: Option[Boolean] = None,
      metadata: Map[String, String] = Map.empty,
      userAgent: String = genUserAgent(
        sys.env.getOrElse("SPARK_CONNECT_USER_AGENT", DEFAULT_USER_AGENT)),
      retryPolicies: Seq[RetryPolicy] = RetryPolicy.defaultPolicies(),
      useReattachableExecute: Boolean = true,
      interceptors: List[ClientInterceptor] = List.empty,
      sessionId: Option[String] = None,
      grpcMaxMessageSize: Int = ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE,
      grpcMaxRecursionLimit: Int = ConnectCommon.CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT,
      allowArrowBatchChunking: Boolean = true,
      preferredArrowChunkSize: Option[Int] = None) {

    private def isLocal = host.equals("localhost")

    def userContext: proto.UserContext = {
      val builder = proto.UserContext.newBuilder()
      if (userId != null) {
        builder.setUserId(userId)
      }
      if (userName != null) {
        builder.setUserName(userName)
      }
      builder.build()
    }

    def credentials: ChannelCredentials = {
      if (isSslEnabled.contains(true) || (token.isDefined && !isLocal)) {
        token match {
          case Some(t) =>
            // With access token added in the http header.
            CompositeChannelCredentials.create(
              TlsChannelCredentials.create,
              new AccessTokenCallCredentials(t))
          case None =>
            TlsChannelCredentials.create()
        }
      } else {
        InsecureChannelCredentials.create()
      }
    }

    def createChannel(): ManagedChannel = {
      val creds = credentials
      val channelBuilder = Grpc.newChannelBuilderForAddress(host, port, creds)

      // Workaround LocalChannelCredentials are added in
      // https://github.com/grpc/grpc-java/issues/9900
      var metadataWithOptionalToken = metadata
      if (!isSslEnabled.contains(true) && isLocal && token.isDefined) {
        metadataWithOptionalToken = metadata + (("Authorization", s"Bearer ${token.get}"))
      }

      if (metadataWithOptionalToken.nonEmpty) {
        channelBuilder.intercept(new MetadataHeaderClientInterceptor(metadataWithOptionalToken))
      }

      interceptors.foreach(channelBuilder.intercept(_))

      channelBuilder.maxInboundMessageSize(grpcMaxMessageSize)
      channelBuilder.build()
    }

    def toSparkConnectClient: SparkConnectClient = new SparkConnectClient(this, createChannel())
  }

  /**
   * A [[CallCredentials]] created from an access token.
   *
   * @param token
   *   A string to place directly in the http request authorization header, for example
   *   "authorization: Bearer <access_token>".
   */
  private[client] class AccessTokenCallCredentials(token: String) extends CallCredentials {
    override def applyRequestMetadata(
        requestInfo: CallCredentials.RequestInfo,
        appExecutor: Executor,
        applier: CallCredentials.MetadataApplier): Unit = {
      appExecutor.execute(() => {
        try {
          val headers = new Metadata()
          headers.put(AUTH_TOKEN_META_DATA_KEY, s"Bearer $token")
          applier.apply(headers)
        } catch {
          case e: Throwable =>
            applier.fail(Status.UNAUTHENTICATED.withCause(e));
        }
      })
    }
  }

  /**
   * A client interceptor to pass extra parameters in http request header.
   *
   * @param metadata
   *   extra metadata placed in the http request header, for example "key: value".
   */
  private[client] class MetadataHeaderClientInterceptor(metadata: Map[String, String])
      extends ClientInterceptor {
    override def interceptCall[ReqT, RespT](
        method: MethodDescriptor[ReqT, RespT],
        callOptions: CallOptions,
        next: Channel): ClientCall[ReqT, RespT] = {
      new ForwardingClientCall.SimpleForwardingClientCall[ReqT, RespT](
        next.newCall(method, callOptions)) {
        override def start(
            responseListener: ClientCall.Listener[RespT],
            headers: Metadata): Unit = {
          metadata.foreach { case (key, value) =>
            headers.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value)
          }
          super.start(responseListener, headers)
        }
      }
    }
  }
}
