from aocd.models import Puzzle
from aocd import submit


def get_sizes_a(dirs: dict, max_size: int) -> int:
    """
    return the size of all dirs which are at most max_size big
    """
    filtered_dirs = [value for value in dirs.values() if value < max_size]
    return sum(filtered_dirs)


def get_sizes_b(dirs: dict, needed_size: int) -> int:
    """
    return the smallest size of a dir which when deleted amounts
    to the needed size of 30M.
    """
    total_size = 70000000
    used_size = dirs[("/",)]
    max_free = needed_size - (total_size - used_size)
    filtered_dirs = [value for value in dirs.values() if value > max_free]
    return min(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 = list()
    dirs = dict()
    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.setdefault(key, 0)  # initialize key for += to work
                    dirs[key] += size  # dirs needs to be initialized
    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_a(dirs, 100000)
    print(f"{answer_a}")
    submit(answer_a, part="a", day=7, year=2022)

    # part b: find the smallest D (huehue) that will free up at least 30M
    answer_b = get_sizes_b(dirs, 30000000)
    print(f"{answer_b}")
    submit(answer_b, part="b", day=7, year=2022)