from collections import defaultdict

from aocd.models import Puzzle
from aocd import submit


def get_sizes(dirs: dict, max_size: int) -> int:
    """
    gather all filesizes from the input
    """
    filtered_dirs = [ value for value in dirs.values() if value < max_size]
    return sum(filtered_dirs)

def parse_input(input_str: str) -> list:
    """
    the input string consists of the following four message types
    - $ cd ..               -> pop dir name from stack
    - $ cd [dir]            -> push dir name to stack
    - $ ls                  -> ignore
    - dir [dir]             -> ignore
    - [size] [file]         -> sum up the sizes of all files in this dir
    parsing is done by implementing a state machine which gathers all
    available information in the following dictionary.
    - {list['/','dir','path']:int(size_sum)}
    """
    path = []
    dirs = defaultdict(int)
    for line in input_str.splitlines():
        match line.split():
            case ('$', 'cd', '..'):
                path.pop()
            case ('$', 'cd', name):
                path.append(name)
            case ('$', 'ls'):
                continue
            case ('dir', _):
                continue
            case (size, file):
                size = int(size)
                for i in range(len(path)):
                    key = tuple(path[:i+1]) # list not hashable
                    dirs[key] += size
    return dirs

if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2022, day=7)
    dirs = parse_input(puzzle.input_data)

    # part a: get the sum of all dirs where the size is at most 100k
    answer_a = get_sizes(dirs, 100000)
    print(f"{answer_a}")
    submit(answer_a, part="a", day=7, year=2022)

    # part b:
    answer_b = get_sizes(dirs, 100)
    print(f"{answer_b}")
    #submit(answer_b, part="b", day=7, year=2022)