package snake

class M04StateSuite extends munit.FunSuite:
  // Small initial states for the "behaves like" checks: two States behave alike when they give the same pair from each.
  private val starts = List(0, 1, 7, 42)
  private def sameOnAll[A](st1: State[Int, A], st2: State[Int, A]): Unit =
    starts.foreach(s => assertEquals(st1.run(s), st2.run(s), s"from state $s"))

  // A counter step: gives the current number times 10 and adds 1 to the state. A def, so that a stub that throws
  // fails each test on its own instead of the whole suite at construction time.
  private def step: State[Int, Int] = State.get[Int].flatMap(n => State.put(n + 1).map(_ => n * 10))

  private def f(a: Int): State[Int, Int] = State.modify[Int](_ * 2).map(_ => a + 1)
  private def g(b: Int): State[Int, (Int, Int)] = State.get[Int].flatMap(s => State.pure((b, s)))

  test("State pure gives its value and leaves the state alone") {
    assertEquals(State.pure[Int, String]("x").run(99), ("x", 99))
  }

  test("State map applies f to the value and keeps the state from the underlying step") {
    assertEquals(step.map(_ + 1).run(4), (41, 5))
  }

  test("State flatMap threads the state through both steps") {
    assertEquals(step.flatMap(a => step.map(b => (a, b))).run(4), ((40, 50), 6))
  }

  test("State get gives the state as the value and leaves it alone") {
    assertEquals(State.get[String].run("board"), ("board", "board"))
  }

  test("State put replaces the state and gives unit") {
    assertEquals(State.put(5).run(99), ((), 5))
  }

  test("State modify applies f to the state and gives unit") {
    assertEquals(State.modify[Int](_ * 2).run(21), ((), 42))
  }

  test("State law: left identity, flatMap(pure(a), f) behaves like f(a)") {
    sameOnAll(State.pure[Int, Int](3).flatMap(f), f(3))
  }

  test("State law: right identity, flatMap(m, pure) behaves like m") {
    sameOnAll(step.flatMap(State.pure), step)
  }

  test("State law: associativity, the grouping of flatMap does not matter") {
    sameOnAll(step.flatMap(f).flatMap(g), step.flatMap(a => f(a).flatMap(g)))
  }

  test("State a counter hands out fresh labels in order") {
    val fresh = State.get[Int].flatMap(n => State.put(n + 1).map(_ => s"L$n"))
    val three = fresh.flatMap(a => fresh.flatMap(b => fresh.map(c => List(a, b, c))))
    assertEquals(three.run(0), (List("L0", "L1", "L2"), 3))
    assertEquals(three.run(7), (List("L7", "L8", "L9"), 10))
  }

  test("State Random is State with the seed as S: the dice give the same values and seeds") {
    // One draw written only with State's get and put: read the seed, compute, store the next seed.
    val draw = State.get[Long].flatMap { seed =>
      val (v, next) = Random.nextUint32.run(seed)
      State.put(next).map(_ => v)
    }
    val d6State = draw.map(v => 1 + (v % 6))
    val twoDiceState = d6State.flatMap(a => d6State.map(b => a + b))
    val d6Random = Random.nextUint32.map(v => 1 + (v % 6))
    val twoDiceRandom = d6Random.flatMap(a => d6Random.map(b => a + b))
    // A Random step turned into a State by hand: the same function, in the other wrapper.
    val drawAsState = State[Long, Long](seed => Random.nextUint32.run(seed))
    List(0L, 1L, 42L, 4294967295L).foreach { seed =>
      assertEquals(draw.run(seed), Random.nextUint32.run(seed), s"draw from $seed")
      assertEquals(draw.run(seed), drawAsState.run(seed), s"draw as a State from $seed")
      assertEquals(d6State.run(seed), d6Random.run(seed), s"d6 from $seed")
      assertEquals(twoDiceState.run(seed), twoDiceRandom.run(seed), s"twoDice from $seed")
    }
  }