freevents.scala

package test

import java.util.UUID

import cats.arrow.NaturalTransformation
import cats.{Id, ~>, Monad}

import scala.annotation.tailrec
import scala.util.Random

object Events {
  case class Handled[E](e: E)

  sealed trait ES[E, A[_], R] {
    def flatMap[R2](f: R => ES[E, A, R2]): ES[E, A, R2] = FlatMap(this, f)

    def map[R2](f: R => R2): ES[E, A, R2] = FlatMap(this, f andThen (x => Pure(x)))

    @tailrec
    final def step: ES[E, A, R] = this match {
      case FlatMap(FlatMap(c, f), g) => c.flatMap(cc => f(cc).flatMap(g)).step
      case FlatMap(Pure(a), f) => f(a).step
      case x => x
    }

    def foldMap[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[R] = step.foldAfterStep(ai, ei)

    protected def foldAfterStep[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[R]

    def handleEvents(modelUpdate: PartialFunction[E, ES[Nothing, A, Unit]], eventListener: PartialFunction[E, ES[E, A, Unit]]): ES[Handled[E], A, R]

    def compile[E2, A2[_]](ainj: A ~> A2, einj: E => E2): ES[E2, A2, R] = {
      type ES2[X] = ES[E2, A2, X]
      foldMap[ES2](new (A ~> ES2) {
        override def apply[X](fa: A[X]) = Suspend(ainj(fa))
      }, e => Emit(einj(e)))
    }
  }

  implicit class ESHandled[E, A[_], R](es: ES[Handled[E], A, R]) {
    def run[M[_]](ai: A ~> M, storeEvent: E => M[Unit])(implicit M: Monad[M]): M[R] =
      es.foldMap(ai, eh => storeEvent(eh.e))
  }

  case class Pure[E, A[_], R](r: R) extends ES[E, A, R] {
    override def handleEvents(mu: PartialFunction[E, ES[Nothing, A, Unit]], el: PartialFunction[E, ES[E, A, Unit]]): ES[Handled[E], A, R] =
      Pure[Handled[E], A, R](r)

    protected def foldAfterStep[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[R] =
      M.pure(r)
  }

  case class Emit[E, A[_]](e: E) extends ES[E, A, Unit] {
    private def extendEvents(p: ES[Nothing, A, Unit]): ES[Handled[E], A, Unit] =
      p.compile[Handled[E], A](NaturalTransformation.id, identity)

    override def handleEvents(mu: PartialFunction[E, ES[Nothing, A, Unit]], el: PartialFunction[E, ES[E, A, Unit]]): ES[Handled[E], A, Unit] = {
      def applyMu: ES[Nothing, A, Unit] = if (mu.isDefinedAt(e)) mu(e) else ES.done
      def applyEl: ES[E, A, Unit] = if (el.isDefinedAt(e)) el(e) else ES.done

      Emit[Handled[E], A](Handled(e))
        .flatMap(_ => extendEvents(applyMu)
          .flatMap(_ => applyEl.handleEvents(mu, el)))
    }

    protected def foldAfterStep[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[Unit] =
      ei(e)
  }

  case class Suspend[E, A[_], R](a: A[R]) extends ES[E, A, R] {
    override def handleEvents(mu: PartialFunction[E, ES[Nothing, A, Unit]], el: PartialFunction[E, ES[E, A, Unit]]): ES[Handled[E], A, R] =
      Suspend[Handled[E], A, R](a)

    protected def foldAfterStep[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[R] =
      ai(a)
  }

  case class FlatMap[E, A[_], R1, R2](c: ES[E, A, R1], f: R1 => ES[E, A, R2]) extends ES[E, A, R2] {
    override def handleEvents(mu: PartialFunction[E, ES[Nothing, A, Unit]], el: PartialFunction[E, ES[E, A, Unit]]): ES[Handled[E], A, R2] = {
      FlatMap[Handled[E], A, R1, R2](c.handleEvents(mu, el), f andThen (_.handleEvents(mu, el)))
    }

    protected def foldAfterStep[M[_]](ai: A ~> M, ei: E => M[Unit])(implicit M: Monad[M]): M[R2] =
      M.flatMap(c.foldMap(ai, ei))(cc => f(cc).foldMap(ai, ei))
  }

  object ES {
    implicit def esMonad[E, A[_]]: Monad[({type x[X] = ES[E, A, X]})#x] = new Monad[({type x[X] = ES[E, A, X]})#x] {
      override def pure[X](x: X) = Pure(x)
      override def flatMap[X, Y](fa: ES[E, A, X])(f: X => ES[E, A, Y]) = fa.flatMap(f)
    }

    def pure[E, A[_], R](r: R): ES[E, A, R] = Pure(r)
    def done[E, A[_]]: ES[E, A, Unit] = pure(())
    def emit[E, A[_]](e: E): ES[E, A, Unit] = Emit(e)
    def suspend[E, A[_], R](a: A[R]): ES[E, A, R] = Suspend(a)
  }
}

object EventsExample extends App {
  import Events._

  object UserModule {
    case class User(id: Long, email: String, password: String)
    case class ApiKey(userId: Long, key: String)

    sealed trait Action[R]
    case class FindUserByEmail(email: String) extends Action[Option[User]]
    case class WriteUser(u: User) extends Action[Unit]
    case class FindApiKeyByUserId(userId: Long) extends Action[Option[ApiKey]]
    case class WriteApiKey(ak: ApiKey) extends Action[Unit]
    case class SendEmail(to: String, body: String) extends Action[Unit]

    sealed trait Event
    case class UserRegistered(u: User) extends Event
    case class ApiKeyCreated(ak: ApiKey) extends Event

    def pure[R](r: R): ES[Event, Action, R] = ES.pure(r)
    def emit(e: Event): ES[Event, Action, Unit] = ES.emit(e)
    def action[E, R](a: Action[R]): ES[E, Action, R] = ES.suspend(a)

    def registerUserCommand(email: String, password: String): ES[Event, Action, Either[String, Unit]] = {
      action(FindUserByEmail(email)).flatMap {
        case None => emit(UserRegistered(User(new Random().nextInt(), email, password))).map(_ => Right(()))
        case Some(user) => pure(Left("User with the given email already exists"))
      }
    }

    val modelUpdate: PartialFunction[Event, ES[Nothing, Action, Unit]] = {
      case UserRegistered(u) => action(WriteUser(u))
      case ApiKeyCreated(ak) => action(WriteApiKey(ak))
    }

    val eventListeners: PartialFunction[Event, ES[Event, Action, Unit]] = {
      case UserRegistered(u) => for {
        _ <- emit(ApiKeyCreated(ApiKey(u.id, UUID.randomUUID().toString)))
        _ <- action(SendEmail(u.email, "Welcome!"))
      } yield ()
    }
  }

  import UserModule._

  val handledCommand = registerUserCommand("adam@example.org", "1234").handleEvents(modelUpdate, eventListeners)

  val result: Either[String, Unit] = handledCommand.run[Id](new (Action ~> Id) {
    override def apply[A](fa: Action[A]) = fa match {
      case FindUserByEmail(email) => println(s"Find user by email: $email"); None
      case WriteUser(u) => println(s"Write user $u")
      case FindApiKeyByUserId(id) => println(s"Find api key by user id: $id"); None
      case WriteApiKey(ak) => println(s"Write api key: $ak")
      case SendEmail(to, body) => println(s"Send email to $to, body: $body")
    }
  }, e => println("Store event: " + e))

  println("Result: " + result)
}