package aoc.day10

import aoc._

type Coord = (Int, Int)

extension (c: Coord) def +(other: Coord): Coord = (c._1 + other._1, c._2 + other._2)

enum Direction:
  case Up, Down, Left, Right

  private val directions = Array[Coord]((-1, 0), (1, 0), (0, -1), (0, 1))

  def apply(current: Coord): Coord = current + directions(ordinal)

  inline def isHorizontal = this == Left || this == Right
  inline def isVertical = !isHorizontal

  def next(pipe: Char): Option[Direction] = Some((pipe, this) match
    case ('|', _) if isVertical   => this
    case ('-', _) if isHorizontal => this
    case ('L', Down)              => Right
    case ('L', Left)              => Up
    case ('7', Up)                => Left
    case ('7', Right)             => Down
    case ('J', Down)              => Left
    case ('J', Right)             => Up
    case ('F', Up)                => Right
    case ('F', Left)              => Down
    case _                        => return None
  )

object Direction:
  val all = Seq(Up, Down, Left, Right)
  def pipe(cur: Direction, next: Direction) =
    Seq('|', '-', '7', 'J', 'F', 'L').find(cur.next(_) == Some(next)).get

val board = lines.toArray

inline def cell(x: Coord) = board(x._1)(x._2)

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

// part1
val (loop, sPipe) = {
  val boardSize = board.map(_.length()).sum

  // look for S
  val sCoord = (0 until board.length).findMap { x =>
    board(x).zipWithIndex.findMap { case (chr, y) => if chr == 'S' then Some((x, y)) else None }
  }.get

  @scala.annotation.tailrec
  def findCycle(accum: Seq[Coord])(current: Coord, direction: Direction): Option[(Seq[Coord], Direction)] =
    if current == sCoord then Some((accum, direction))
    else if accum.length > boardSize then None
    else
      direction.next(cell(current)) match
        case Some(nx) =>
          val nxCoord = nx(current)
          if inside(nxCoord) then findCycle(current +: accum)(nxCoord, nx) else None
        case None => None
  // find the loop
  Direction.all
    .map(d => (d(sCoord), d))
    .filter(v => inside(v._1))
    .findMap((cur, dir) =>
      findCycle(Seq(sCoord))(cur, dir)
        .map((loop, toS) => (loop.toSet, Direction.pipe(toS, dir)))
    )
    .get
}

def part1 =
  println(loop.size / 2)

// part 2

@scala.annotation.tailrec
def loop(accum: Int)(x: Int, y: Int, insideLoop: Boolean): Int =
  if !inside((x, y)) then accum
  else if !loop.contains((x, y)) then loop(accum + (if insideLoop then 1 else 0))(x + 1, y + 1, insideLoop)
  else
    val r = if board(x)(y) == 'S' then sPipe else board(x)(y)
    loop(accum)(
      x + 1,
      y + 1,
      if Seq('|', '-', 'J', 'F').contains(r) then !insideLoop else insideLoop
    )

def part2 =
  val res =
    (0 until board.length).map(x => loop(0)(x, 0, false)).sum
      + (1 until board(0).length).map(y => loop(0)(0, y, false)).sum
  println(res)

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