#! /usr/bin/env python3
import itertools
import re
from collections import namedtuple

line_matcher = re.compile(r"(\w+ \w+) bags contain ([\w+ ,]+)+.")
child_matcher = re.compile(r"([0-9]+) (\w+ \w+) bag[s]{0,1}")


BagRule = namedtuple("BagRule", ["parent_bag", "bag_counts"])


def parse_line(line):
    result = line_matcher.match(line)
    if result is None:
        # print(f"{line}: No match")
        return None
    else:
        parent_bag = result.group(1)
        bag_counts = {}
        # print(result.groups())
        for bags in result.group(2).split(","):
            if bags == "no other bags":
                break
            child_result = child_matcher.search(bags.strip())
            # print(bags.strip(), child_result.groups())
            if child_result is None:
                print(f"Invalid child bags: {bags}")
            bag_counts[child_result.group(2)] = int(child_result.group(1))

        return BagRule(parent_bag, bag_counts)


def get_direct_carriers(contained_in, bags):
    return set(itertools.chain.from_iterable((
        contained_in.get(bag, [])
        for bag in bags
    )))


def part1():
    contained_in = {}
    with open("input.txt") as f:
        for line in f:
            rule = parse_line(line)
            for bag in rule.bag_counts.keys():
                contained_in.setdefault(bag, []).append(rule.parent_bag)

    my_bag = "shiny gold"
    holding_bags = set()
    last_level = set((my_bag,))
    while last_level:
        last_level = get_direct_carriers(contained_in, last_level)
        holding_bags |= last_level

    print(f"Holding bags: {len(holding_bags)}")


def count_children(contents, bag):
    rule = contents.get(bag)
    if rule is None:
        return 0
    return sum(
        count * (1 + count_children(contents, child_bag))
        for child_bag, count in rule.bag_counts.items()
    )


def part2():
    contents = {}
    with open("input.txt") as f:
        for line in f:
            rule = parse_line(line)
            contents[rule.parent_bag] = rule

    children = count_children(contents, "shiny gold")
    print(f"Holds {children} bags")


if __name__ == "__main__":
    part1()
    part2()