from __future__ import annotations
import math
from collections import defaultdict, deque
from tools.aoc import AOCDay
from typing import Any


class Module:
    def __init__(self, name: str, machine: Machine):
        self.name = name
        self.machine = machine
        self.state = False
        self.last_input = defaultdict(bool)
        self.inputs = []
        self.outputs = []

    def add_output(self, output: str):
        self.outputs.append(output)

    def add_input(self, input: str):
        self.inputs.append(input)

    def send(self, signal: bool):
        self.state = signal
        for output in self.outputs:
            self.machine.send(self.name, output, signal)

    def receive(self, sender: str, signal: bool):
        raise NotImplementedError()


class Broadcaster(Module):
    def receive(self, sender: str, signal: bool):
        self.send(signal)


class FlipFlop(Module):
    def receive(self, sender: str, signal: bool):
        if not signal:
            self.send(not self.state)


class Conjunction(Module):
    def receive(self, sender: str, signal: bool):
        self.last_input[sender] = signal
        self.send(sum(self.last_input.values()) != len(self.inputs))


class Machine:
    def __init__(self):
        self.queue = deque()
        self.send_signals = {True: 0, False: 0}
        self.modules = {}

    def add_module(self, module: Module):
        self.modules[module.name] = module

    def update_connections(self):
        modules = list(self.modules.values())
        for module in modules:
            module.inputs = []

        for module in modules:
            for output in module.outputs:
                if output not in self.modules:
                    self.modules[output] = FlipFlop(output, self)
                self.modules[output].add_input(module.name)

    def send(self, sender: str, destination: str, signal: bool):
        self.queue.append((sender, destination, signal))

    def push_button(self, monitor: str = None):
        self.send("button", "broadcaster", False)
        monitor_highs = []
        while self.queue:
            sender, destination, signal = self.queue.popleft()
            if destination == monitor and signal:
                monitor_highs.append(sender)
            self.send_signals[signal] += 1
            self.modules[destination].receive(sender, signal)

        return monitor_highs

    def get_signal_value(self) -> int:
        return self.send_signals[True] * self.send_signals[False]


class Day(AOCDay):
    inputs = [
        [
            (32000000, "input20_test"),
            (11687500, "input20_test2"),
            (856482136, "input20_dennis"),
            (1020211150, "input20"),
        ],
        [
            (224046542165867, "input20_dennis"),
            (238815727638557, "input20"),
        ],
    ]

    def parse_input(self) -> Machine:
        machine = Machine()
        for line in self.getInput():
            module_name, targets = line.split(" -> ")
            if module_name.startswith("broad"):
                module = Broadcaster(module_name, machine)
            elif module_name.startswith("%"):
                module = FlipFlop(module_name[1:], machine)
            elif module_name.startswith("&"):
                module = Conjunction(module_name[1:], machine)
            else:
                assert False, "unknown module type: %s" % module_name

            for target in targets.split(", "):
                module.add_output(target)

            machine.add_module(module)

        machine.update_connections()
        return machine

    def part1(self) -> Any:
        machine = self.parse_input()
        for _ in range(1000):
            machine.push_button()
        return machine.get_signal_value()

    def part2(self) -> Any:
        machine = self.parse_input()
        count = 0
        input_highs = {}
        to_monitor = machine.modules["rx"].inputs[0]
        wait_for = len(machine.modules[to_monitor].inputs)
        while True:
            count += 1
            found_highs = machine.push_button(monitor=to_monitor)

            for high in found_highs:
                input_highs[high] = count

            if len(input_highs) == wait_for:
                return math.lcm(*list(input_highs.values()))


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