from __future__ import annotations
import itertools
import math
from typing import Optional

from .aoc import Aoc2, AocSameParser


class Map:
    def __init__(self, inpt: str) -> None:
        lines = inpt.splitlines()

        self.directions: list[int] = [0 if rl == "L" else 1 for rl in lines[0]]

        self.rules = {line[:3]: (line[7:10], line[12:15]) for line in lines[2:]}


class Day08(AocSameParser[Map], Aoc2[Map, Map]):
    def parseinput(self, inpt: str) -> Map:
        return Map(inpt)

    def part1(self, inpt: Map) -> int:
        current = "AAA"

        for steps, direction in enumerate(itertools.cycle(inpt.directions)):
            if current == "ZZZ":
                return steps
            else:
                current = inpt.rules[current][direction]
        raise RuntimeError("End of the rainbow")

    def find_loop(self, list_of_nrs: list[int]) -> Optional[tuple[int, int]]:
        for index in range(len(list_of_nrs)):
            for index in range(len(list_of_nrs[index + 1 :])):
                diff = list_of_nrs[index + 1] - list_of_nrs[index]
                if list_of_nrs[index + 1] + diff in list_of_nrs:
                    return list_of_nrs[index], diff
        return None

    def part2(self, inpt: Map) -> int:
        current = [node for node in inpt.rules.keys() if node.endswith("A")]
        steps_to_z = [[] for _ in current]
        loops: list[Optional[tuple[int, int]]] = [None for _ in current]

        for steps, direction in enumerate(itertools.cycle(inpt.directions)):
            for n, node in enumerate(current):
                if node.endswith("Z"):
                    steps_to_z[n].append(steps)
                loops[n] = self.find_loop(steps_to_z[n])

            if not any(loop is None for loop in loops):
                break
            else:
                current = [inpt.rules[node][direction] for node in current]

        # Turns out, each path directly loops :/
        return math.lcm(*(loop[0] for loop in loops if loop is not None))


def main() -> None:
    day08 = Day08(2023, 8)
    day08._example_code_nr1 = 1
    day08._example_code_nr2 = -1
    day08.run()


if __name__ == "__main__":
    main()