from tools.aoc import AOCDay
from typing import Any


class Node:
    def __init__(self, data: list) -> None:
        self.metadata: list = data[-data[1]:]
        self.children: list = []

    def get_metadata_sum(self):
        return sum(self.metadata) + sum(x.get_metadata_sum() for x in self.children)

    def get_value(self):
        if not self.children:
            return sum(self.metadata)
        else:
            my_value = 0
            for x in self.metadata:
                if x - 1 < len(self.children):
                    my_value += self.children[x - 1].get_value()

            return my_value

    def print(self, depth: int = 0):
        print(" " * depth, self.metadata)
        for x in self.children:
            x.print(depth + 2)


class Day(AOCDay):
    inputs = [
        [
            (138, "input8_test"),
            (40746, "input8"),
        ],
        [
            (66, "input8_test"),
            (37453, "input8"),
        ]
    ]

    def get_node_tree(self) -> Node:
        inp = self.getInputAsArraySplit(' ', int)
        children = []
        while inp:
            index = 0
            parent_index = 0
            while index < len(inp) and inp[index]:
                if isinstance(inp[index], str):
                    index += 1
                else:
                    parent_index = index
                    index += 2
            sub_children = []
            while index < len(inp) - 2 and str(inp[index + 2]).startswith("A"):
                try:
                    sub_children.append(children[int(inp.pop(index + 2)[1:])])
                except IndexError:
                    print(children, index)
            node = Node(inp[index:index+2+inp[index + 1]])
            node.children = sub_children
            children.append(node)
            inp[parent_index] -= 1
            inp = inp[:index] + ['A'+str(len(children) - 1)] + inp[index+2+inp[index + 1]:]
            if len(inp) == 1:
                del children
                return node

    def part1(self) -> Any:
        return self.get_node_tree().get_metadata_sum()

    def part2(self) -> Any:
        return self.get_node_tree().get_value()


if __name__ == '__main__':
    day = Day(2018, 8)
    day.run(verbose=True)