package be.vandewalleh.aoc.days

import be.vandewalleh.aoc.utils.input.Day
import be.vandewalleh.aoc.utils.input.Input
import be.vandewalleh.aoc.utils.input.Lines
import be.vandewalleh.aoc.utils.input.createDay
import org.eclipse.collections.api.factory.Stacks
import org.eclipse.collections.api.multimap.list.ImmutableListMultimap
import org.eclipse.collections.api.stack.MutableStack
import org.eclipse.collections.impl.factory.Multimaps

data class Bag(val count: Int, val color: String)

@Day(7)
class Day07(@Lines val input: Input<List<String>>) {

    private val map: ImmutableListMultimap<String, Bag>

    init {
        val mutableMap = Multimaps.mutable.list.empty<String, Bag>()

        val colorRegex = "^(\\w+ \\w+)".toRegex()
        val requirementRegex = "(\\d+) (\\w+ \\w+) bag".toRegex()

        for (line in input.value) {
            val outerColor = colorRegex.find(line)!!.groupValues[1]
            for (match in requirementRegex.findAll(line)) {
                val (_, count, color) = match.groupValues
                mutableMap.put(outerColor, Bag(count.toInt(), color))
            }
        }

        map = mutableMap.toImmutable()
    }

    private fun bagSequence(rootColor: String): Sequence<Bag> = sequence {
        val stack: MutableStack<Bag> = Stacks.mutable.ofAll(map.get(rootColor))
        while (stack.notEmpty()) {
            val current = stack.pop().also { yield(it) }
            map[current.color]?.let {
                it.forEach { (count, color) ->
                    stack.push(Bag(current.count * count, color))
                }
            }
        }
    }

    fun part1() = map.keySet().count { bagSequence(it).any { it.color == "shiny gold" } }

    fun part2() = bagSequence("shiny gold").sumBy { it.count }

}

fun main() = with(createDay<Day07>()) {
    println(part1())
    println(part2())
}