package aoc.day8

import aoc._
import scala.collection.mutable
import scala.collection.immutable.NumericRange.Exclusive

enum Instruction { case Left, Right }

case class Node(name: String, edges: Array[String]):
  assert(edges.length == 2)
  def go(i: Instruction) = edges(i.ordinal)

object Parser extends CommonParser:
  import Instruction._
  val instruction = ("R" ^^ { _ => Right }) | ("L" ^^ { _ => Left })
  val instructions = rep(instruction)

  val name = """[\dA-Z]{3}""".r
  val node = name ~ "= (" ~ name ~ "," ~ name ~ ")" ^^ { case (name ~ _ ~ a ~ _ ~ b ~ _) =>
    Node(name, Array(a, b))
  }

// part 1

val instructions = Parser.parse(Parser.instructions, lines.next()).get.toArray
def nextOf(i: Int) = if i + 1 == instructions.length then 0 else i + 1
val map = lines
  .drop(1) // empty line
  .map(Parser.parse(Parser.node, _).get)
  .map(n => n.name -> n)
  .toMap

def part1 =
  val insRep =
    Iterator.unfold(0)(i => Some(instructions(i), nextOf(i)))
  val steps = insRep
    .scanLeft("AAA")((pos, ins) => map(pos).go(ins))
    .zipWithIndex
    .find(_._1 == "ZZZ")
    .map(_._2)
  println(steps.get)

// part 2
case class State(node: String, insPos: Int):
  def isGood = node.endsWith("Z")
  def next = State(map(node).go(instructions(insPos)), nextOf(insPos))

case class Loop(loopSize: Long, goodPositions: Set[Long]):
  def shiftBy(stepsTaken: Long) =
    val inLoop = stepsTaken % loopSize
    Loop(loopSize, goodPositions.map(v => (v + loopSize - inLoop) % loopSize))

  def scaleTo(newSize: Long) =
    assert(newSize % loopSize == 0)
    val range = Exclusive(0L, newSize, loopSize)
    Loop(newSize, goodPositions.flatMap(v => range.map(r => r + v)))

def getLoop(from: State) =
  val visited = mutable.Map[State, Long]()
  @scala.annotation.tailrec
  def visit(
      s: State,
      idx: Long
  ): (Long, Long) = // returns loop size and steps before loop
    visited.get(s) match
      case None =>
        visited += (s -> idx)
        visit(s.next, idx + 1)
      case Some(value) => (idx - value, value)
  val (loopSize, beforeLoop) = visit(from, 0)
  val goodPos = visited.toIterator
    .filter(_._2 >= beforeLoop)
    .filter(_._1.isGood)
    .map(_._2 - beforeLoop)
    .toSet
  (beforeLoop, Loop(loopSize, goodPos))

def part2 =
  val start =
    map.values.filter(_.name.endsWith("A")).map(n => State(n.name, 0)).toSeq
  val loopInfos = start.map(getLoop)
  val toManuallySimulate = loopInfos.map(_._1).max
  traverse(start, toManuallySimulate) match
    case Right(v) => println(v)
    case Left(states) =>
      val loops = loopInfos.map((before, loop) => loop.shiftBy(toManuallySimulate - before))
      val size = loops.foldLeft(1L)((v, loop) => lcm(v, loop.loopSize))
      if loops.forall(l => l.goodPositions == Set(l.loopSize - 3))
      then // dirty hack from input
        println(toManuallySimulate + size - 3)
      else
        val scaled = loops.map(_.scaleTo(size))
        val res = scaled
          .map(_.goodPositions)
          .foldLeft(scaled.head.goodPositions)((s, pos) => s.intersect(pos))
          .min + toManuallySimulate
        println(res)

def lcm(a: Long, b: Long) =
  @scala.annotation.tailrec
  def gcd(a: Long, b: Long): Long =
    if b == 0 then a
    else gcd(b, a % b)
  a / gcd(a, b) * b

def traverse(states: Seq[State], toSimulate: Long): Either[Seq[State], Long] =
  @scala.annotation.tailrec
  def go(states: Seq[State], steps: Long): Either[Seq[State], Long] =
    if states.forall(_.isGood) then Right(steps)
    else if steps == toSimulate then Left(states)
    else go(states.map(_.next), steps + 1)
  go(states, 0)

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