package aoc
package day12

import scala.annotation.tailrec
import scala.collection.mutable

import util.chaining.*

val dayNumber = "12"

@main def part1: Unit =
  println(part1(loadInput(dayNumber)))

@main def part2: Unit =
  println(part2(loadInput(dayNumber)))

def part1(input: String): String =
  def generateArrangements(cond: String): List[String] =
    @tailrec def aux(acc: List[String], remaining: String): List[String] =
      remaining.headOption match
        case None => acc
        case Some('?') =>
          aux(
            (acc.map(_ + '#') ++ acc.map(_ + '.')),
            remaining.tail
          )
        case Some('#') =>
          aux(acc.map(_ + '#'), remaining.tail)
        case Some('.') =>
          aux(acc.map(_ + '.'), remaining.tail)
        case _ => ??? // Can't happen
    aux(List(""), cond)

  def isValid(arr: String, groups: Array[Int]): Boolean =
    arr
      .split('.')
      .filter(_.nonEmpty)
      .map(_.length)
      .corresponds(groups)(_ == _)

  input
    .split('\n')
    .map { case s"$cond $groups" =>
      generateArrangements(cond) -> groups.split(',').map(_.toInt)
    }
    .map((arrs, groups) => arrs.count(isValid(_, groups)))
    .sum
    .toString

def part2(input: String): String =
  def countArrangements(
      cond: String,
      groups: List[Int]
  ): Long =
    val cache = mutable.Map.empty[(List[Char], List[Int], Int), Long]

    def countCached(cond: List[Char], groups: List[Int], curGroup: Int): Long =
      cache.getOrElseUpdate(
        (cond, groups, curGroup),
        count(cond, groups, curGroup)
      )

    def count(cond: List[Char], groups: List[Int], curGroup: Int): Long =
      if cond.isEmpty then
        if curGroup == 0 && groups.isEmpty || groups == List(curGroup) then 1
        else 0
      else
        cond.head match
          case '?' =>
            countCached('#' +: cond.tail, groups, curGroup) +
              countCached('.' +: cond.tail, groups, curGroup)
          case '#' =>
            if groups.headOption.contains(curGroup) then 0
            else countCached(cond.tail, groups, curGroup + 1)
          case '.' =>
            if groups.headOption.contains(curGroup) then
              countCached(cond.tail, groups.tail, 0)
            else if curGroup == 0 then countCached(cond.tail, groups, 0)
            else 0
          case _ => ??? // Can't happen

    countCached(cond.toList, groups, 0)

  input
    .split('\n')
    .map { case s"$cond $groups" =>
      cond -> groups.split(',').map(_.toInt).toList
    }
    .map((cond, groups) =>
      countArrangements(
        List.fill(5)(cond).mkString("?"),
        List.fill(5)(groups).flatten
      )
    )
    .sum
    .toString