import re
from functools import lru_cache
from typing import Iterator, Any, Generator, Dict

from aoc import BaseAssignment

rule_matcher = re.compile(
    r'(?P<bag>\w+\ \w+) bags contain (?P<contains>no other bags|(((, )?\d \w+\ \w+ bags?)+))\.')
contains_matcher = re.compile(r'(?P<count>\d+) (?P<bag>\w+\ \w+) bags?')


class Assignment(BaseAssignment):
    def get_bag_contents(self, contents: str) -> dict:
        content_dict = {}
        if contents != 'no other bags':
            for content in contents.split(', '):
                count = contains_matcher.match(content).groupdict()
                content_dict[count['bag']] = int(count['count'])

        return content_dict

    def read_input(self, example=False) -> Dict:
        self.mapping = {}

        for line in super().read_input(example):
            match = rule_matcher.match(line)
            rule = match.groupdict()
            self.mapping[rule['bag']] = self.get_bag_contents(rule['contains'])


class AssignmentOne(Assignment):
    @lru_cache
    def contains_shiny_gold_bag(self, bag):
        if 'shiny gold' in self.mapping[bag]:
            return True

        return any([
            self.contains_shiny_gold_bag(containing_bag)
            for containing_bag in self.mapping[bag].keys()
        ])

    def run(self, input) -> Any:
        nr_of_possibilities = 0
        for bag, contains in self.mapping.items():
            if self.contains_shiny_gold_bag(bag):
                nr_of_possibilities += 1
        return nr_of_possibilities


class AssignmentTwo(Assignment):
    def nr_of_bags_inside(self, bag):
        return sum([
            count + (count * self.nr_of_bags_inside(containing_bag))
            for containing_bag, count in self.mapping[bag].items()
        ])

    def run(self, input) -> Any:
        return self.nr_of_bags_inside('shiny gold')