from tools.aoc import AOCDay


def get_fuel(count: int, part2: bool = False):
    if not part2:
        return count // 3 - 2
    else:
        fuel_sum = 0
        count = count // 3 - 2
        while count > 0:
            fuel_sum += count
            count = count // 3 - 2

        return fuel_sum


class Day(AOCDay):
    inputs = [
        [
            (2, "test_input01"),
            (2, "test_input01_2"),
            (654, "test_input01_3"),
            (33583, "test_input01_4"),
            (3282386, "input01")
        ],
        [
            (2, "test_input01_2"),
            (966, "test_input01_3"),
            (50346, "test_input01_4"),
            (4920708, "input01")
        ]
    ]

    def get_fuel(self, part2: bool = False):
        if len(self.input) == 1:
            return get_fuel(self.getInput(int), part2)
        else:
            fuel_sum = 0
            for x in self.getInput(int):
                fuel_sum += get_fuel(x, part2)

            return fuel_sum

    def part1(self):
        return self.get_fuel()

    def part2(self):
        return self.get_fuel(True)


if __name__ == '__main__':
    day = Day(2019, 1)
    day.run(verbose=True)