use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, line_ending},
    combinator::{map, map_res, value},
    multi::separated_list1,
    sequence::{delimited, pair, preceded, separated_pair, tuple},
    IResult,
};
use std::{
    collections::{BTreeMap, 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::default();

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

    Ok(())
}

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

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

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

#[derive(Debug, Clone)]
enum ModuleState {
    FlipFlop(bool),
    Conjunction(BTreeMap<ModuleId, bool>),
}

#[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_id: ModuleId,
    rx_got_a_low_pulse: bool,
}

impl Default for RxWatcher {
    fn default() -> Self {
        RxWatcher {
            pulses: VecDeque::default(),
            rx_module_id: ModuleId::from_short_alphanumeric("rx").unwrap(),
            rx_got_a_low_pulse: false,
        }
    }
}

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

impl Circuit {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                separated_list1(line_ending, Module::parser),
                delimited(
                    line_ending,
                    preceded(tag("broadcaster -> "), Module::outs_parser),
                    line_ending,
                ),
                separated_list1(line_ending, Module::parser),
            )),
            |(modules1, broadcaster, modules2)| {
                let mut modules: Vec<Module> =
                    modules1.into_iter().chain(modules2.into_iter()).collect();
                let modules_snapshot = modules.clone();

                for module in &mut modules {
                    if let ModuleState::Conjunction(ref mut conjunction_state) = module.state {
                        *conjunction_state = modules_snapshot
                            .iter()
                            .filter(|input| input.outs.contains(&module.id))
                            .map(|input| (input.id, false))
                            .collect();
                    }
                }

                Circuit {
                    broadcaster,
                    modules: modules.into_iter().map(|m| (m.id, m)).collect(),
                }
            },
        )(input)
    }

    fn push_the_button(&mut self, pulse_tracker: &mut impl PulseTracker) {
        pulse_tracker.button_pulse();
        for b in &self.broadcaster {
            pulse_tracker.push(Pulse {
                state: false,
                input: ModuleId::broadcaster(),
                output: *b,
            });
        }

        while let Some(pulse) = pulse_tracker.pop() {
            if let Some(module) = self.modules.get_mut(&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) => {
                        current.insert(pulse.input, pulse.state);
                        Some(current.values().any(|s| *s == false))
                    }
                };
                if let Some(new_pulse_state) = new_pulse_state {
                    for out in &module.outs {
                        pulse_tracker.push(Pulse {
                            state: new_pulse_state,
                            input: module.id,
                            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.rx_module_id {
            self.rx_got_a_low_pulse = true;
        }
        self.pulses.push_back(pulse);
    }

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

impl Module {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_pair(
                pair(
                    alt((
                        value(ModuleState::FlipFlop(false), tag("%")),
                        value(ModuleState::Conjunction(BTreeMap::new()), tag("&")),
                    )),
                    map_res(alpha1, |id: &str| ModuleId::from_short_alphanumeric(id)),
                ),
                tag(" -> "),
                Module::outs_parser,
            ),
            |((state, id), outs)| Module { 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 from_short_alphanumeric(s: &str) -> Result<ModuleId, ParseIntError> {
        u32::from_str_radix(s, 36).map(ModuleId)
    }
}