from tools.aoc import AOCDay
from typing import Any


def get_available_steps(steps: dict) -> list:
    return [x for x in steps if len(steps[x]) == 0]


class Day(AOCDay):
    inputs = [
        [
            ("CABDFE", "input7_test"),
            ("JMQZELVYXTIGPHFNSOADKWBRUC", "input7"),
        ],
        [
            (15, "input7_test"),
            (1133, "input7"),
        ]
    ]

    def get_step_blocker(self) -> dict:
        steps = {}
        for line in self.getInput():
            tmp = line.split(" ")
            step_done, step_begin = tmp[1], tmp[7]
            if step_begin not in steps:
                steps[step_begin] = []
            if step_done not in steps:
                steps[step_done] = []

            steps[step_begin].append(step_done)

        return steps

    def part1(self) -> Any:
        steps = self.get_step_blocker()
        step_order = ""
        while next_steps := get_available_steps(steps):
            next_step = min(next_steps)
            step_order += next_step
            del steps[next_step]
            for x in steps:
                if next_step in steps[x]:
                    steps[x].remove(next_step)

        return step_order

    def part2(self) -> Any:
        steps = self.get_step_blocker()
        second = 0
        time_add = 0 if self._current_test_file.endswith("test") else 60
        worker_count = 2 if self._current_test_file.endswith("test") else 5
        step_order = ""
        current_work = []
        while second == 0 or current_work:
            new_work = []
            for work in current_work:
                if work[1] - 1 == 0:
                    step_order += work[0]
                    for x in steps:
                        if work[0] in steps[x]:
                            steps[x].remove(work[0])
                else:
                    new_work.append([work[0], work[1] - 1])

            current_work = new_work
            second += 1

            if len(current_work) < worker_count:
                available_worker = worker_count - len(current_work)
                avaiable_work = get_available_steps(steps)
                while avaiable_work and available_worker:
                    next_step = min(avaiable_work)
                    del steps[next_step]
                    avaiable_work.remove(next_step)
                    current_work.append([next_step, ord(next_step) - 64 + time_add])
                    available_worker -= 1

        return second - 1


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