from collections import deque
from tools.aoc import AOCDay
from typing import Any


class File:
    def __init__(self, name: str, size: int):
        self.name = name
        self.size = size


class Dir:
    def __init__(self, name: str, parent: 'Dir' = None):
        self.name = name
        self.parent = parent
        self.dirs = []
        self.files = []

    def getSize(self):
        files = sum(x.size for x in self.files)
        return files + sum(x.getSize() for x in self.dirs)


class Day(AOCDay):
    inputs = [
        [
            (95437, "input7_test"),
            (1367870, "input7"),
        ],
        [
            (24933642, "input7_test"),
            (549173, "input7"),
        ]
    ]

    def get_dir_struct(self) -> Dir:
        current_dir = None
        for line in self.getInput():
            if line.startswith("$"):  # command
                cmd = line[2:].split(" ")
                if cmd[0] == 'cd':
                    if cmd[1] == "..":
                        current_dir = current_dir.parent
                    else:
                        current_dir = Dir(cmd[1], current_dir)
                        if current_dir.parent is not None:
                            current_dir.parent.dirs.append(current_dir)

                continue

            file = line.split(" ")
            if file[0] != 'dir':
                file_size = int(file[0])
                current_dir.files.append(File(file[1], file_size))

        while current_dir.parent:
            current_dir = current_dir.parent

        return current_dir

    def get_sizes(self) -> (int, int):
        root = self.get_dir_struct()
        need_free = 30_000_000 - (70_000_000 - root.getSize())
        smallest = 1e20
        queue = deque()
        queue.append(root)
        sum_size = 0
        while queue:
            current_dir = queue.popleft()
            dir_size = current_dir.getSize()
            if dir_size <= 100_000:
                sum_size += dir_size
            if need_free <= dir_size < smallest:
                smallest = dir_size
            for sub_dir in current_dir.dirs:
                queue.append(sub_dir)

        return sum_size, smallest

    def part1(self) -> Any:
        ret, _ = self.get_sizes()
        return ret

    def part2(self) -> Any:
        _, ret = self.get_sizes()
        return ret


if __name__ == '__main__':
    day = Day(2022, 7)
    day.run(verbose=True)