package snake

import Fixtures.*
import Check.*

class M12CheckSuite extends munit.FunSuite:

  // ---- generators

  test("elementOf picks the element at the index drawn by intBelow") {
    val xs = List("a", "b", "c", "d", "e")
    val (i, seed) = Random.intBelow(5).run(0)
    assertEquals(elementOf(xs).run(0), (xs(i), seed))
  }

  test("elementOf of a one-element array always gives that element, and still uses up a draw") {
    assertEquals(elementOf(List("only")).run(0), ("only", 1831565813L))
  }

  test("listOf(gen, 0) is the empty list and leaves the seed alone") {
    assertEquals(listOf(Random.nextUint32, 0).run(99), (List.empty[Long], 99L))
  }

  test("listOf threads the seed from each draw to the next") {
    assertEquals(listOf(Random.nextUint32, 3).run(0), (List(1144304738L, 1416247L, 958946056L), 1199730143L))
  }

  test("listOf of a constant generator repeats the value and leaves the seed alone") {
    assertEquals(listOf(Random.pure("x"), 4).run(7), (List("x", "x", "x", "x"), 7L))
  }

  test("genActions(n) gives n actions, the same ones for the same seed") {
    val (actions, _) = genActions(30).run(5)
    assertEquals(actions.length, 30)
    assert(actions.forall {
      case Action.Tick | Action.Turn(_) => true
      case _                            => false
    })
    assertEquals(genActions(30).run(5), genActions(30).run(5))
  }

  // ---- forAll

  test("forAll reports success with the number of runs") {
    assertEquals(forAll(Random.intBelow(10), 200, 1)(x => x >= 0 && x < 10), CheckResult.Passed(200))
  }

  test("forAll with zero runs succeeds without drawing anything") {
    assertEquals(forAll(Random.nextUint32, 0, 1)(_ => false), CheckResult.Passed(0))
  }

  test("forAll reports the first failing draw, its value and the seed it started from") {
    val gen = Random.intBelow(10)
    val prop = (x: Int) => x != 7
    val (draws, _) = listOf(gen, 200).run(3)
    val first = draws.indexWhere(x => !prop(x))
    assert(first > 0, "the test needs a failure that is not the very first draw")
    forAll(gen, 200, 3)(prop) match
      case CheckResult.Failed(run, seed, value) =>
        assertEquals((run, value), (first, 7))
        assertEquals(gen.run(seed)._1, 7)
      case other => fail(s"expected a failure, got $other")
  }

  test("forAll stops at the first failure: the property is not asked about later values") {
    val gen = Random.intBelow(10)
    val (draws, _) = listOf(gen, 200).run(3)
    val first = draws.indexWhere(_ == 7)
    var calls = 0 // test code may count; the code under test may not
    val result = forAll(gen, 200, 3) { x => calls = calls + 1; x != 7 }
    result match
      case CheckResult.Failed(run, _, _) => assertEquals(run, first)
      case other                         => fail(s"expected a failure, got $other")
    assertEquals(calls, first + 1)
  }

  test("forAll is reproducible: same generator, property, runs and seed give the same report") {
    assertEquals(forAll(Random.intBelow(1000), 500, 11)(_ % 97 != 0), forAll(Random.intBelow(1000), 500, 11)(_ % 97 != 0))
  }

  // ---- shrinking

  test("removals lists every way to drop one element, in order") {
    assertEquals(removals(List(1, 2, 3)), List(List(2, 3), List(1, 3), List(1, 2)))
  }

  test("removals of an empty list is empty, and of a single element is one empty list") {
    assertEquals(removals(List.empty[Int]), List.empty[List[Int]])
    assertEquals(removals(List("a")), List(List.empty[String]))
  }

  test("minimize keeps exactly the elements the failure needs") {
    assertEquals(minimize(List(1, 2, 3, 4, 5, 6, 7, 8, 9, 10))(xs => xs.contains(3) && xs.contains(7)), List(3, 7))
  }

  test("minimize keeps the order and the duplicates the failure needs") {
    assertEquals(
      minimize(List("b", "a", "c", "a", "d", "a"))(xs => xs.count(_ == "a") >= 2 && xs.contains("c")),
      List("c", "a", "a")
    )
  }

  test("minimize of a list that fails whatever you remove is the empty list") {
    assertEquals(minimize(List(1, 2, 3))(_ => true), List.empty[Int])
  }

  test("minimize returns a list that still fails and has no failing removal (1-minimal)") {
    val fails = (xs: List[Int]) => xs.sum >= 10
    val small = minimize(List(5, 1, 4, 2, 6, 3))(fails)
    assert(fails(small))
    assert(removals(small).forall(ys => !fails(ys)))
  }

  // ---- invariants

  cases("m12-violations").foreach { c =>
    test(s"violations: ${c("name").str}") {
      assertEquals(violations(gameState(c("args")(0))), list(c("expected"))(_.str))
    }
  }

  // ---- moveGaps

  // A whole state whose snake has its head at (x, y); moveGaps only looks at the head.
  private def at(x: Int, y: Int): GameState =
    GameState(20, 20, List(Point(x, y), Point(x - 1, y)), Direction.Right, Direction.Right, None, 0L, 0, Status.Playing,
      0, None)

  test("moveGaps measures the distance between consecutive moves of the head") {
    assertEquals(moveGaps(List(at(1, 0), at(1, 0), at(2, 0), at(2, 0), at(2, 0), at(3, 0), at(3, 1))), List(3, 1))
  }

  test("moveGaps of fewer than two moves is empty") {
    assertEquals(moveGaps(Nil), List.empty[Int])
    assertEquals(moveGaps(List(at(1, 0))), List.empty[Int])
    assertEquals(moveGaps(List(at(1, 0), at(1, 0), at(2, 0), at(2, 0))), List.empty[Int])
  }

  test("moveGaps counts a move at index 1") {
    assertEquals(moveGaps(List(at(1, 0), at(2, 0), at(2, 0), at(3, 0))), List(2))
  }

  test("moveGaps treats an empty snake as a head of its own") {
    def withSnake(snake: List[Point]): GameState = at(0, 0).copy(snake = snake)
    val states = List(Nil, Nil, List(Point(1, 0)), List(Point(1, 0)), List(Point(2, 0))).map(withSnake)
    assertEquals(moveGaps(states), List(2))
  }

  // ---- the payoff: properties of the game, and deliberately broken games ("mutants") that they must catch

  private type Update = (GameState, Action) => GameState

  // Every state of the game played with `upd` from `start`. (For the real `update` this is `history`.)
  private def statesWith(upd: Update, start: GameState, actions: List[Action]): List[GameState] =
    actions.scanLeft(start)(upd)
  private def sane(upd: Update)(actions: List[Action]): Boolean =
    statesWith(upd, Spawn.newGame(1), actions).forall(s => violations(s).isEmpty)

  test("statesWith(update, ...) in these tests is the history of Module 5") {
    val (actions, _) = genActions(25).run(9)
    assertEquals(statesWith(Game.update, Spawn.newGame(1), actions), Game.history(Spawn.newGame(1), actions))
  }

  test("the real game keeps every invariant on 300 random games of 40 actions") {
    assertEquals(forAll(genActions(40), 300, 2024)(sane(Game.update)), CheckResult.Passed(300))
  }

  // Mutant 1: `turn` checks against nextDir instead of dir (version A of Puzzle 2 in Module 3).
  private def turnAgainstNextDir(state: GameState, action: Action): GameState = action match
    case Action.Turn(dir) =>
      if state.status == Status.Playing && dir != Model.opposite(state.nextDir) then state.copy(nextDir = dir)
      else state
    case other => Game.update(state, other)

  test("the quick double turn bug is found and shrunk to its two key presses") {
    val prop = sane(turnAgainstNextDir)
    forAll(genActions(40), 300, 2024)(prop) match
      case CheckResult.Failed(_, _, value) =>
        minimize(value)(actions => !prop(actions)) match
          case List(Action.Turn(first), Action.Turn(second)) =>
            assert(List(Direction.Up, Direction.Down).contains(first))
            assertEquals(second, Direction.Left)
          case other => fail(s"expected two turns, got $other")
      case other => fail(s"expected a failure, got $other")
  }

  // Mutant 2: `advance` forgets to check the walls.
  private def advanceWithoutWalls(state: GameState): GameState =
    if state.status != Status.Playing then state
    else
      val head = Rules.move(state.snake.head, state.nextDir)
      val eating = state.apple.contains(head)
      val body = if eating then state.snake else Util.dropLast(state.snake)
      if Util.contains(body, head, _ == _) then state.copy(dir = state.nextDir, status = Status.GameOver)
      else
        state.copy(
          snake = head :: body,
          dir = state.nextDir,
          apple = if eating then None else state.apple,
          score = state.score + (if eating then 1 else 0)
        )
  // What `update` does after a step since Exercise 10.6: take the queued turn. The mutants below that replace `Tick`
  // do the same, so that each one differs from the real game in exactly one thing.
  private def takeQueued(s: GameState): GameState = s.queued match
    case Some(dir) if s.status == Status.Playing => s.copy(nextDir = dir, queued = None)
    case _                                       => s
  private def noWalls(state: GameState, action: Action): GameState = action match
    case Action.Tick => takeQueued(Spawn.spawnApple(advanceWithoutWalls(state)))
    case other       => Game.update(state, other)

  test("a game without walls is found and shrunk to the shortest way off the board: ten ticks") {
    val prop = sane(noWalls)
    forAll(genActions(40), 300, 2024)(prop) match
      case CheckResult.Failed(_, _, value) =>
        assertEquals(minimize(value)(actions => !prop(actions)), List.fill(10)(Action.Tick))
      case other => fail(s"expected a failure, got $other")
  }

  // Mutant 4: a `Tick` that forgets the turn queue of Exercise 10.6 (the `Frame` step still takes it).
  private def tickForgetsQueue(state: GameState, action: Action): GameState = action match
    case Action.Tick => Spawn.tick(state)
    case other       => Game.update(state, other)

  test("a Tick that forgets the queued turn is caught by queued-not-reversed and shrunk to four actions") {
    val prop = sane(tickForgetsQueue)
    forAll(genActions(40), 300, 2024)(prop) match
      case CheckResult.Failed(_, _, value) =>
        assertEquals(
          minimize(value)(actions => !prop(actions)),
          List(Action.Turn(Direction.Up), Action.Turn(Direction.Right), Action.Tick, Action.Turn(Direction.Left))
        )
      case other => fail(s"expected a failure, got $other")
  }

  // Mutant 5: the queue accepts any key that differs from the pending turn, even the one straight back from it.
  private def queueAnyKey(state: GameState, action: Action): GameState = action match
    case Action.Turn(dir) if state.status == Status.Playing && state.nextDir != state.dir && dir != state.nextDir =>
      state.copy(queued = Some(dir))
    case other => Game.update(state, other)

  test("a queue that accepts a reversal is found and shrunk to its two key presses") {
    val prop = sane(queueAnyKey)
    forAll(genActions(40), 300, 2024)(prop) match
      case CheckResult.Failed(_, _, value) =>
        assertEquals(
          minimize(value)(actions => !prop(actions)),
          List(Action.Turn(Direction.Up), Action.Turn(Direction.Down))
        )
      case other => fail(s"expected a failure, got $other")
  }

  // Mutant 3: the frame rule this course used at first: step whenever the *total* frame count is a multiple of movesEvery.
  private def totalFrameCount(state: GameState, action: Action): GameState = action match
    case Action.Frame =>
      if state.status != Status.Playing then state
      else
        val counted = state.copy(frame = state.frame + 1)
        if counted.frame % Game.movesEvery(counted.score) == 0 then Spawn.tick(counted) else counted
    case other => Game.update(state, other)

  // A long corridor: the snake is one step from its third apple (score 2 -> 3, so the speed changes), `frame` frames into the game.
  private def corridor(frame: Int): GameState =
    GameState(60, 1, List(Point(1, 0), Point(0, 0)), Direction.Right, Direction.Right, Some(Point(2, 0)), 0L, 2,
      Status.Playing, frame, None)
  private val frames = List.fill(40)(Action.Frame)
  private def steady(upd: Update)(frame: Int): Boolean =
    val states = statesWith(upd, corridor(frame), frames)
    val scoreAtMove = states.zip(states.drop(1)).collect {
      case (before, after) if after.snake.head != before.snake.head => after.score
    }
    moveGaps(states).zip(scoreAtMove).forall((gap, score) => gap == Game.movesEvery(score))

  test("the real game moves exactly every movesEvery(score) frames, also across a speed-up") {
    assertEquals(forAll(Random.intBelow(8), 100, 7)(steady(Game.update)), CheckResult.Passed(100))
  }

  test("counting frames from the start of the game makes some steps too short, and the property finds one") {
    forAll(Random.intBelow(56), 100, 7)(steady(totalFrameCount)) match
      case CheckResult.Failed(_, _, value) => assertEquals(steady(totalFrameCount)(value), false)
      case other                           => fail(s"expected a failure, got $other")
  }

  // ---- random states (Exercise 12.6)

  private def connected(snake: List[Point]): Boolean =
    snake.zip(snake.drop(1)).forall((a, b) => math.abs(b.x - a.x) + math.abs(b.y - a.y) == 1)
  private def board(width: Int, height: Int): GameState = Model.initialState(0).copy(width = width, height = height)

  test("genSnake draws the head first, then walks from the tail (seed 0, 5x5, at most 6 cells)") {
    assertEquals(
      genSnake(5, 5, 6).run(0),
      (List(Point(3, 2), Point(3, 1), Point(3, 0), Point(2, 0), Point(1, 0), Point(0, 0)), 4231026099L)
    )
  }

  test("genSnake is pinned for a few more seeds") {
    assertEquals(
      genSnake(5, 5, 6).run(1)._1,
      List(Point(2, 3), Point(3, 3), Point(4, 3), Point(4, 4), Point(3, 4), Point(2, 4))
    )
    assertEquals(
      genSnake(5, 5, 6).run(42)._1,
      List(Point(1, 0), Point(1, 1), Point(2, 1), Point(2, 0), Point(3, 0), Point(3, 1))
    )
  }

  test("genSnake with maxLen 1 is just the head, drawn with two draws") {
    val (x, s1) = Random.intBelow(5).run(0)
    val (y, s2) = Random.intBelow(5).run(s1)
    assertEquals(genSnake(5, 5, 1).run(0), (List(Point(x, y)), s2))
  }

  test("genSnake stops early when the walk is stuck: a 1x1 board holds one cell") {
    assertEquals(genSnake(1, 1, 5).run(3)._1, List(Point(0, 0)))
  }

  test("genSnake always gives a snake on the board, connected, without overlap, at most maxLen long") {
    assertEquals(
      forAll(genSnake(4, 3, 8), 300, 5) { snake =>
        snake.length >= 1 && snake.length <= 8 && connected(snake)
        && snake.forall(p => Rules.inBounds(board(4, 3), p))
        && snake.distinct.length == snake.length
      },
      CheckResult.Passed(300)
    )
  }

  test("every generated state is sane: 500 states of a 5x5 board have no violations") {
    assertEquals(forAll(genState(5, 5), 500, 2024)(s => violations(s).isEmpty), CheckResult.Passed(500))
  }

  test("generated states are Playing, have nothing queued, and vary in length and apple") {
    val (states, _) = listOf(genState(5, 5), 500).run(2024)
    assert(states.forall(s => s.status == Status.Playing && s.queued.isEmpty && s.width == 5 && s.height == 5))
    assert(states.forall(s =>
      if s.snake.length == 1 then s.dir == Direction.Right else Rules.move(s.snake(1), s.dir) == s.snake.head
    ))
    assert(states.exists(_.snake.length == 1) && states.exists(_.snake.length >= 10))
    assert(states.exists(_.apple.isEmpty) && states.exists(_.apple.isDefined))
    assert(states.exists(s => s.nextDir != s.dir))
  }

  test("generated states start with seed 0, score 0 and frame 0") {
    val (states, _) = listOf(genState(5, 5), 50).run(2024)
    assert(states.forall(s => s.seed == 0L && s.score == 0 && s.frame == 0))
  }

  // ---- a second implementation as an oracle (Exercise 12.7)

  cases("m03-advance").foreach { c =>
    test(s"advanceSpec agrees with Module 3's advance fixture: ${c("name").str}") {
      assertEquals(advanceSpec(gameState(c("args")(0))), gameState(c("expected")))
    }
  }

  test("advance and advanceSpec agree on 500 generated states of a 5x5 board") {
    assertEquals(forAll(genState(5, 5), 500, 2024)(s => Rules.advance(s) == advanceSpec(s)), CheckResult.Passed(500))
  }

  // Mutant 6: Puzzle 1 of Module 3, "the new head must not be on the old snake" (it forgets that the tail moves away).
  private def advanceOldSnakeCheck(state: GameState): GameState =
    if state.status != Status.Playing then state
    else
      val head = Rules.move(state.snake.head, state.nextDir)
      val eating = state.apple.contains(head)
      val body = if eating then state.snake else Util.dropLast(state.snake)
      if !Rules.inBounds(state, head) || Util.contains(state.snake, head, _ == _) then
        state.copy(dir = state.nextDir, status = Status.GameOver)
      else
        state.copy(
          snake = head :: body,
          dir = state.nextDir,
          apple = if eating then None else state.apple,
          score = state.score + (if eating then 1 else 0)
        )

  test("no invariant sees the wrong collision check of Puzzle 1: every state it produces is sane") {
    assertEquals(
      forAll(genState(5, 5), 500, 2024)(s => violations(advanceOldSnakeCheck(s)).isEmpty),
      CheckResult.Passed(500)
    )
  }

  test("comparing with advanceSpec catches the wrong collision check: the head was moving onto the tail") {
    forAll(genState(5, 5), 500, 2024)(s => advanceOldSnakeCheck(s) == advanceSpec(s)) match
      case CheckResult.Failed(run, _, s) =>
        assertEquals(run, 21)
        assertEquals(Rules.move(s.snake.head, s.nextDir), s.snake.last)
        assertEquals(advanceOldSnakeCheck(s).status, Status.GameOver)
        assertEquals(advanceSpec(s).status, Status.Playing)
      case other => fail(s"expected a failure, got $other")
  }