import java.io.File

data class Node(val id: Int) {
    val children = mutableListOf<Node>()
    val metadata = mutableListOf<Int>()

    val value: Int by lazy {
        if (children.isEmpty()) metadata.sum()
        else {
            metadata.sumBy { childRef ->
                children
                    .take(childRef)
                    .takeIf { it.size == childRef }
                    ?.lastOrNull()
                    ?.value ?: 0
            }
        }
    }
}

val allNodes = mutableSetOf<Node>()

fun main() {

    val inputString = File("input8.txt").readText().trim()

    val input = inputString
        .split(Regex("\\s+"))
        .map { it.toInt() }
        .iterator()

    val nodeIds = generateSequence(1) { it + 1 }.iterator()

    val root = readNode(nodeIds, input)

    //A
    println(allNodes.flatMap { it.metadata }.sum())

    //B
    println(root.value)
}

fun readNode(nodeIds: Iterator<Int>, input: Iterator<Int>): Node {
    val node = Node(nodeIds.next())
    val childrenCount = input.next()
    val metadataCount = input.next()

    allNodes += node

    repeat(childrenCount) {
        node.children += readNode(nodeIds, input)
    }

    repeat(metadataCount) {
        node.metadata += input.next()
    }

    return node
}