use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::{HashMap, VecDeque};

#[derive(Default)]
pub struct Data {
    fwd: HashMap<String, Vec<(usize, String)>>,
    rev: HashMap<String, Vec<(usize, String)>>,
}

impl Data {
    pub fn new() -> Data {
        Default::default()
    }
}

#[aoc_generator(day7)]
#[inline(always)]
pub fn input_generator(input: &str) -> Data {
    let mut ret = Data::new();
    for line in input.lines() {
        let mut line = line.to_owned();
        line = line.replace("contain", "");
        line = line.replace("bags", "");
        line = line.replace("bag", "");
        line = line.replace(",", "");
        line = line.replace(".", "");
        while line.contains("  ") {
            line = line.replace("  ", " ");
        }
        line = line.trim().to_owned();
        let mut l: VecDeque<String> = line
            .split_ascii_whitespace()
            .map(|s| s.to_owned())
            .collect();
        let outer_color: String = [l.pop_front().unwrap(), l.pop_front().unwrap()].join(" ");
        while let (Some(n), Some(c1), Some(c2)) = (l.pop_front(), l.pop_front(), l.pop_front()) {
            let n: usize = n.parse().unwrap();
            let inner_color = vec![c1.to_owned(), c2.to_owned()].join(" ");
            ret.fwd
                .entry(outer_color.clone())
                .or_default()
                .push((n, inner_color.clone()));
            ret.rev
                .entry(inner_color.clone())
                .or_default()
                .push((n, outer_color.clone()));
        }
    }
    ret
}

#[aoc(day7, part1)]
#[inline(always)]
pub fn solve_part1(input: &Data) -> usize {
    let mut cnt = 0;
    let mut q = VecDeque::new();
    q.push_back("shiny gold".to_owned());
    while let Some(c) = q.pop_front() {
        for (_, nc) in input.rev.get(&c).iter().map(|v| v.iter()).flatten() {
            cnt += 1;
            q.push_back(nc.clone());
        }
    }
    return cnt;
}

#[aoc(day7, part2)]
#[inline(always)]
pub fn solve_part2(input: &Data) -> usize {
    let mut cnt = 0;
    let mut q = VecDeque::new();
    q.push_back((1, "shiny gold".to_owned()));
    while let Some((o_n, c)) = q.pop_front() {
        for (n, nc) in input.fwd.get(&c).iter().map(|v| v.iter()).flatten() {
            cnt += n * o_n;
            q.push_back((n * o_n, nc.clone()));
        }
    }
    return cnt;
}