use.scala

package use

import kyo.*

// --- usage --- //

trait Logger[-S]:
  def log(level: String)(message: String): Unit < S

trait Console[-S]:
  def read: String < S
  def printLine(line: String): Unit < S

val comp: String < (Use[Console] & Use[Logger]) =
  for
    console <- Use.get[Console]
    logger  <- Use.get[Logger]
    _       <- logger.log("Debug")("Start teletype example")
    _       <- console.printLine("Hello, World!")
    _       <- console.printLine("What is your name?")
    name    <- console.read
    _       <- console.printLine(s"Hello $name!")
  yield name

// --- handling --- //

trait LoggerAndConsole[-S] extends Logger[S], Console[S]

object sout extends LoggerAndConsole[Sync]:
  override def read: String < Sync                              = Console.readLine.orPanic
  override def printLine(line: String): Unit < Sync             = Console.printLine(line)
  override def log(level: String)(message: String): Unit < Sync = Console.printLine(s"[$level]: $message")

object noOp extends LoggerAndConsole[Any]:
  override def read: String < Any                              = "Pierre"
  override def printLine(line: String): Unit < Any             = ()
  override def log(level: String)(message: String): Unit < Any = ()

object noOpLog extends Logger[Any]:
  override def log(level: String)(message: String): Unit < Any = ()

object App extends KyoApp:
  // normal program
  run:
    Use.run(sout)(comp)

  // noOp
  run:
    Use.run(noOp)(comp)

  // handle Use[Log] first, then Use[Console] & Use[Log]
  run:
    comp.handle(Use.run(noOpLog), Use.run(sout))

// --- implementation --- //

import kyo.kernel.ContextEffect

// Custom effect that can be defined outside of Kyo
// It's based on ContextEffect (which is not an ArrowEffect)
// The same thing use by Env
//
// We can use R[Any] <: R[S] here, because when handled,
//   S will be added to the pending set.
sealed trait Use[+R[-_]] extends ContextEffect[TypeMap[R[Any]]]

object Use:
  private trait AnyT[-A]

  // We can erase R in the Tag because the implementation rely on a TypeMap that will do all the work
  private def erasedTag[R[-_]] = Tag[Use[AnyT]].asInstanceOf[Tag[Use[R]]]

  // R[Use[R]] will force the use of R to add Use[R] to the pending set ... that has already Use[R]
  // it also solves cases where the effect is lifted (in a Stream, in a (A < S) < S2, ...
  // so we don't lose the type marker for Use[R] if it's not use directly
  // (if you would have used Context function for that, it's an use case for Caprese)
  def use[R[-_]](using frame: Frame)[A, S1](f: R[Use[R]] => A < S1)(using tag: Tag[R[Any]]): A < (Use[R] & S1) =
    ContextEffect.suspendWith(erasedTag[R]): map =>
      f(map.asInstanceOf[TypeMap[R[Any]]].get(using tag))

  // run with R[S1] will provide R[?] to the A < (Use[R] & S2), and add S1 in the pending set
  def run[R[-_], S1](r: R[S1])[A, S2](a: A < (Use[R] & S2))(using tag: Tag[R[Any]], frame: Frame): A < (S1 & S2) =
    val env: TypeMap[R[Any]] = TypeMap(r.asInstanceOf[R[Any]])
    ContextEffect.handle(erasedTag[R], env, _.union(env))(a)

  def get[R[-_]](using frame: Frame, tag: Tag[R[Any]]): R[Use[R]] < Use[R] = use[R](identity)