/*
 * 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.texera.amber.engine.architecture.worker

import com.twitter.util.Future
import org.apache.texera.amber.core.executor.{
  ExecFactory,
  OpExecInitInfo,
  OpExecSource,
  OpExecWithClassName,
  OpExecWithCode
}
import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
  AsyncRPCContext,
  DebugCommandRequest,
  EmptyRequest,
  EvaluatePythonExpressionRequest
}
import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{EmptyReturn, EvaluatedValue}
import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceFs2Grpc
import org.apache.texera.amber.engine.architecture.worker.promisehandlers._
import org.apache.texera.amber.engine.common.AmberLogging
import org.apache.texera.amber.engine.common.rpc.AsyncRPCHandlerInitializer
import org.apache.texera.amber.operator.source.cache.CacheSourceOpExec

import java.net.URI

class DataProcessorRPCHandlerInitializer(val dp: DataProcessor)
    extends AsyncRPCHandlerInitializer(dp.asyncRPCClient, dp.asyncRPCServer)
    with WorkerServiceFs2Grpc[Future, AsyncRPCContext]
    with AmberLogging
    with InitializeExecutorHandler
    with OpenExecutorHandler
    with PauseHandler
    with AddPartitioningHandler
    with QueryStatisticsHandler
    with ResumeHandler
    with StartHandler
    with EndHandler
    with StartChannelHandler
    with EndChannelHandler
    with AssignPortHandler
    with AddInputChannelHandler
    with FlushNetworkBufferHandler
    with RetrieveStateHandler
    with PrepareCheckpointHandler
    with FinalizeCheckpointHandler
    with UpdateExecutorHandler {
  val actorId: ActorVirtualIdentity = dp.actorId

  var cachedTotalWorkerCount = 0

  override def debugCommand(
      request: DebugCommandRequest,
      ctx: AsyncRPCContext
  ): Future[EmptyReturn] = ???

  override def evaluatePythonExpression(
      request: EvaluatePythonExpressionRequest,
      ctx: AsyncRPCContext
  ): Future[EvaluatedValue] = ???

  override def retryCurrentTuple(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] =
    ???

  override def noOperation(request: EmptyRequest, ctx: AsyncRPCContext): Future[EmptyReturn] = ???

  def setupExecutor(execInitInfo: OpExecInitInfo, workerIdx: Int, workerCount: Int): Unit = {
    dp.executor = execInitInfo match {
      case OpExecWithClassName(className, descString) =>
        ExecFactory.newExecFromJavaClassName(className, descString, workerIdx, workerCount)
      case OpExecWithCode(code, _) =>
        ExecFactory.newExecFromJavaCode(code)
      case OpExecSource(storageUri, _) =>
        new CacheSourceOpExec(URI.create(storageUri))
      case OpExecInitInfo.Empty =>
        throw new IllegalArgumentException("Empty executor initialization info")
    }
  }
}
