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


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

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

        return inst, nodes

    def part1(self) -> Any:
        inst, nodes = self.parse_input()
        cur_node = "AAA"
        cur_step = 0
        while cur_node != "ZZZ":
            dir = inst[cur_step % len(inst)]
            if dir == "L":
                cur_node = nodes[cur_node][0]
            else:
                cur_node = nodes[cur_node][1]

            cur_step += 1

        return cur_step

    def part2(self) -> Any:
        inst, nodes = self.parse_input()
        i_len = len(inst)
        check_nodes = [x for x in nodes if x.endswith("A")]
        z_steps = []
        for node in check_nodes:
            cur_step = 0
            cur_node = 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
            z_steps.append(cur_step)

        return math.lcm(*z_steps)


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