zerosleeps

import unittest
from utils import get_raw_input


class Directory:
    def __init__(self, parent=None):
        self.parent = parent
        self.directories = {}
        self.files = {}

    def subdirectory(self, name):
        return self.directories.setdefault(name, Directory(parent=self))

    def add_file(self, size, name):
        return self.files.setdefault(name, int(size))

    @property
    def size(self):
        return sum(self.files.values()) + sum(
            [dir.size for dir in self.directories.values()]
        )

    @property
    def subdirectories(self):
        return list(self.directories.values()) + [
            sd
            for d in [
                directory.subdirectories for directory in self.directories.values()
            ]
            for sd in d
        ]


def build_file_system(raw_input):
    root_directory = Directory()
    current_directory = root_directory

    for line in raw_input:
        if line == "$ cd /":
            current_directory = root_directory
        elif line == "$ cd ..":
            current_directory = current_directory.parent
        elif line.startswith("$ cd"):
            current_directory = current_directory.subdirectory(line.split()[-1])
        elif line.startswith("dir"):
            current_directory.subdirectory(line.split()[-1])
        elif line[0].isdigit():
            current_directory.add_file(*line.split())

    return root_directory


def part_one(file_system):
    return sum(
        [
            directory.size
            for directory in file_system.subdirectories
            if directory.size <= 100000
        ]
    )


def part_two(file_system):
    min_space_to_find = 30000000 - (70000000 - file_system.size)
    candidate_directory_sizes = [
        directory.size
        for directory in file_system.subdirectories
        if directory.size >= min_space_to_find
    ]
    return min(candidate_directory_sizes)


if __name__ == "__main__":
    raw_input = get_raw_input("day_07_input.txt")
    file_system = build_file_system(raw_input)
    print(f"Part one: {part_one(file_system)}")
    print(f"Part two: {part_two(file_system)}")


class TestExamples(unittest.TestCase):
    def setUp(self):
        self.input = [
            "$ cd /",
            "$ ls",
            "dir a",
            "14848514 b.txt",
            "8504156 c.dat",
            "dir d",
            "$ cd a",
            "$ ls",
            "dir e",
            "29116 f",
            "2557 g",
            "62596 h.lst",
            "$ cd e",
            "$ ls",
            "584 i",
            "$ cd ..",
            "$ cd ..",
            "$ cd d",
            "$ ls",
            "4060174 j",
            "8033020 d.log",
            "5626152 d.ext",
            "7214296 k",
        ]
        self.file_system = build_file_system(self.input)

    def test_size_of_e(self):
        self.assertEqual(self.file_system.directories["a"].directories["e"].size, 584)

    def test_size_of_a(self):
        self.assertEqual(self.file_system.directories["a"].size, 94853)

    def test_size_of_d(self):
        self.assertEqual(self.file_system.directories["d"].size, 24933642)

    def test_size_of_root_directory(self):
        self.assertEqual(self.file_system.size, 48381165)

    def test_part_one_example(self):
        self.assertEqual(part_one(self.file_system), 95437)

    def test_part_two_example(self):
        self.assertEqual(part_two(self.file_system), 24933642)