import math
from tools.aoc import AOCDay
from typing import Any


def get_steps_to_z(inst: str, nodes: dict, from_node: str) -> int:
    i_len = len(inst)
    cur_step = 0
    cur_node = from_node
    while not cur_node.endswith("Z"):
        if inst[cur_step % i_len] == "L":
            cur_node = nodes[cur_node][0]
        else:
            cur_node = nodes[cur_node][1]

        cur_step += 1

    return cur_step


class Day(AOCDay):
    inputs = [
        [
            (2, "input8_test1"),
            (6, "input8_test2"),
            (18023, "input8_dennis"),
            (12643, "input8"),
        ],
        [
            (6, "input8_test3"),
            (14449445933179, "input8_dennis"),
            (13133452426987, "input8"),
        ],
    ]

    def parse_input(self) -> (str, dict):
        inst = self.getInput()[0]
        nodes = {}
        for line in self.getInput()[2:]:
            node, paths = line.split(" = ")
            pl, pr = paths.split(", ")
            nodes[node] = (pl[1:], pr[:-1])

        return inst, nodes

    def part1(self) -> Any:
        inst, nodes = self.parse_input()
        return get_steps_to_z(inst, nodes, "AAA")

    def part2(self) -> Any:
        inst, nodes = self.parse_input()
        check_nodes = [x for x in nodes if x.endswith("A")]
        z_steps = [get_steps_to_z(inst, nodes, x) for x in check_nodes]
        return math.lcm(*z_steps)


if __name__ == "__main__":
    day = Day(2023, 8)
    day.run(verbose=True)