package aoc.day23

import aoc._
import aoc.direction._
import scala.collection.mutable

enum Cell:
  case Path, Forest
  case Slope(dir: Dir)
import Cell._

object Parser extends CommonParser:
  import Dir._
  val slope = ('^' ^^^ Up | 'v' ^^^ Down | '<' ^^^ Left | '>' ^^^ Right) ^^ (Slope(_))
  val cell = '.' ^^^ Path | '#' ^^^ Forest | slope

  val line = rep(cell) ^^ (_.toArray)

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

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

extension [K](m: mutable.Set[K])
  inline def withFound[T](k: K)(inline body: => T): T =
    m += k
    val res = body
    m -= k
    res

// part 1

object longestPath:
  var bestPathFound = 0
  val visited = mutable.Set.empty[(Int, Int)]
  val intersections = mutable.Map.empty[(Int, Int), Seq[((Int, Int), Int)]]
  val cache = mutable.Map.empty[(Long, Int, Int), Option[Int]]

  def options(followSlopes: Boolean)(x: Int, y: Int) =
    val nexts = board(x)(y) match
      case Path       => Dir.all.iterator
      case Forest     => Iterator.empty
      case Slope(dir) => if followSlopes then Iterator(dir) else Dir.all.iterator
    nexts
      .map(_(x, y))
      .filter(inside(_, _))
      .filter(board(_)(_) != Forest)
      .toSeq

  @scala.annotation.tailrec
  def walkUntilIntersection(
      followSlopes: Boolean
  )(last: (Int, Int), x: Int, y: Int, accum: Int): Option[((Int, Int), Int)] =
    if x == board.size - 1 then return Some(((x, y), accum))
    val options = this.options(followSlopes)(x, y).filter(_ != last)
    options match
      case Seq()         => None
      case Seq((nx, ny)) => walkUntilIntersection(followSlopes)((x, y), nx, ny, accum + 1)
      case _             => Some(((x, y), accum))

  def getChoices(followSlopes: Boolean)(x: Int, y: Int) =
    intersections.getOrElseUpdate(
      (x, y), {
        this.options(followSlopes)(x, y).flatMap(walkUntilIntersection(followSlopes)((x, y), _, _, 1))
      }
    )

  def apply(followSlopes: Boolean)(
      x: Int,
      y: Int,
      accum: Int = 0
  ): Option[Int] =
    if x == board.size - 1 then Some(accum)
    else
      getChoices(followSlopes)(x, y) match
        case Seq()   => None
        case options =>
          // println(s"$x $y $accum $options")
          visited.withFound((x, y)):
            val res = options
              .filter((p, _) => !visited.contains(p))
              .flatMap:
                case ((nx, ny), dist) =>
                  longestPath(followSlopes)(nx, ny, accum + dist)
              .maxOption
            res

def part1 =
  println(longestPath(true)(0, board(0).indexWhere(_ != Forest)))

def part2 =
  println(longestPath(false)(0, board(0).indexWhere(_ != Forest)))

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