worker.scala

package object worker {

  case class QueueUrl(value: String) extends AnyVal

  trait SerDe[A] {
    def serialize(a: A): String
    def deserialize(s: String): A
  }

  trait Ackable {
    import com.mfglabs.metadsp.SQSClient.client

    val msgQueueUrl: QueueUrl
    val receiptHandle: String
    val initialVisibilityTimeout: FiniteDuration

    def ack(): Future[Unit] = {
      client.deleteMessage(msgQueueUrl.value, receiptHandle)
    }

    def changeVisibilityTimeout(visibility: FiniteDuration): Future[Unit] = {
      client.changeMessageVisibility(msgQueueUrl.value, receiptHandle, visibility.toSeconds.toInt)
    }
  }

  type OneShotJob[A, B] = (A, Ackable) => Future[(B, QueueUrl)]

  type StreamJob[A, B] = Flow[(A, Ackable), (B, QueueUrl), Unit]

  abstract class Worker[A: SerDe, B: SerDe](inputQueue: QueueUrl) {
    val inputSerDe = implicitly[SerDe[A]]
    val ouputSerDe = implicitly[SerDe[B]]

    val flow: Flow[(A, Ackable), (B, QueueUrl), Unit]

    val source: Source[(A, Ackable), Unit] = {
      SQSClient.stream
        .receiveMessageAsStream(inputQueue.value, autoAck = false, attributes =  Seq("VisibilityTimeout"))
        .map { msg =>
          val ackable: Ackable = new Ackable {
            val msgQueueUrl = inputQueue
            val initialVisibilityTimeout = FiniteDuration(msg.getAttributes().get("VisibilityTimeout").toInt, SECONDS)
            val receiptHandle = msg.getReceiptHandle
          }
          (inputSerDe.deserialize(msg.getBody), ackable)
        }
        .mapMaterialized(_ => ())
    }

    val sink: Flow[(B, QueueUrl), SendMessageResult, Unit] = {
      Flow[(B, QueueUrl)].map { case (out, queue) =>
        val req = new SendMessageRequest()
        req.setMessageBody(ouputSerDe.serialize(out))
        req.setQueueUrl(queue.value)
        req
      }
      .via(SQSClient.stream.sendMessageAsStream)
    }

    def run()(implicit fm: FlowMaterializer) = source.via(flow).via(sink).runForeach(_ => ())
  }

  object StreamJob {
    def fromOneShotJob[I, O](maxConcurrency: Int, job: OneShotJob[I, O]): StreamJob[I, O] =
      Flow[(I, Ackable)].via(FlowExt.mapAsyncWithBoundedConcurrency(maxConcurrency) { case (i, ackable) =>
        job(i, ackable)
      })

    def fromOneShotJob2[I, O, I1 <: I, O1 <: O, I2 <: I, O2 <: O](maxConcurrency: Int,
                                                                  job1: OneShotJob[I1, O1], job2: OneShotJob[I2, O2]): StreamJob[I, O] =
      Flow[(I, Ackable)].via(FlowExt.mapAsyncWithBoundedConcurrency(maxConcurrency) {
        case (i1: I1, ackable: Ackable) => job1(i1, ackable)
        case (i2: I2, ackable: Ackable) => job2(i2, ackable)
      })

    def fromOneShotJob3[I, O, I1 <: I, O1 <: O, I2 <: I, O2 <: O, I3 <: I, O3 <: O](maxConcurrency: Int,
                                                                                    job1: OneShotJob[I1, O1], job2: OneShotJob[I2, O2],
                                                                                    job3: OneShotJob[I3, O3]): StreamJob[I, O] =
      Flow[(I, Ackable)].via(FlowExt.mapAsyncWithBoundedConcurrency(maxConcurrency) {
        case (i1: I1, ackable: Ackable) => job1(i1, ackable)
        case (i2: I2, ackable: Ackable) => job2(i2, ackable)
        case (i3: I3, ackable: Ackable) => job3(i3, ackable)
      })
  }

  object Worker {
    def fromOneShotJob[I, O](maxConcurrency: Int, inputQueue: QueueUrl, job: OneShotJob[I, O]) = new Worker[I, O](inputQueue) {
      override val flow = StreamJob.fromOneShotJob(maxConcurrency, job)
    }

    def fromOneShotJob2[I, O, I1 <: I, O1 <: O, I2 <: I, O2 <: O](maxConcurrency: Int, inputQueue: QueueUrl,
                                                                  job1: OneShotJob[I1, O1], job2: OneShotJob[I2, O2]) =
      new Worker[I, O](inputQueue) {
        override val flow = StreamJob.fromOneShotJob2(maxConcurrency, job1, job2)
      }

    def fromOneShotJob3[I, O, I1 <: I, O1 <: O, I2 <: I, O2 <: O, I3 <: I, O3 <: O](maxConcurrency: Int, inputQueue: QueueUrl,
                                                                  job1: OneShotJob[I1, O1], job2: OneShotJob[I2, O2],
                                                                  job3: OneShotJob[I3, O3]) =
      new Worker[I, O](inputQueue) {
        override val flow = StreamJob.fromOneShotJob3(maxConcurrency, job1, job2, job3)
      }

    def fromStreamJob[I, O](inputQueue: QueueUrl, job: StreamJob[I, O]) = new Worker[I, O](inputQueue) {
      override val flow = job
    }

    class ADTBasedFanOut2[I, I1 <: I, I2 <: I] extends FlexiRoute[(I, Ackable), FanOutShape2[(I, Ackable), (I1, Ackable), (I2, Ackable)]](
      new FanOutShape2(FanOutShape.Name[(I, Ackable)]("TypeBasedRoute2")), OperationAttributes.name("TypeBasedRoute2")) {
        import FlexiRoute._

        override def createRouteLogic(p: PortT) = new RouteLogic[(I, Ackable)] {
          override def initialState = State[Any](DemandFromAll(p.out0, p.out1)) { (ctx, _, element) =>
            element match {
              case (i1: I1, ackable: Ackable) => ctx.emit(p.out0)(i1, ackable)
              case (i2: I2, ackable: Ackable) => ctx.emit(p.out1)(i2, ackable)
            }
            SameState
          }
        }
      }

    def fromStreamJob2[I, O, I1 <: I, O1 <: O, I2 <: I, O2 <: O](inputQueue: QueueUrl, job1: StreamJob[I1, O1],
                                                                 job2: StreamJob[I2, O2]) = {
      new Worker[I, O](inputQueue) {
        override val flow = Flow() { implicit b =>
          import FlowGraph.Implicits._
          val flexiRoute = b.add(new ADTBasedFanOut2[I, I1, I2])
          val merge = b.add(Merge[(O, QueueUrl)](2))

          flexiRoute.out0 ~> job1 ~> merge.in(0)
          flexiRoute.out1 ~> job2 ~> merge.in(1)

          (flexiRoute.in, merge.out)
        }
      }
    }

  }

}