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

import org.apache.pekko.actor.{ActorRef, ActorSystem, PoisonPill, Props}
import org.apache.pekko.testkit.{TestKit, TestProbe}
import io.grpc.MethodDescriptor
import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{
  GetActorRef,
  NetworkAck,
  NetworkMessage,
  RegisterActorRef
}
import org.apache.texera.amber.engine.architecture.control.utils.TrivialControlTester
import org.apache.texera.amber.engine.architecture.rpc.controlcommands._
import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{
  IntResponse,
  ReturnInvocation,
  StringResponse
}
import org.apache.texera.amber.engine.architecture.rpc.testerservice.RPCTesterGrpc._
import org.apache.texera.amber.engine.common.ambermessage.WorkflowFIFOMessage
import org.apache.texera.amber.engine.common.ambermessage.WorkflowMessage.getInMemSize
import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.ControlInvocation
import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER
import org.scalatest.wordspec.AnyWordSpecLike
import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}

import scala.collection.mutable
import scala.concurrent.duration._

class TrivialControlSpec
    extends TestKit(ActorSystem("TrivialControlSpec"))
    with AnyWordSpecLike
    with BeforeAndAfterEach
    with BeforeAndAfterAll {

  override def afterAll(): Unit = {
    TestKit.shutdownActorSystem(system)
  }

  def testControl[T](
      numActors: Int,
      eventPairs: ((MethodDescriptor[_, _], ControlRequest), T)*
  ): Unit = {
    val (events, expectedValues) = eventPairs.unzip
    val (probe, idMap) = setUp(numActors, events: _*)
    var flag = 0
    while (flag < expectedValues.length) {
      probe.receiveOne(10.seconds) match {
        case null =>
          throw new AssertionError(
            s"timeout: received $flag of ${expectedValues.length} expected returns"
          )
        case GetActorRef(id, replyTo) =>
          replyTo.foreach { actor =>
            actor ! RegisterActorRef(id, idMap(id))
          }
        case NetworkMessage(
              msgID,
              workflowMsg @ WorkflowFIFOMessage(_, _, ReturnInvocation(id, returnValue))
            ) =>
          probe.sender() ! NetworkAck(
            msgID,
            getInMemSize(workflowMsg),
            0L // no queued credit
          )
          assert(returnValue.asInstanceOf[T] == expectedValues(id.toInt))
          flag += 1
        case _ =>
        //skip
      }
    }
    idMap.foreach { x =>
      x._2 ! PoisonPill
    }
  }

  def setUp(
      numActors: Int,
      cmd: (MethodDescriptor[_, _], ControlRequest)*
  ): (TestProbe, mutable.HashMap[ActorVirtualIdentity, ActorRef]) = {
    val probe = TestProbe()
    val idMap = mutable.HashMap[ActorVirtualIdentity, ActorRef]()
    for (i <- 0 until numActors) {
      val id = ActorVirtualIdentity(s"$i")
      val ref =
        probe.childActorOf(Props(new TrivialControlTester(id)))
      idMap(id) = ref
    }
    idMap(CONTROLLER) = probe.ref
    var seqNum = 0
    cmd.foreach {
      case (methodName, msg) =>
        probe.send(
          idMap(ActorVirtualIdentity("0")),
          NetworkMessage(
            seqNum,
            WorkflowFIFOMessage(
              ChannelIdentity(CONTROLLER, ActorVirtualIdentity("0"), isControl = true),
              seqNum,
              ControlInvocation(
                methodName,
                msg,
                AsyncRPCContext(CONTROLLER, ActorVirtualIdentity("0")),
                seqNum
              )
            )
          )
        )
        seqNum += 1
    }
    (probe, idMap)
  }

  "testers" should {

    "execute Ping Pong" in {
      testControl(2, ((METHOD_SEND_PING, Ping(1, 5, ActorVirtualIdentity("1"))), IntResponse(5)))
    }

    "execute Ping Pong 2 times" in {
      testControl(
        2,
        ((METHOD_SEND_PING, Ping(1, 4, ActorVirtualIdentity("1"))), IntResponse(4)),
        ((METHOD_SEND_PING, Ping(10, 13, ActorVirtualIdentity("1"))), IntResponse(13))
      )
    }

    "execute Chain" in {
      testControl(
        10,
        (
          (METHOD_SEND_CHAIN, Chain((1 to 9).map(i => ActorVirtualIdentity(i.toString)))),
          StringResponse("9")
        )
      )
    }

    "execute Collect" in {
      testControl(
        4,
        (
          (METHOD_SEND_COLLECT, Collect((1 to 3).map(i => ActorVirtualIdentity(i.toString)))),
          StringResponse("finished")
        )
      )
    }

    "execute RecursiveCall" in {
      testControl(1, ((METHOD_SEND_RECURSION, Recursion(0)), StringResponse("0")))
    }

    "execute MultiCall" in {
      testControl(
        10,
        (
          (METHOD_SEND_MULTI_CALL, MultiCall((1 to 9).map(i => ActorVirtualIdentity(i.toString)))),
          StringResponse("finished")
        )
      )
    }

    "execute NestedCall" in {
      testControl(1, ((METHOD_SEND_NESTED, Nested(5)), StringResponse("Hello World!")))
    }

    "execute ErrorCall" in {
      assertThrows[RuntimeException] {
        testControl(1, ((METHOD_SEND_ERROR_COMMAND, ErrorCommand()), ()))
      }

    }
  }

}
