import aoclib
import re

DAY = 7
TEST_SOLUTION_PART1 = 4
TEST_SOLUTION_PART2 = 32
MY_BAG = "shiny gold"
content_re = re.compile(r"(\d+) (\w+ \w+) bag.*")


def getBagContent(bag_line):
    """
    think of this like a shopping list
    "dark gold" needs:
      - "posh crimson": 3
      - "mirrored lavender": 3
    """
    color, contents = bag_line.split(" bags contain ")
    return_dict = {color: {}}
    if contents == 'no other bags.':
        return return_dict

    for content in contents.split(", "):
        match_groups = re.match(content_re, content)
        if not match_groups:
            print("ERROR: content '%s' cannot be matched" % content)

        return_dict[color][match_groups.group(2)] = int(match_groups.group(1))

    return return_dict


def getPossibleBagsForMyColor(color_dict, my_color):
    if my_color not in color_dict:
        return set()

    color_set = set(color_dict[my_color])
    for color in color_dict[my_color]:
        extend_set = getPossibleBagsForMyColor(color_dict, color)
        for extra_color in extend_set:
            color_set.add(extra_color)

    return color_set


def getBagContentCount(color_dict, my_color):
    if my_color not in color_dict or not color_dict[my_color]:
        return 0

    count = 0
    for color in color_dict[my_color]:
        sub_count = getBagContentCount(color_dict, color)
        count = count + color_dict[my_color][color] + color_dict[my_color][color] * sub_count

    return count


@aoclib.print_execution_time
def part1(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
    possible_bags = {}
    for bag_line in my_input:
        content_dict = getBagContent(bag_line)
        for bag_color in content_dict:
            for content_color in content_dict[bag_color]:
                if content_color not in possible_bags:
                    possible_bags[content_color] = [bag_color]
                else:
                    possible_bags[content_color].append(bag_color)

    return len(getPossibleBagsForMyColor(possible_bags, MY_BAG))


@aoclib.print_execution_time
def part2(test_mode=False):
    my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
    big_content_dict = {}
    for bag_line in my_input:
        content_dict = getBagContent(bag_line)
        for color in content_dict:
            big_content_dict[color] = content_dict[color]

    return getBagContentCount(big_content_dict, MY_BAG)