from pathlib import Path
from typing import NamedTuple
from itertools import cycle
import math


class Node(NamedTuple):
    id: str
    left: str
    right: str


def parse(input: Path) -> tuple[str, dict[str, Node]]:
    with input.open() as f:

        steps = f.readline().strip()
        nodes: dict[str, Node] = {}

        _ = f.readline()

        for line in f:
            node = Node(
                line[0:3],
                line[7:10],
                line[12:15],
            )
            nodes[node.id] = node

    return steps, nodes


def part2(input: Path) -> int:
    steps, nodes = parse(input)

    start_nodes: list[str] = [
        node for node in nodes.keys() if node.endswith("A")
    ]

    periods: dict[str, int] = {}
    for start_node in start_nodes:
        current = start_node
        num_steps = 0
        for step in cycle(steps):
            if current.endswith("Z"):
                print(f"Found a node for {start_node} at {current} in {num_steps}")
                periods[start_node] = num_steps
                # This assumes there is only one possible end and that it begins a
                # cycle after that. Apparently true for this data set.
                break

            num_steps += 1

            if step == "L":
                current = nodes[current].left
            elif step == "R":
                current = nodes[current].right
            else:
                raise ValueError(f"Unknown step {step}")

    return math.lcm(*[p for p in periods.values()])


def part2_brute(input: Path) -> int:
    steps, nodes = parse(input)

    current: list[str] = [
        node for node in nodes.keys() if node.endswith("A")
    ]

    num_steps = 0
    for step in steps:
        if all(node.endswith("Z") for node in current):
            break

        num_steps += 1

        for i, node in enumerate(current):
            # if node.endswith("Z"):
            #     print(f"Node {initial[i]} reached {node} in {num_steps-1} steps")
            if step == "L":
                current[i] = nodes[node].left
            elif step == "R":
                current[i] = nodes[node].right
            else:
                raise ValueError(f"Unknown step {step}")

    return num_steps


def part1(input: Path) -> int:
    steps_s, nodes = parse(input)
    steps = cycle(steps_s)

    current = "AAA"
    num_steps = 0
    for step in steps:
        if current == "ZZZ":
            break

        num_steps += 1

        if step == "L":
            current = nodes[current].left
        elif step == "R":
            current = nodes[current].right
        else:
            raise ValueError(f"Unknown step {step}")

    return num_steps


if __name__ == "__main__":
    # input = Path("sample.txt")
    input = Path("input.txt")

    result1 = part1(input)
    print("part1", result1)

    result2 = part2(input)
    print("part2", result2)