import scala.io.Source
import scala.util.chaining.*

object Day14 extends App:

  val input = Source
    .fromURL(getClass.getResource("day14Input.txt"))
    .mkString
    .split('\n')
    .toList

  enum Block:
    case Sand, Rock, Air
  object Block:
    val show: Block => String =
      case Sand => "o"
      case Rock => "#"
      case Air  => "."

  type Point = (Int, Int)
  extension (point: Point)
    def x = point._1
    def y = point._2

  final case class Grid(blocks: Map[Point, Block]):
    val blocksWithDefault = blocks.withDefaultValue(Block.Air)

    val lowestLevel = blocks.keySet.map(_.y).max

    def show =
      (for
        x <- blocks.keySet.map(_.x).min to blocks.keySet.map(_.x).max
        y <- blocks.keySet.map(_.y).min to blocks.keySet.map(_.y).max
      yield ((x, y), blocksWithDefault((x, y))))
        .groupBy(_._1.y)
        .toList
        .sortBy(_._1)
        .map(_._2.sortBy(_._1.x).map(p => Block.show(p._2)).mkString)
        .mkString("\n")

    def applyGravityToSand(sand: Point): Point =
      if blocksWithDefault(sand.x, sand.y + 1) == Block.Air then (sand.x, sand.y + 1)
      else if blocksWithDefault(sand.x - 1, sand.y + 1) == Block.Air then (sand.x - 1, sand.y + 1)
      else if blocksWithDefault(sand.x + 1, sand.y + 1) == Block.Air then (sand.x + 1, sand.y + 1)
      else sand

    def dropSand(sand: Point): Option[Grid] =
      val finalPosition = LazyList
        .iterate(sand)(applyGravityToSand)
        .sliding(2)
        .takeWhile {
          case LazyList(p1, p2) => p1 != p2 && p2.y <= lowestLevel
          case _                => false
        }
        .toList
        .lastOption
        .flatMap(_.lastOption)
        .getOrElse(sand)
      if finalPosition == sand || finalPosition.y >= lowestLevel then None
      else Some(Grid(blocks + ((finalPosition, Block.Sand))))
  end Grid

  val grid: Grid =
    input
      .flatMap(
        _.split(" -> ")
          .map(_.split(",")).map { case Array(x, y) => (x.toInt, y.toInt) }.sliding(2).flatMap {
            case Array((x1, y1), (x2, y2)) =>
              if x2 - x1 == 0 then
                for y <- y1 to y2 by (Math.signum(y2 - y1).toInt) yield ((x2, y), Block.Rock)
              else for x <- x1 to x2 by (Math.signum(x2 - x1).toInt) yield ((x, y2), Block.Rock)
          }
      )
      .toMap
      .pipe(Grid.apply)

  val part1 = LazyList
    .iterate(Option(grid))(_.flatMap(_.dropSand((500, 0))))
    .takeWhile(_.isDefined)
    .map(_.get)

  println("Part 1:")
  println(part1.last.show)
  println(part1.length - 1)

  val grid2 = Grid(
    grid.blocks ++ (
      for x <- 300 to 700
      yield ((x, grid.lowestLevel + 2), Block.Rock)
    )
  )

  val part2 = LazyList
    .iterate(Option(grid2))(_.flatMap(_.dropSand((500, 0))))
    .takeWhile(_.exists(_.blocksWithDefault(500, 0) != Block.Sand))
    .map(_.get)
  println("Part 2:")
  println(part2.last.show)
  println(part2.length)
end Day14