test.scala

package object worker {

  trait SerDe[MsgType] {
    def serialize(msgType: MsgType): String
    def deserialize(s: String): MsgType
  }
  
  case class QueueUrl(value: String) extends AnyVal
  
  abstract class Queue[-A](implicit serde: SerDe[A]) {
    val queueUrl: QueueUrl
  
    case class Ackable(receiptHandle: String, initialVisibilityTimeout: FiniteDuration) {
      import SQSClient.client
  
      def ack(): Future[Unit] = {
        client.deleteMessage(queueUrl.value, receiptHandle)
      }
  
      def changeVisibilityTimeout(visibility: FiniteDuration): Future[Unit] = {
        client.changeMessageVisibility(queueUrl.value, receiptHandle, visibility.toSeconds.toInt)
      }
    }
  
    class ToAck(val ackables: Seq[Ackable])
    object ToAck {
      def apply(ackables: Seq[Ackable]) = new ToAck(ackables)
      def apply(ackable: Ackable) = ToAck(Seq(ackable))
    }
  
    def from[B <: A]: Source[(B, Ackable), Unit] = { // hack due to contravariance. B =:= A
      SQSClient.stream
        .receiveMessageAsStream(queueUrl.value, autoAck = false, attributes =  Seq("VisibilityTimeout"))
        .map { msg =>
          val ackable = Ackable(msg.getReceiptHandle, FiniteDuration(msg.getAttributes().get("VisibilityTimeout").toInt, SECONDS))
          (serde.deserialize(msg.getBody).asInstanceOf[B], ackable) // TODO: ???
        }
        .mapMaterialized(_ => ())
    }
  
    def to: Flow[(A, ToAck), Unit, Unit] = {
      Flow[(A, ToAck)].map { case (out, toAck) =>
        val req = new SendMessageRequest()
        req.setMessageBody(serde.serialize(out))
        req.setQueueUrl(queueUrl.value)
        (req, toAck)
      }
      .via(FlowExtLocal.tupleVia(SQSClient.stream.sendMessageAsStream, Flow[ToAck].map(identity)))
      .mapConcat { case (_, toAck) => toAck.ackables.to[scala.collection.immutable.Seq] }
      .mapAsync(_.ack())
    }
  }
  
  object Queue {
    object BlackHoleQueue extends Queue[Unit] {
      val queueUrl = QueueUrl("backhole-queue")
  
      override def from[B <: Unit]: Source[(B, Ackable), Unit] = Source.empty()
  
      override def to: Flow[(Unit, ToAck), Unit, Unit] =
        Flow[(Unit, ToAck)]
          .mapConcat { case (_, toAck) => toAck.ackables.to[scala.collection.immutable.Seq] }
          .mapAsync(_.ack())
    }
  }
  
  trait Worker0[In] {
    val in: Queue[In]
    val blackHoleOut = Queue.BlackHoleQueue
    val flow: Flow[(In, in.Ackable), (Unit, blackHoleOut.ToAck), Unit]
  
    def run()(implicit fm: FlowMaterializer) = in.from[In].via(flow).via(blackHoleOut.to).runWith(Sink.ignore())
  }
  
  trait Worker1[In, Out] {
    val in: Queue[In]
    val out: Queue[Out]
    val flow: Flow[(In, in.Ackable), (Out, out.ToAck), Unit]
  
    def run()(implicit fm: FlowMaterializer) = in.from[In].via(flow).via(out.to).runWith(Sink.ignore())
  }
  
  trait Worker2[In, Out1, Out2] {
    val in: Queue[In]
    val out1: Queue[Out1]
    val out2: Queue[Out2]
  
    val flow1: Flow[(In, in.Ackable), (Out1, out1.ToAck), Unit]
    val flow2: Flow[(In, in.Ackable), (Out2, out2.ToAck), Unit]
  
    def run()(implicit fm: FlowMaterializer) = {
      val graph = FlowGraph.closed() { implicit builder: FlowGraph.Builder =>
        import FlowGraph.Implicits._
        val bcast = builder.add(Broadcast[(In, in.Ackable)](2))
  
        in.from[In] ~> bcast ~> flow1 ~> out1.to ~> Sink.ignore
        bcast ~> flow2 ~> out2.to ~> Sink.ignore
      }
  
      graph.run()
    }
}


// USAGE

/**
 * Queues
 */

object DataTransferQueue extends Queue[DataTransferMessage] {
  val queueUrl = QueueUrl("todo")
}

object RedshiftLoadQueue extends Queue[RedshiftLoadMessage] {
  val queueUrl = QueueUrl("todo")
}

/**
 * Messages
 */

sealed trait DataTransferMessage
case class AppNexusPullPerformanceReport(startDate: DateTime, endDate: DateTime) extends DataTransferMessage
case class MediaMathPullPerformanceReport(startDate: DateTime, endDate: DateTime) extends DataTransferMessage

sealed trait RedshiftLoadMessage
case class AppNexusLoadPerformanceReport(extractKey: String) extends RedshiftLoadMessage
case class MediaMathLoadPerformanceReport(extractKey: String) extends RedshiftLoadMessage

sealed trait RedshiftComputeMessage
case class FullPerformanceReportComputeAndWebDbLoad(webJobsToNotify: Seq[Long]) extends RedshiftComputeMessage

/**
 * Workers
 */

object DataTransferWorker extends Worker1[DataTransferMessage, RedshiftLoadMessage] {
  val in = DataTransferQueue
  val out = RedshiftLoadQueue

  val flow = Flow[(DataTransferMessage, in.Ackable)].map {
    case (msg: AppNexusPullPerformanceReport, ackable) => AppNexusJobs.pullPerformanceReport(msg, ackable)
    case (msg: MediaMathPullPerformanceReport, ackable) => MediaMathJobs.pullPerformanceReport(msg, ackable)
  }
}