package aoc
package day14

import scala.annotation.tailrec
import scala.collection.mutable

val dayNumber = "14"

@main def part1: Unit =
  println(part1(loadInput(dayNumber)))

@main def part2: Unit =
  println(part2(loadInput(dayNumber)))

def part1(input: String): String =
  val lines = input.split('\n')
  val border = lines.head.map(_ => '#').mkString + "##"
  val grid = lines
    .map(l => s"#$l#")
    .prepended(border)
    .appended(border)
    .map(_.to(mutable.ArraySeq))
  for
    y <- grid.indices
    x <- grid.head.indices
  do
    if grid(y)(x) != 'O' then ()
    else
      (y - 1 to (0, -1))
        .takeWhile(grid(_)(x) == '.')
        .lastOption
        .foreach(newY =>
          grid(y)(x) = '.'
          grid(newY)(x) = 'O'
        )

  grid
    .map(_.count(_ == 'O'))
    .reverse
    .zipWithIndex
    .map((count, row) => count * row)
    .sum
    .toString

def part2(input: String): String =
  val lines = input.split('\n')
  val border = lines.head.map(_ => '#').mkString + "##"
  val grid = lines
    .map(l => s"#$l#")
    .prepended(border)
    .appended(border)
    .map(_.to(mutable.ArraySeq))

  val coords =
    for
      y <- grid.indices
      x <- grid.head.indices
    yield (x, y)

  enum Directions(
      val coordinates: Seq[(Int, Int)],
      val rangeToEdge: (Int, Int) => Seq[(Int, Int)]
  ):
    case North
        extends Directions(
          coords,
          (x, y) => (y - 1 to (0, -1)).map((x, _))
        )
    case West
        extends Directions(
          coords.sortBy(_._1),
          (x, y) => (x - 1 to (0, -1)).map((_, y))
        )
    case South
        extends Directions(
          coords.reverse,
          (x, y) => (y + 1 until grid.length).map((x, _))
        )
    case East
        extends Directions(
          coords.sortBy(_._1).reverse,
          (x, y) => (x + 1 until grid.length).map((_, y))
        )

  def cycle =
    for d <- Directions.values do
      for (x, y) <- d.coordinates do
        if grid(y)(x) != 'O' then ()
        else
          d.rangeToEdge(x, y)
            .takeWhile((xR, yR) => grid(yR)(xR) == '.')
            .lastOption
            .foreach((xR, yR) =>
              grid(y)(x) = '.'
              grid(yR)(xR) = 'O'
            )

  def cycleUntilRepeat: (Int, Int) =
    val seen = mutable.Map.empty[String, Int]

    @tailrec def aux(iter: Int): (Int, Int) =
      cycle
      val gridString = grid.map(_.mkString).mkString("\n")
      seen.get(gridString) match
        case Some(prevIterations) => (iter, iter - prevIterations)
        case None =>
          seen.put(gridString, iter)
          aux(iter + 1)

    aux(1)

  val (repeatedIndex, repeatedLength) = cycleUntilRepeat
  val cyclesLeft = (1_000_000_000 - repeatedIndex) % repeatedLength

  for _ <- (0 until cyclesLeft) do cycle

  grid
    .map(_.count(_ == 'O'))
    .reverse
    .zipWithIndex
    .map((count, row) => count * (row))
    .sum
    .toString