import scala.io.Source

sealed trait Token
case class Simple(s: String) extends Token
case class Expansion(t: Stream[Token], repCount: Int) extends Token

def read_marker(s: String): (Int, Int, String) = {
  val (size, rest) = s.span(_ != 'x')
  val (cnt, after) = rest.tail.span(_ != ')')
  (size.toInt, cnt.toInt, after.tail)
}

def expand_to_marker(s: String): (String, String) = {
  val (before, rest) = s.span(_ != '(')
  (before.filterNot(_.isWhitespace), if(rest.length > 0) rest.tail else "")
}

def expand(s: String, recurse: Boolean = false): Stream[Token] =
  if(s.length == 0) {
    Stream.empty
  } else {
    val (expandedStr, marker) = expand_to_marker(s)
    if(marker.length == 0) {
      Simple(expandedStr) #:: Stream.empty
    } else {
      val (sz, cnt, rest) = read_marker(marker)
      val patt = if(recurse) Expansion(expand(rest.take(sz), recurse), cnt) else Expansion(Simple(rest.take(sz)) #:: Stream.empty, cnt)
      Simple(expandedStr) #:: patt #:: expand(rest.drop(sz), recurse)
    }
  }

def repString(s: String, cnt: Int): String = {
  val sb = new StringBuilder("")
  var c = cnt
  while(c > 0) {
    sb ++= s
    c -= 1
  }
  sb.toString
}

def expanded(s: Stream[Token]): String = (s.map {
  case Simple(s) => s
  case Expansion(s, count) => repString(expanded(s), count)
}).mkString

def expanded_count(s: Stream[Token]): BigInt = (s.map {
  case Simple(s) => BigInt(s.length)
  case Expansion(s, count) => expanded_count(s) * BigInt(count)
}).sum

val test_cases1 = List(
  "ADVENT", "A(1x5)BC", "(3x3)XYZ", "A(2x2)BCD(2x2)EFG", "(6x1)(1x3)A", "X(8x2)(3x3)ABCY"
)

val test_cases2 = List(
  "(3x3)XYZ", "X(8x2)(3x3)ABCY", "(27x12)(20x12)(13x14)(7x10)(1x12)A", "(25x3)(3x3)ABC(2x3)XY(5x2)PQRSTX(18x9)(3x2)TWO(5x7)SEVEN"
)

test_cases1.foreach { tc =>
  val exp = expanded(expand(tc, false))
  println(s"$tc: (${exp.length})")
}

println("Part two")

test_cases2.foreach { tc =>
  val exp = expanded_count(expand(tc, true))
  println(s"$tc: (${exp})")
}

val input = Source.fromFile("input.9").mkString
println(
  expanded(expand(input)).length
)
println(
  expanded_count(expand(input, true))
)