package snake

import SafeIO.*

class M07bSafeIOSuite extends munit.FunSuite:
  private def rectAt(x: Int): Shape = Shape.Rect(x, 0, 1, 1, "#000000")
  private val r1 = rectAt(0)
  private val t1 = Shape.Text(1, 2, "hi", "#ffffff", 16, "left")
  private def drawOp(shape: Shape): Op = Op.Draw(shape)
  private val seedOp: Op = Op.RandomSeed
  private def boom[A, B](a: A): B = throw new RuntimeException("boom")
  private val N = 100000

  // The value of a RandomSeed answer (the tests' interpreters always answer it with a Seed).
  private def seedOf(r: OpResult): Long = r match
    case OpResult.Seed(value) => value
    case _                    => -1L // Done, or Module 10's Best: never the answer to a RandomSeed

  // Which constructor a program is. Functions cannot be compared, so the tests check a node's shape by its tag and
  // then observe the program through `trace`. (A tag, not the program itself: printing a 100 000-deep chain would
  // overflow the stack.)
  private def tag(p: SafeIO[Any]): String = p match
    case Pure(_)       => "Pure"
    case Effect(_, _)  => "Effect"
    case FlatMap(_, _) => "FlatMap"

  // --- building -------------------------------------------------------------------------------------------------

  test("pure builds a Pure") { assertEquals(pure(3), Pure(3)) }

  test("effect(op) is an Effect whose continuation is pure") {
    effect(seedOp) match
      case Effect(op, k) =>
        assertEquals(op, seedOp)
        assertEquals(k(OpResult.Seed(5)), Pure(OpResult.Seed(5)))
      case other => fail(s"expected an Effect, got ${tag(other)}")
  }

  test("flatMap only builds a FlatMap node") {
    val m = pure(1)
    val p = m.flatMap(n => pure(n + 1))
    p match
      case FlatMap(io, _) => assertEquals[SafeIO[Any], SafeIO[Any]](io, m)
      case other          => fail(s"expected a FlatMap, got ${tag(other)}")
    assertEquals(trace(p, 0), Traced(2, Nil))
  }

  test("building flatMap(pure(1), throwingF) does not throw") {
    assertEquals(tag(pure(1).flatMap(boom[Int, SafeIO[Int]])), "FlatMap")
  }

  test("building map(io, throwingF) does not throw") {
    assertEquals(tag(effect(seedOp).map(boom[OpResult, Int])), "FlatMap")
  }

  test("tracing map(io, throwingF) runs f and throws") {
    interceptMessage[RuntimeException]("boom")(trace(effect(seedOp).map(boom[OpResult, Int]), 0))
  }

  // --- step: the three rewrite rules ----------------------------------------------------------------------------

  // An Effect that reads the seed and produces it as a Long.
  private val readSeed: SafeIO[Long] = Effect(seedOp, r => Pure(seedOf(r)))

  test("step leaves a Pure unchanged") {
    val p = pure(1)
    assertEquals(step(p), p)
  }

  test("step leaves an Effect unchanged") {
    val p = effect(seedOp)
    assertEquals(step(p), p)
  }

  test("step of FlatMap(Pure(a), f) is f(a)") {
    assertEquals(step(pure(3).flatMap(n => pure(n * 2))), Pure(6))
  }

  test("step of FlatMap(Effect(op, k), f) is an Effect with the same op that continues with FlatMap(k(r), f)") {
    step(readSeed.flatMap(n => pure(n * 2))) match
      case Effect(op, k) =>
        assertEquals(op, seedOp)
        val next = k(OpResult.Seed(5))
        next match
          case FlatMap(io, _) => assertEquals[SafeIO[Any], SafeIO[Any]](io, Pure(5L))
          case other          => fail(s"expected a FlatMap, got ${tag(other)}")
        assertEquals(trace(next, 0), Traced(10L, Nil))
      case other => fail(s"expected an Effect, got ${tag(other)}")
  }

  test("step re-associates FlatMap(FlatMap(m, g), f) into FlatMap(m, x => FlatMap(g(x), f))") {
    val m = readSeed
    val s = step(m.flatMap(n => pure(n + 1)).flatMap(n => pure(n * 2)))
    s match
      case FlatMap(io, _) => assertEquals[SafeIO[Any], SafeIO[Any]](io, m)
      case other          => fail(s"expected a FlatMap, got ${tag(other)}")
    assertEquals(trace(s, 4), Traced(10L, List(seedOp)))
  }

  test("step does not call the functions it re-associates") {
    assertEquals(tag(step(effect(seedOp).flatMap(boom[OpResult, SafeIO[Int]]).flatMap(boom[Int, SafeIO[Int]]))), "FlatMap")
  }

  test("step of a left-nested chain is one rewrite, not a whole walk") {
    // After one step, the top is still a FlatMap: step did constant work, the rest is left for later steps.
    val chain = (1 to N).foldLeft(pure(0))((acc, _) => acc.flatMap(n => pure(n + 1)))
    assertEquals(tag(step(chain)), "FlatMap")
  }

  // --- the interpreter ------------------------------------------------------------------------------------------

  test("trace of pure(3) has result 3 and no ops") {
    assertEquals(trace(pure(3), 0), Traced(3, Nil))
  }

  test("trace answers RandomSeed with the seed and Draw with Done") {
    assertEquals(
      trace(effect(seedOp).flatMap(s => effect(drawOp(r1)).map(_ => s)), 9),
      Traced(OpResult.Seed(9), List(seedOp, drawOp(r1)))
    )
    assertEquals(trace(effect(drawOp(r1)), 9), Traced(OpResult.Done, List(drawOp(r1))))
  }

  test("trace refuses operations of later modules") {
    intercept[IllegalArgumentException](trace(effect(Op.LoadBest), 0))
  }

  private def drawing(shape: Shape): SafeIO[Unit] = effect(drawOp(shape)).map(_ => ())

  test("sequence of nothing has no ops") {
    assertEquals(trace(sequence(Nil), 0), Traced((), Nil))
  }

  test("sequence runs effects left to right") {
    assertEquals(
      trace(sequence(List(effect(seedOp).map(_ => ()), drawing(t1), drawing(r1))), 0).ops,
      List(seedOp, drawOp(t1), drawOp(r1))
    )
  }

  // --- stack safety ---------------------------------------------------------------------------------------------

  private def leftDraws(n: Int): SafeIO[Any] =
    (0 until n).foldLeft[SafeIO[Any]](pure(()))((acc, i) => acc.flatMap(_ => effect(drawOp(rectAt(i)))))
  private def rightDraws(n: Int): SafeIO[Any] =
    (0 until n).foldRight[SafeIO[Any]](pure(()))((i, rest) => effect(drawOp(rectAt(i))).flatMap(_ => rest))

  test("100 000 left-nested flatMaps of draws run, in order") {
    val Traced(result, ops) = trace(leftDraws(N), 0)
    assertEquals(result, OpResult.Done)
    assertEquals(ops.length, N)
    assertEquals((ops.head, ops.last), (drawOp(rectAt(0)), drawOp(rectAt(N - 1))))
  }

  test("100 000 right-nested flatMaps of draws run, in order") {
    val ops = trace(rightDraws(N), 0).ops
    assertEquals(ops.length, N)
    assertEquals((ops.head, ops.last), (drawOp(rectAt(0)), drawOp(rectAt(N - 1))))
  }

  test("100 000 left-nested maps run without any op") {
    assertEquals(trace((1 to N).foldLeft(pure(0))((acc, _) => acc.map(_ + 1)), 0), Traced(N, Nil))
  }

  test("sequence of 100 000 draws runs") {
    assertEquals(trace(sequence(List.tabulate(N)(i => drawing(rectAt(i)))), 0).ops.length, N)
  }

  test("a while loop over step is a complete interpreter for 100 000 left-nested flatMaps") {
    // What an impure runtime would do (a loop is allowed there): rewrite until the top is Pure or Effect, perform.
    var performed = List.empty[Op]
    var current = leftDraws(N)
    while tag(current) != "Pure" do
      current = step(current)
      current match
        case Effect(op, k) =>
          performed = op :: performed
          current = k(OpResult.Done)
        case _ => ()
    assertEquals(performed.length, N)
  }

  // --- the monad laws, observed through trace -------------------------------------------------------------------

  private val programs: List[(String, () => SafeIO[Long])] = List(
    ("a pure value", () => pure(4L)),
    ("one effect", () => readSeed),
    ("a draw followed by reading the seed", () => effect(drawOp(r1)).flatMap(_ => readSeed))
  )
  private def f(n: Long): SafeIO[Long] = effect(drawOp(rectAt(n.toInt))).flatMap(_ => pure(n * 2))
  private def g(n: Long): SafeIO[Long] = effect(seedOp).map(s => seedOf(s) + n)

  programs.foreach { (name, program) =>
    test(s"left identity holds for $name") {
      val a = trace(program(), 5).result
      assertEquals(trace(pure(a).flatMap(f), 5), trace(f(a), 5))
    }
    test(s"right identity holds for $name") {
      assertEquals(trace(program().flatMap(pure), 5), trace(program(), 5))
    }
    test(s"associativity holds for $name") {
      assertEquals(trace(program().flatMap(f).flatMap(g), 5), trace(program().flatMap(x => f(x).flatMap(g)), 5))
    }
  }

  // --- the same programs under Module 7's IO --------------------------------------------------------------------

  // Each program is written once against this interface and built with both IOs.
  private trait Builder[F[_]]:
    def pure[A](a: A): F[A]
    def effect(op: Op): F[OpResult]
    def map[A, B](io: F[A])(f: A => B): F[B]
    def flatMap[A, B](io: F[A])(f: A => F[B]): F[B]
    def sequence(ios: List[F[Unit]]): F[Unit]
    def trace[A](io: F[A], seed: Long): Traced[A]
    def perform(op: Op): F[Unit] = map(effect(op))(_ => ())

  private object Safe extends Builder[SafeIO]:
    def pure[A](a: A): SafeIO[A] = SafeIO.pure(a)
    def effect(op: Op): SafeIO[OpResult] = SafeIO.effect(op)
    def map[A, B](io: SafeIO[A])(f: A => B): SafeIO[B] = io.map(f)
    def flatMap[A, B](io: SafeIO[A])(f: A => SafeIO[B]): SafeIO[B] = io.flatMap(f)
    def sequence(ios: List[SafeIO[Unit]]): SafeIO[Unit] = SafeIO.sequence(ios)
    def trace[A](io: SafeIO[A], seed: Long): Traced[A] = SafeIO.trace(io, seed)

  private object Simple extends Builder[IO]:
    def pure[A](a: A): IO[A] = IO.pure(a)
    def effect(op: Op): IO[OpResult] = IO.effect(op)
    def map[A, B](io: IO[A])(f: A => B): IO[B] = io.map(f)
    def flatMap[A, B](io: IO[A])(f: A => IO[B]): IO[B] = io.flatMap(f)
    def sequence(ios: List[IO[Unit]]): IO[Unit] = IO.sequence(ios)
    def trace[A](io: IO[A], seed: Long): Traced[A] = IO.trace(io, seed)

  // A program for any Builder, traced with seed 7.
  private trait Shared:
    def traced[F[_]](M: Builder[F]): Traced[?]

  private val shared: List[(String, Shared)] = List(
    ("seed, then a draw that depends on it", new Shared:
      def traced[F[_]](M: Builder[F]): Traced[?] =
        M.trace(M.flatMap(M.map(M.effect(seedOp))(seedOf))(s => M.map(M.effect(drawOp(rectAt((s % 20).toInt))))(_ => s)), 7)
    ),
    ("a sequence of three effects", new Shared:
      def traced[F[_]](M: Builder[F]): Traced[?] =
        M.trace(M.sequence(List(M.perform(drawOp(t1)), M.perform(seedOp), M.perform(drawOp(r1)))), 7)
    ),
    ("200 left-nested flatMaps", new Shared:
      def traced[F[_]](M: Builder[F]): Traced[?] =
        M.trace((0 until 200).foldLeft(M.pure(0))((acc, i) => M.flatMap(acc)(n => M.map(M.effect(drawOp(rectAt(i))))(_ => n + 1))), 7)
    ),
    ("nested on both sides", new Shared:
      def traced[F[_]](M: Builder[F]): Traced[?] =
        val left = M.flatMap(M.map(M.effect(seedOp))(seedOf))(s => M.flatMap(M.effect(drawOp(r1)))(_ => M.pure(s + 1)))
        M.trace(M.flatMap(left)(n => M.map(M.sequence(List(M.perform(drawOp(t1)), M.perform(seedOp))))(_ => n * 3)), 7)
    )
  )

  shared.foreach { (name, program) =>
    test(s"trace equals Module 7's trace for $name") {
      assertEquals(program.traced(Safe), program.traced(Simple))
    }
  }