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

from .aoc import Aoc2, AocSameParser


@dataclass
class Map:
    directions: list[int]
    rules: dict[str, tuple[str, str]]

    @staticmethod
    def from_input(inpt: str) -> Map:
        lines = inpt.splitlines()

        directions: list[int] = [0 if rl == "L" else 1 for rl in lines[0]]
        rules = {line[:3]: (line[7:10], line[12:15]) for line in lines[2:]}
        return Map(directions, rules)


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

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

        for steps, direction in enumerate(itertools.cycle(inpt.directions)):
            if current == "ZZZ":
                return steps
            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, nr in enumerate(list_of_nrs):
            for nr2 in list_of_nrs[index + 1 :]:
                diff = nr2 - nr
                if nr2 + diff in list_of_nrs:
                    return nr, diff
        return None

    def part2(self, inpt: Map) -> int:
        current = [node for node in inpt.rules.keys() if node.endswith("A")]
        steps_to_z: list[list[int]] = [[] 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
            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, example_code_nr1=1, example_code_nr2=-1)
    day08.run()


if __name__ == "__main__":
    main()