import logging
import os
import sys
from collections import defaultdict
from math import lcm
from typing import Literal, TypeAlias

VERBOSE = os.getenv("AOC_VERBOSE") == "True"
logging.basicConfig(level=logging.INFO if VERBOSE else logging.WARNING)


ModuleType: TypeAlias = Literal["broadcaster", "conjunction", "flip-flop"]
PulseType: TypeAlias = Literal["high", "low"]

modules: dict[str, tuple[ModuleType, list[str]]] = {}

lines = sys.stdin.read().splitlines()

for line in lines:
    name, outputs_s = line.split(" -> ")
    outputs = outputs_s.split(", ")
    if name == "broadcaster":
        modules["broadcaster"] = ("broadcaster", outputs)
    else:
        modules[name[1:]] = (
            "conjunction" if name.startswith("&") else "flip-flop",
            outputs,
        )


def process(
    start: tuple[str, str, PulseType],
    flip_flop_states: dict[str, Literal["on", "off"]],
    conjunction_states: dict[str, dict[str, PulseType]],
) -> tuple[dict[PulseType, int], dict[str, dict[PulseType, int]]]:
    pulses: list[tuple[str, str, PulseType]] = [start]
    counts: dict[PulseType, int] = {"low": 0, "high": 0}
    inputs: dict[str, dict[PulseType, int]] = defaultdict(lambda: {"low": 0, "high": 0})

    logging.info("starting process... ")

    while pulses:
        input, name, pulse = pulses.pop(0)
        logging.info(f"{input} -{pulse}-> {name}")
        counts[pulse] += 1

        inputs[name][pulse] += 1

        if name not in modules:
            continue

        type, outputs = modules[name]

        if type == "broadcaster":
            ...

        elif type == "flip-flop":
            if pulse == "high":
                continue

            if flip_flop_states[name] == "off":
                flip_flop_states[name] = "on"
                pulse = "high"
            else:
                flip_flop_states[name] = "off"
                pulse = "low"

        else:
            conjunction_states[name][input] = pulse

            if all(state == "high" for state in conjunction_states[name].values()):
                pulse = "low"
            else:
                pulse = "high"

        pulses.extend((name, output, pulse) for output in outputs)

    return counts, inputs


with open("./day20.dot", "w") as fp:
    fp.write("digraph G {\n")
    fp.write("rx [shape=circle, color=red, style=filled];\n")
    for name, (type, outputs) in modules.items():
        if type == "conjunction":
            shape = "diamond"
        elif type == "flip-flop":
            shape = "box"
        else:
            shape = "circle"
        fp.write(f"{name} [shape={shape}];\n")
    for name, (type, outputs) in modules.items():
        for output in outputs:
            fp.write(f"{name} -> {output};\n")
    fp.write("}\n")

# part 1
flip_flop_states: dict[str, Literal["on", "off"]] = {
    name: "off" for name, (type, _) in modules.items() if type == "flip-flop"
}
conjunction_states: dict[str, dict[str, PulseType]] = {
    name: {input: "low" for input, (_, outputs) in modules.items() if name in outputs}
    for name, (type, _) in modules.items()
    if type == "conjunction"
}
counts: dict[PulseType, int] = {"low": 0, "high": 0}
for _ in range(1000):
    result, _ = process(
        ("button", "broadcaster", "low"), flip_flop_states, conjunction_states
    )
    for pulse in ("low", "high"):
        counts[pulse] += result[pulse]
answer_1 = counts["low"] * counts["high"]
print(f"answer 1 is {answer_1}")

# part 2

# reset states
for name in flip_flop_states:
    flip_flop_states[name] = "off"

for name in conjunction_states:
    for input in conjunction_states[name]:
        conjunction_states[name][input] = "low"

# find the conjunction connected to rx
to_rx = [name for name, (_, outputs) in modules.items() if "rx" in outputs]
assert len(to_rx) == 1, "cannot handle multiple module inputs for rx"
assert (
    modules[to_rx[0]][0] == "conjunction"
), "can only handle conjunction as input to rx"

to_rx_inputs = [name for name, (_, outputs) in modules.items() if to_rx[0] in outputs]
assert all(
    modules[i][0] == "conjunction" and len(modules[i][1]) == 1 for i in to_rx_inputs
), "can only handle inversion as second-order inputs to rx"


count = 1
cycles: dict[str, int] = {}
second: dict[str, int] = {}
while len(second) != len(to_rx_inputs):
    _, inputs = process(
        ("button", "broadcaster", "low"), flip_flop_states, conjunction_states
    )

    for node in to_rx_inputs:
        if inputs[node]["low"] == 1:
            if node not in cycles:
                cycles[node] = count
            elif node not in second:
                second[node] = count

    count += 1

assert all(
    second[k] == cycles[k] * 2 for k in to_rx_inputs
), "cannot only handle cycles starting at the beginning"

answer_2 = lcm(*cycles.values())
print(f"answer 2 is {answer_2}")