package aoc.day12

import aoc._
import scala.collection.mutable

case class Record(input: String, groups: List[Int]):
  lazy val possibleFills = loop(0, groups)

  def unfold(times: Int) =
    Record((1 to times).map(_ => input).mkString("?"), (1 to times).flatMap(_ => groups).toList)

  private inline def canEmpty(pos: Char) = pos == '.' || pos == '?'
  private inline def canFill(pos: Char) = pos == '#' || pos == '?'

  private val memo = mutable.Map.empty[(Int, Int), Long]
  private def loop(from: Int, groups: List[Int]): Long =
    if from == input.length() then (if groups.isEmpty then 1 else 0)
    else
      memo.getOrElseUpdate(
        (from, groups.length), {
          // don't fill this group
          val noFill = if canEmpty(input(from)) then loop(from + 1, groups) else 0
          // fill this group
          val doFill = groups match
            case head :: next if from + head <= input.length =>
              val slice = input.slice(from, from + head)
              if slice.forall(canFill(_)) then
                if from + head == input.length() then loop(from + head, next)
                else if canEmpty(input(from + head)) then loop(from + head + 1, next)
                else 0
              else 0
            case _ => 0
          noFill + doFill
        }
      )

object Parser extends CommonParser:
  val groups = repsep(num, ",")
  val record = """[?#.]+""".r ~ groups ^^ { case (input ~ groups) => Record(input, groups) }

  def apply(s: String) = parse(record, s).get

val records = lines.map(Parser(_)).toSeq

def part1 =
  val result = records.map(_.possibleFills).sum
  println(result)

def part2 =
  val result = records.map(_.unfold(5).possibleFills).sum
  println(result)

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