package aoc.day16

import aoc._
import scala.collection.mutable

enum Dir:
  import Dir._
  case Up, Down, Left, Right
  def apply(x: Int, y: Int) =
    val k = this.ordinal
    (x + dx(k), y + dy(k))

object Dir:
  private[Dir] val dx = Array(-1, 1, 0, 0)
  private[Dir] val dy = Array(0, 0, -1, 1)
import Dir._

enum Cell:
  case Slash, BSlash, Hori, Vert, Empty

  def apply(d: Dir) = this match
    case BSlash =>
      d match
        case Up    => Seq(Left)
        case Down  => Seq(Right)
        case Left  => Seq(Up)
        case Right => Seq(Down)
    case Slash =>
      d match
        case Up    => Seq(Right)
        case Down  => Seq(Left)
        case Left  => Seq(Down)
        case Right => Seq(Up)
    case Hori =>
      d match
        case Left | Right => Seq(d)
        case Up | Down    => Seq(Left, Right)
    case Vert =>
      d match
        case Up | Down    => Seq(d)
        case Left | Right => Seq(Up, Down)
    case Empty => Seq(d)

import Cell._

object Parser extends CommonParser:
  val cell =
    ('/' ^^^ Slash) | ('\\' ^^^ BSlash) | ('-' ^^^ Hori) | ('|' ^^^ Vert) | ('.' ^^^ Empty)
  val line = rep1(cell)

val board =
  lines
    .map(Parser.parse(Parser.line, _).get)
    .map(_.toArray)
    .toArray

def inside(x: Int, y: Int) =
  x >= 0 && x < board.length && y >= 0 && y < board(0).length

class Visitor:
  private val visited = mutable.Set.empty[(Int, Int, Dir)]

  private val queue = mutable.Queue.empty[(Int, Int, Dir)]

  def clear() =
    visited.clear()

  def go(x: Int, y: Int, dir: Dir): Unit =
    queue += ((x, y, dir))
    loop()

  @scala.annotation.tailrec
  private def loop(): Unit =
    if queue.isEmpty then ()
    else
      val (x, y, dir) = queue.dequeue()
      visit(x, y, dir)
      loop()

  private def visit(x: Int, y: Int, dir: Dir): Unit =
    if visited.contains((x, y, dir)) then return ()
    visited += ((x, y, dir))
    val cell = board(x)(y)
    cell(dir)
      .map(k => (k(x, y), k))
      .filter { case ((x, y), _) => inside(x, y) }
      .foreach { case ((x, y), k) => queue.enqueue((x, y, k)) }

  def visitedCells =
    visited
      .map((x, y, dir) => (x, y))

// part 1

def part1 =
  val v = Visitor()
  v.go(0, 0, Right)
  println(v.visitedCells)
  println(v.visitedCells.size)

// part 2

def part2 =
  val v = Visitor()
  val choices = (0 until board.length).flatMap { row =>
    Seq((row, 0, Right), (row, board(0).length - 1, Left))
  } ++
    (0 until board(0).length).flatMap { col =>
      Seq((0, col, Down), (board.length - 1, col, Up))
    }
  val res = choices.map { case (x, y, dir) =>
    v.clear()
    v.go(x, y, dir)
    v.visitedCells.size
  }.max
  println(res)

@main def Day16(part: Int) = part match
  case 1 => part1
  case 2 => part2