use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, line_ending},
    combinator::{map, map_res, value},
    multi::separated_list1,
    sequence::{delimited, pair, separated_pair},
    IResult,
};
use std::{collections::VecDeque, fs, num::ParseIntError};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_20.txt")?;
    let circuit = Circuit::parser(&input).unwrap().1;

    {
        let mut circuit = circuit.clone();
        let mut pulse_tracker = PulseCounter::default();

        for _ in 0..1000 {
            circuit.push_the_button(&mut pulse_tracker);
        }
        dbg!(pulse_tracker.low_pulse_count * pulse_tracker.high_pulse_count);
    }

    {
        let mut circuit = circuit.clone();
        let mut pulse_tracker = RxWatcher::new(
            circuit
                .modules
                .iter()
                .find(|m| m.id == ModuleId::rx())
                .unwrap()
                .index,
            circuit
                .modules
                .iter()
                .find(|m| m.id == ModuleId::from_short_alphanumeric("th").unwrap())
                .unwrap()
                .index,
        );

        while !pulse_tracker.rx_got_a_low_pulse {
            pulse_tracker.i += 1;
            circuit.push_the_button(&mut pulse_tracker);
        }
        dbg!(pulse_tracker.i, pulse_tracker);
    }

    Ok(())
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct ModuleId(u32);

#[derive(Debug, Clone)]
struct Circuit {
    modules: Vec<Module>,
}

#[derive(Debug, Clone)]
struct ParsedModule {
    id: ModuleId,
    outs: Vec<ModuleId>,
    state: ModuleState,
}

#[derive(Debug, Clone)]
struct Module {
    id: ModuleId,
    index: usize,
    outs: Vec<usize>,
    state: ModuleState,
}

#[derive(Debug, Clone)]
enum ModuleState {
    FlipFlop(bool),
    Conjunction(u64),
    Sink,
}

#[derive(Debug, Default)]
struct PulseCounter {
    pulses: VecDeque<Pulse>,
    low_pulse_count: u64,
    high_pulse_count: u64,
}

#[derive(Debug)]
struct RxWatcher {
    pulses: VecDeque<Pulse>,
    rx_module_index: usize,
    rx_got_a_low_pulse: bool,
    logging_module_index: usize,
    i: usize,
}

impl RxWatcher {
    fn new(rx_module_index: usize, logging_module_index: usize) -> Self {
        RxWatcher {
            pulses: VecDeque::default(),
            rx_module_index,
            rx_got_a_low_pulse: false,
            logging_module_index,
            i: 0,
        }
    }
}

#[derive(Debug)]
struct Pulse {
    state: bool,
    input: usize,
    output: usize,
}

impl Circuit {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            pair(
                delimited(
                    tag("broadcaster -> "),
                    ParsedModule::outs_parser,
                    line_ending,
                ),
                separated_list1(line_ending, ParsedModule::parser),
            ),
            |(broadcaster, parsed_modules)| {
                let sink_index = parsed_modules.len() + 1;
                let rx_sink_index = parsed_modules.len() + 2;
                let mut modules: Vec<Module> = parsed_modules
                    .iter()
                    .enumerate()
                    .map(|(index, module)| Module {
                        id: module.id,
                        index: index + 1,
                        outs: module
                            .outs
                            .iter()
                            .map(|out_id| {
                                parsed_modules
                                    .iter()
                                    .enumerate()
                                    .find(|(_, out)| &out.id == out_id)
                                    .map_or_else(
                                        || {
                                            if out_id == &ModuleId::rx() {
                                                rx_sink_index
                                            } else {
                                                sink_index
                                            }
                                        },
                                        |(index, _)| index + 1,
                                    )
                            })
                            .collect(),
                        state: module.state.clone(),
                    })
                    .collect();
                modules.push(Module {
                    id: ModuleId::sink(),
                    index: sink_index,
                    outs: Vec::new(),
                    state: ModuleState::Sink,
                });
                modules.push(Module {
                    id: ModuleId::rx(),
                    index: rx_sink_index,
                    outs: Vec::new(),
                    state: ModuleState::Sink,
                });
                modules.insert(
                    0,
                    Module {
                        id: ModuleId::broadcaster(),
                        index: 0,
                        outs: broadcaster
                            .iter()
                            .map(|mid| modules.iter().find(|m| &m.id == mid).unwrap().index)
                            .collect(),
                        state: ModuleState::Sink,
                    },
                );

                let modules_snapshot = modules.clone();

                for module in &mut modules {
                    if let ModuleState::Conjunction(ref mut conjunction_state) = module.state {
                        for input in &modules_snapshot {
                            if input.outs.contains(&module.index) {
                                let mask = !(1 << input.index);
                                *conjunction_state &= mask;
                            }
                        }
                    }
                }

                Circuit { modules }
            },
        )(input)
    }

    fn push_the_button(&mut self, pulse_tracker: &mut impl PulseTracker) {
        pulse_tracker.button_pulse();
        for b in &self.modules[0].outs {
            pulse_tracker.push(Pulse {
                state: false,
                input: 0,
                output: *b,
            });
        }

        while let Some(pulse) = pulse_tracker.pop() {
            let module = &mut self.modules[pulse.output];
            let new_pulse_state: Option<bool> = match module.state {
                ModuleState::FlipFlop(ref mut current) => {
                    if !pulse.state {
                        *current = !*current;
                        Some(*current)
                    } else {
                        None
                    }
                }
                ModuleState::Conjunction(ref mut current) => {
                    let mask = 1 << pulse.input;
                    if pulse.state {
                        *current |= mask;
                    } else {
                        *current &= !mask;
                    }
                    Some(*current != u64::MAX)
                }
                ModuleState::Sink => None,
            };
            if let Some(new_pulse_state) = new_pulse_state {
                for out in &module.outs {
                    pulse_tracker.push(Pulse {
                        state: new_pulse_state,
                        input: module.index,
                        output: *out,
                    });
                }
            }
        }
    }
}

trait PulseTracker {
    fn button_pulse(&mut self);
    fn push(&mut self, pulse: Pulse);
    fn pop(&mut self) -> Option<Pulse>;
}

impl PulseTracker for PulseCounter {
    fn button_pulse(&mut self) {
        self.low_pulse_count += 1;
    }

    fn push(&mut self, pulse: Pulse) {
        if pulse.state {
            self.high_pulse_count += 1;
        } else {
            self.low_pulse_count += 1;
        }
        self.pulses.push_back(pulse);
    }

    fn pop(&mut self) -> Option<Pulse> {
        self.pulses.pop_front()
    }
}

impl PulseTracker for RxWatcher {
    fn button_pulse(&mut self) {}

    fn push(&mut self, pulse: Pulse) {
        if pulse.state && pulse.output == self.logging_module_index {
            println!("{}: {} into {}", self.i, pulse.input, pulse.output);
        }
        if !pulse.state && pulse.output == self.rx_module_index {
            self.rx_got_a_low_pulse = true;
        }
        self.pulses.push_back(pulse);
    }

    fn pop(&mut self) -> Option<Pulse> {
        self.pulses.pop_front()
    }
}

impl ParsedModule {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_pair(
                pair(
                    alt((
                        value(ModuleState::FlipFlop(false), tag("%")),
                        value(ModuleState::Conjunction(u64::MAX), tag("&")),
                    )),
                    map_res(alpha1, |id: &str| ModuleId::from_short_alphanumeric(id)),
                ),
                tag(" -> "),
                ParsedModule::outs_parser,
            ),
            |((state, id), outs)| ParsedModule { id, state, outs },
        )(input)
    }

    fn outs_parser(input: &str) -> IResult<&str, Vec<ModuleId>> {
        separated_list1(
            tag(", "),
            map_res(alpha1, |s: &str| ModuleId::from_short_alphanumeric(s)),
        )(input)
    }
}

impl ModuleId {
    fn broadcaster() -> ModuleId {
        ModuleId(0)
    }

    fn sink() -> ModuleId {
        ModuleId(u32::MAX)
    }

    fn rx() -> ModuleId {
        ModuleId::from_short_alphanumeric("rx").unwrap()
    }

    fn from_short_alphanumeric(s: &str) -> Result<ModuleId, ParseIntError> {
        u32::from_str_radix(s, 36).map(ModuleId)
    }
}