package aoc.day19

import aoc._

enum Rate:
  case X, M, A, S
import Rate._

case class StateT[T](s: Array[T]):
  assert(s.size == 4)
  def apply(r: Rate) = s(r.ordinal)
  inline def update(r: Rate)(inline upd: T => T) =
    val t = s.clone()
    t(r.ordinal) = upd(this(r))
    StateT(t)

type State = StateT[Int]

// [low, high)
case class Range(low: Int, high: Int):
  def isEmpty = low >= high
  inline def size = (high - low)
  inline def withMin(min: Int) = Range(low.max(min), high)
  inline def withMax(max: Int) = Range(low, high.min(max))
object Range:
  lazy val empty = Range(1, 4001)

type Ranges = StateT[Range]

extension (rs: Ranges)
  def isEmpty = rs.s.exists(_.isEmpty)

  def size = rs.s.map(_.size.toLong).product

  def passing(r: Rule) = r match
    case Rule.Lt(r, v, _) =>
      rs.update(r)(_.withMax(v))
    case Rule.Gt(r, v, _) =>
      rs.update(r)(_.withMin(v + 1))

  def failing(r: Rule) = r match
    case Rule.Lt(r, v, _) =>
      rs.update(r)(_.withMin(v))
    case Rule.Gt(r, v, _) =>
      rs.update(r)(_.withMax(v + 1))

enum Rule:
  case Lt(r: Rate, v: Int, andThenTarget: String)
  case Gt(r: Rate, v: Int, andThenTarget: String)

  def test(s: State) = this match
    case Lt(r, v, _) => s(r) < v
    case Gt(r, v, _) => s(r) > v

  inline def andThen = this match
    case Lt(_, _, andThen) => andThen
    case Gt(_, _, andThen) => andThen

case class Work(rules: Seq[Rule], andThen: String):
  def apply(s: State) =
    rules
      .find(_.test(s))
      .map(_.andThen)
      .getOrElse(andThen)

extension (m: Map[String, Work])
  @scala.annotation.tailrec
  def test(cur: String)(s: State): Boolean =
    cur match
      case "A" => true
      case "R" => false
      case _   => test(m(cur)(s))(s)

object Parser extends CommonParser:
  val name = """[a-zA-Z]+""".r
  val rate =
    'x' ^^^ X | 'm' ^^^ M | 'a' ^^^ A | 's' ^^^ S

  val rule =
    val constr = '<' ^^^ (Rule.Lt(_, _, _)) | '>' ^^^ (Rule.Gt(_, _, _))
    rate ~ constr ~ num ~ ":" ~ name ^^ { case rate ~ constr ~ num ~ _ ~ target =>
      constr(rate, num, target)
    }
  val work =
    val rules = rep(rule <~ ",")
    name ~ "{" ~ rules ~ name ~ "}" ^^ { case name ~ _ ~ rules ~ andThen ~ _ =>
      (name -> Work(rules, andThen))
    }

  val state =
    val stat = rate ~ "=" ~ num ^^ { case rate ~ _ ~ num => (rate, num) }
    "{" ~ repsep(stat, ",") ~ "}" ^^ { case _ ~ stats ~ _ =>
      StateT(stats.foldLeft(Array(0, 0, 0, 0)) { case (arr, (rate, num)) =>
        arr(rate.ordinal) = num; arr
      })
    }

val (works, states) =
  lines.toSeq.splitBy(_ == "") match
    case Seq(works, states) =>
      (works.map(Parser.parse(Parser.work, _).get).toMap, states.map(Parser.parse(Parser.state, _).get))
    case _ => throw Exception("Expected two groupings")

// part 1

def part1 = println(
  states
    .filter(works.test("in"))
    .map(_.s.sum)
    .sum
)

// part 2

extension (m: Map[String, Work])
  def combinations(cur: String)(visited: Set[String], ranges: Ranges): Long =
    cur match
      case "A"                 => ranges.size
      case "R"                 => 0
      case _ if ranges.isEmpty => 0
      case _ =>
        val work = m(cur)
        var curRanges = ranges
        val passings = work.rules.map: rule =>
          val passing = curRanges.passing(rule)
          curRanges = curRanges.failing(rule)
          (passing, rule.andThen)
        ((curRanges, work.andThen) +: passings)
          .filter((_, t) => !visited.contains(t))
          .map((r, t) => combinations(t)(visited + t, r))
          .sum

def part2 =
  println(works.combinations("in")(Set.empty, StateT(Array.fill(4)(Range.empty))))

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