use std::collections::{HashMap, VecDeque};

use aoc_runner_derive::{aoc, aoc_generator};
use num::Integer;

#[derive(Debug, Clone)]
enum Module {
    FlipFlop { state: bool, dests: Vec<String> },
    Conjunction { state: HashMap<String, bool>, dests: Vec<String> },
    Broadcast { dests: Vec<String> },
}

#[aoc_generator(day20)]
fn parse(input: &str) -> HashMap<String, Module> {
    let mut input: HashMap<_, _> = input.lines().map(|line| {
        let (name, dests) = line.split_once("->").unwrap();
        let dests = parse_destinations(dests);
        let module = parse_module(name.trim(), dests);
        (name.trim_matches(['&', '%', ' '].as_slice()).to_string(), module)
    }).collect();
    collect_inputs(&mut input);
    input
}

fn parse_module(name: &str, dests: Vec<String>) -> Module {
    match name {
        "broadcaster" => Module::Broadcast { dests },
        name if name.starts_with('%') => Module::FlipFlop { state: false, dests },
        name if name.starts_with('&') => Module::Conjunction { state: HashMap::new(), dests },
        _ => panic!("Invalid Module {}", name)
    }

}

fn collect_inputs(modules: &mut HashMap<String, Module>) {
    for (name, module) in modules.clone() {
        let dests = match module {
            Module::FlipFlop { state: _, dests: dest } => dest,
            Module::Conjunction { state: _, dests: dest } => dest,
            Module::Broadcast { dests: dest } => dest,
        };

        for dest_name in dests {
            if let Some(dest_mod) = modules.get_mut(&dest_name) {
                match dest_mod {
                    Module::Conjunction { state, dests: _ } => state.insert(name.to_string(), false),
                    _ => continue
                };
            }
        }
    }
}

fn parse_destinations(dests: &str) -> Vec<String> {
    dests.split(',')
        .map(|x| x.trim().to_string())
        .collect::<Vec<String>>()
}

#[aoc(day20, part1)]
fn part1(modules: &HashMap<String, Module>) -> usize {
    let mut modules = modules.clone();
    let mut high: usize = 0;
    let mut low: usize = 0;
    for _ in 0..1000 {
        let (high_pulses, low_pulses) = push_the_button(&mut modules);
        high += high_pulses;
        low += low_pulses;
    }

    high * low
}

fn push_the_button(modules: &mut HashMap<String, Module>) -> (usize, usize) {
    let mut process_queue: VecDeque<(String, bool, String)> = VecDeque::new();
    process_queue.push_back(("broadcaster".to_string(), false, "".to_string()));
    let mut low_pulses: usize = 0;
    let mut high_pulses: usize = 0;

    while let Some((module_name, input, sender)) = process_queue.pop_front() {
        match input {
            true => high_pulses += 1,
            false => low_pulses += 1,
        }
        if let Some(module) = modules.get_mut(&module_name) {
            match module {
                Module::FlipFlop { state, dests } => {
                    if !input {
                        let flipped = !(*state);
                        *state = flipped;
                        for dest in dests {
                            process_queue.push_back((dest.clone(), state.clone(), module_name.clone()));
                        }
                    }
                },
                Module::Conjunction { state, dests } => {
                    let stored = state.get_mut(&sender.to_string()).unwrap();
                    *stored = input;

                    let mut send = true;
                    if state.values().all(|x| *x) {
                        send = false;
                    }
                    for dest in dests {
                        process_queue.push_back((dest.clone(), send, module_name.clone()));
                    }
                },
                Module::Broadcast { dests } => {
                    for dest in dests {
                        process_queue.push_back((dest.clone(), input, module_name.clone()));
                    }
                },
            }
        }
    }

    return (high_pulses, low_pulses)
}

#[aoc(day20, part2)]
fn part2(og_modules: &HashMap<String, Module>) -> usize {
    let mut lcm = vec![];
    let before = ["tr", "xm", "dr", "nh"];

    for module_name in before {
        let mut modules = og_modules.clone();
        let mut count: usize = 1;
        while !push_the_button_part2(&mut modules, module_name.to_string()) {
            count += 1;
        }
        lcm.push(count);
    }

    println!("{:?}", lcm);

    let lcm = lcm.iter()
        .cloned()
        .reduce(|a, b| a.lcm(&b))
        .unwrap();

    lcm
}

fn push_the_button_part2(modules: &mut HashMap<String, Module>, to_find: String) -> bool {
    let mut process_queue: VecDeque<(String, bool, String)> = VecDeque::new();
    process_queue.push_back(("broadcaster".to_string(), false, "".to_string()));

    while let Some((module_name, input, sender)) = process_queue.pop_front() {
        if module_name == to_find && input == false {
            return true
        }
        if let Some(module) = modules.get_mut(&module_name) {
            match module {
                Module::FlipFlop { state, dests } => {
                    if !input {
                        let flipped = !(*state);
                        *state = flipped;
                        for dest in dests {
                            process_queue.push_back((dest.clone(), state.clone(), module_name.clone()));
                        }
                    }
                },
                Module::Conjunction { state, dests } => {
                    let stored = state.get_mut(&sender.to_string()).unwrap();
                    *stored = input;

                    let mut send = true;
                    if state.values().all(|x| *x) {
                        send = false;
                    }
                    for dest in dests {
                        process_queue.push_back((dest.clone(), send, module_name.clone()));
                    }
                },
                Module::Broadcast { dests } => {
                    for dest in dests {
                        process_queue.push_back((dest.clone(), input, module_name.clone()));
                    }
                },
            }
        }
    }

    return false
}


#[cfg(test)]
mod tests {
    use super::*;

    const EX_1: &str = r"broadcaster -> a, b, c
%a -> b
%b -> c
%c -> inv
&inv -> a";

const EX_2: &str = r"broadcaster -> a
%a -> inv, con
&inv -> b
%b -> con
&con -> rx";

    #[test]
    fn part1_example1() {
        assert_eq!(part1(&parse(EX_1)), 32000000);
    }



    #[test]
    fn part1_example2() {
        assert_eq!(part1(&parse(EX_2)), 11687500);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EX_2)), 1);
    }
}