package snake

import scala.annotation.tailrec

/**
 * m07b safeio (optional): a stack-safe IO. Same idea as m07's IO, plus one more constructor, so that `flatMap` only
 * builds data and no program shape can overflow the stack, however deeply its `flatMap`s are nested.
 *   - `Pure(a)`          done, with a value
 *   - `Effect(op, k)`    perform `op`, then continue with `k(result)`
 *   - `FlatMap(io, f)`   run `io`, then continue with `f` of its result. `io` produces some type `X` that only `f`
 *                        knows about: the case has its own type parameters, so `X` stays hidden inside the node.
 * `Op`, `OpResult` and `Traced` are the ones from IO.scala. Independent of the game and of `IO`.
 */
enum SafeIO[+A]:
  case Pure(a: A)
  case Effect(op: Op, k: OpResult => SafeIO[A])
  case FlatMap[X, B](io: SafeIO[X], f: X => SafeIO[B]) extends SafeIO[B]

  // Only builds a node: never calls f, not even when this is a Pure.
  def flatMap[B](f: A => SafeIO[B]): SafeIO[B] = FlatMap(this, f)

  def map[B](f: A => B): SafeIO[B] = flatMap(a => Pure(f(a)))

object SafeIO:

  def pure[A](a: A): SafeIO[A] = Pure(a)

  def effect(op: Op): SafeIO[OpResult] = Effect(op, pure)

  // A left fold on purpose: with SafeIO the nesting no longer matters.
  def sequence(ios: List[SafeIO[Unit]]): SafeIO[Unit] = ios.foldLeft(pure(()))((acc, io) => acc.flatMap(_ => io))

  // One rewrite towards a program whose top is a Pure or an Effect. Constant work, no recursion.
  //   FlatMap(Pure(a), f)          -> f(a)
  //   FlatMap(Effect(op, k), f)    -> Effect(op, r => FlatMap(k(r), f))
  //   FlatMap(FlatMap(m, g), f)    -> FlatMap(m, x => FlatMap(g(x), f))     (re-association)
  // Pure and Effect are already in that shape and come back unchanged.
  def step[A](io: SafeIO[A]): SafeIO[A] = io match
    case FlatMap(Pure(x), f)       => f(x)
    case FlatMap(Effect(op, k), f) => Effect(op, r => FlatMap(k(r), f))
    case FlatMap(FlatMap(m, g), f) => FlatMap(m, x => FlatMap(g(x), f))
    case other                     => other

  // Module 7's two operations only. Module 10's LoadBest and SaveBest share the `Op` type but are not part of this
  // module: a SafeIO program that uses them is a mistake, and saying so beats inventing an answer.
  private def answer(op: Op, seed: Long): OpResult = op match
    case Op.RandomSeed                => OpResult.Seed(seed)
    case Op.Draw(_)                   => OpResult.Done
    case Op.LoadBest | Op.SaveBest(_) => throw new IllegalArgumentException(s"SafeIO.trace does not interpret $op")

  // The interpreter as a loop: `loop` calls itself only as its very last action, which `@tailrec` turns into a jump,
  // so the stack stays flat. Every FlatMap is rewritten by `step`, every Effect answered; ops collect at the front.
  def trace[A](io: SafeIO[A], seed: Long): Traced[A] =
    @tailrec def loop(current: SafeIO[A], performed: List[Op]): Traced[A] = current match
      case Pure(a)       => Traced(a, performed.reverse)
      case Effect(op, k) => loop(k(answer(op, seed)), op :: performed)
      case other         => loop(step(other), performed)
    loop(io, Nil)