package snake

/** What `Check.forAll` reports: every value passed, or the first one that did not (with the seed that draws it again). */
enum CheckResult[+A]:
  case Passed(runs: Int)
  case Failed(run: Int, seed: Long, value: A)

/**
 * m12 check: property-based testing built from the course's own parts.
 * A generator of test data is nothing new: it is a Random[A] from Random.scala (seed => (A, nextSeed)).
 */
object Check:

  /** One of the elements of the non-empty list `xs`: the one at index `intBelow(xs.length)`. */
  def elementOf[A](xs: List[A]): Random[A] = Random.intBelow(xs.length).map(xs(_))

  /** `n` values drawn from `gen`, one after another; the seed is threaded from each draw to the next. */
  def listOf[A](gen: Random[A], n: Int): Random[List[A]] =
    if n <= 0 then Random.pure(Nil)
    else gen.flatMap(a => listOf(gen, n - 1).map(rest => a :: rest))

  // Ticks appear three times so that the snake moves more often than it turns.
  private val actionPool: List[Action] = List(
    Action.Tick, Action.Tick, Action.Tick,
    Action.Turn(Direction.Up), Action.Turn(Direction.Down), Action.Turn(Direction.Left), Action.Turn(Direction.Right)
  )

  /** A generator of `n` random actions (given). */
  def genActions(n: Int): Random[List[Action]] = listOf(elementOf(actionPool), n)

  /**
   * Check `prop` on `runs` values drawn from `gen`, starting from `seed`; stop at the first value that fails.
   * Recursion carries the seed: no counter, no loop. `tailrec` makes the compiler check that it cannot overflow the stack.
   */
  def forAll[A](gen: Random[A], runs: Int, seed: Long)(prop: A => Boolean): CheckResult[A] =
    @annotation.tailrec
    def go(i: Int, s: Long): CheckResult[A] =
      if i >= runs then CheckResult.Passed(runs)
      else
        val (value, next) = gen.run(s)
        if prop(value) then go(i + 1, next) else CheckResult.Failed(i, s, value)
    go(0, seed)

  /** Every list that is `xs` with exactly one element removed, in the order of the removed index. */
  def removals[A](xs: List[A]): List[List[A]] = xs.indices.toList.map(i => xs.take(i) ++ xs.drop(i + 1))

  /** Keep removing one element for as long as the list still fails; the result is 1-minimal. */
  def minimize[A](xs: List[A])(fails: List[A] => Boolean): List[A] =
    removals(xs).find(fails) match
      case Some(smaller) => minimize(smaller)(fails)
      case None          => xs

  private def neighbours(a: Point, b: Point): Boolean = math.abs(a.x - b.x) + math.abs(a.y - b.y) == 1

  // Each invariant: its name, and a predicate that is true when the state is fine.
  private val invariants: List[(String, GameState => Boolean)] = List(
    ("snake-not-empty", s => s.snake.nonEmpty),
    ("snake-in-bounds", s => s.snake.forall(p => Rules.inBounds(s, p))),
    ("snake-no-overlap", s => s.snake.distinct.length == s.snake.length),
    ("snake-connected", s => s.snake.zip(s.snake.drop(1)).forall((a, b) => neighbours(a, b))),
    ("apple-in-bounds", s => s.apple.forall(a => Rules.inBounds(s, a))),
    ("apple-not-on-snake", s => s.apple.forall(a => !Util.contains(s.snake, a, _ == _))),
    ("turn-not-reversed", s => s.nextDir != Model.opposite(s.dir)),
    ("queued-not-reversed", s => !s.queued.contains(Model.opposite(s.nextDir)))
  )

  /** The names of the invariants that `state` breaks, in the order of the table above. */
  def violations(state: GameState): List[String] =
    invariants.collect { case (name, holds) if !holds(state) => name }

  /** The numbers of states between consecutive moves of the snake's head. */
  def moveGaps(states: List[GameState]): List[Int] =
    val moved = states.zip(states.drop(1)).zipWithIndex.collect {
      case ((before, after), i) if after.snake.headOption != before.snake.headOption => i + 1
    }
    moved.zip(moved.drop(1)).map((a, b) => b - a)

  // ---- random states (Exercise 12.6) and a second implementation of the step as an oracle (Exercise 12.7)

  private val directions = List(Direction.Up, Direction.Down, Direction.Left, Direction.Right)

  // `Rules.inBounds` wants a whole state; here there is only a board size.
  private def onBoard(width: Int, height: Int, p: Point): Boolean = p.x >= 0 && p.x < width && p.y >= 0 && p.y < height

  // Extend `snake` at its tail by up to `left` more cells, each a random free in-bounds neighbour of the last cell.
  private def grow(width: Int, height: Int, snake: List[Point], left: Int): Random[List[Point]] =
    if left <= 0 then Random.pure(snake)
    else
      val options = directions
        .map(d => Rules.move(snake.last, d))
        .filter(p => onBoard(width, height, p) && !Util.contains(snake, p, _ == _))
      if options.isEmpty then Random.pure(snake)
      else elementOf(options).flatMap(p => grow(width, height, snake :+ p, left - 1))

  /** A random snake, head first: a random head, then a random walk that never leaves the board or crosses itself. */
  def genSnake(width: Int, height: Int, maxLen: Int): Random[List[Point]] =
    Random.intBelow(width).flatMap(x =>
      Random.intBelow(height).flatMap(y => grow(width, height, List(Point(x, y)), maxLen - 1))
    )

  // The direction the head moved in to get from the second cell to the first (Right for a one-cell snake).
  private def headingOf(snake: List[Point]): Direction = snake match
    case head :: second :: _ => directions.find(d => Rules.move(second, d) == head).getOrElse(Direction.Right)
    case _                   => Direction.Right

  /** A random sane Playing state on a `width` x `height` board (given). */
  def genState(width: Int, height: Int): Random[GameState] =
    Random.intBelow(width * height).flatMap { n =>
      genSnake(width, height, n + 1).flatMap { snake =>
        val dir = headingOf(snake)
        val state = GameState(width, height, snake, dir, dir, None, 0L, 0, Status.Playing, 0, None)
        val free = Spawn.freeCells(state)
        elementOf(directions.filter(_ != Model.opposite(dir))).flatMap { nextDir =>
          Random.intBelow(2).flatMap { coin =>
            val apple: Random[Option[Point]] =
              if coin == 0 || free.isEmpty then Random.pure(None) else elementOf(free).map(Some(_))
            apple.map(a => state.copy(nextDir = nextDir, apple = a))
          }
        }
      }
    }

  /**
   * The step of Module 3, written from the rules instead of from `advance`: build the snake the step would give,
   * then ask `violations` whether that snake is on the board and does not overlap itself.
   */
  def advanceSpec(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 snake = head :: (if eating then state.snake else Util.dropLast(state.snake))
      val broken = violations(state.copy(snake = snake)).exists(v => v == "snake-in-bounds" || v == "snake-no-overlap")
      if broken then state.copy(dir = state.nextDir, status = Status.GameOver)
      else
        state.copy(
          snake = snake,
          dir = state.nextDir,
          apple = if eating then None else state.apple,
          score = state.score + (if eating then 1 else 0)
        )