use std::fmt;

use multimap::MultiMap;

fn main() {
    let input = include_str!("../data/input");
    let rules = parse_rules(input);
    print!("{}", FmtRules(rules.clone()));

    let mut reverted_rules: Rules = Default::default();
    rules.iter_all().for_each(|(container, contained)| {
        for c in contained {
            reverted_rules.insert(c.1, (1, *container));
        }
    });
    // dbg!(reverted_rules.clone());

    let shiny_gold_spec = ("shiny", "gold"); 
    let mut shiny_gold_containers = find_containers(
        &reverted_rules, 
        &shiny_gold_spec
    );
    shiny_gold_containers.0.sort();
    shiny_gold_containers.0.dedup();
    // shiny_gold_containers also contain shiny_gold sadly...
    dbg!(&shiny_gold_containers);
    println!("Number of bags that can contain shiny gold ones : {}", shiny_gold_containers.0.len()-1);


    //////// PART 2 /////////
    let shiny_gold_contained = find_containers(&rules, &shiny_gold_spec);
    // Remove 1 because the shiny gold bag actually doesn't contain itself ;)
    println!("Number of bags contained by shiny gold bags : {}", shiny_gold_contained.1 - 1);
}

type BagSpec<'a> = (&'a str, &'a str);

type Rules<'a> = MultiMap<BagSpec<'a>, (usize, BagSpec<'a>)>;
// Needed as you cannot implement a trait for an abstract type if not in current mod...
struct FmtRules<'a>(Rules<'a>);

impl<'a> fmt::Display for FmtRules<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for ((adjective,colour), rules) in self.0.iter_all() {
            write!(f, "{} {} bags contain ", adjective, colour)?;
            if rules.is_empty() {
                write!(f, "no other bags.")?;
            } else {
                for (idx, (nb, (adj, colour))) in rules.iter().enumerate() {
                    if idx > 0 { write!(f, ", ")? }
                    write!(f, "{} {} {} {}", nb, adj, colour, if *nb > 1 { "bags" } else { "bag" })?;
                }
                writeln!(f, ".")?;
            }
        }
        Ok(())
    }
}

fn parse_rules(input: &str) -> Rules<'_> {
    let mut rules: Rules = Default::default();

    peg::parser! {
        pub(crate) grammar parser() for str {
            pub(crate) rule root(r: &mut Rules<'input>)
                = (line(r) "." whitespace()*)* ![_]
            
            rule line(r: &mut Rules<'input>)
                = bag:bag_spec() " contain " rules:rules() {
                    if let Some(rules) = rules {
                        for rule in rules {
                            r.insert(bag, rule);
                        }
                    }
                }

            rule bag_spec() -> BagSpec<'input>
                = adjective:name() " " colour:name() " bag" "s"? {
                    (adjective, colour)
                }
            
            rule rules() -> Option<Vec<(usize, BagSpec<'input>)>>
                = rules:rules1()+ {Some(rules)}
                  / "no other bags" { None }
            
            rule rules1() -> (usize, BagSpec<'input>)
                = n:number() " " bag:bag_spec() ", "? {
                    (n, bag)
                }
            
            rule name() -> &'input str
                = $((!whitespace()[_])+)
            
            rule number() -> usize
                = n:$(['0'..='9']+) { n.parse().unwrap() }
            
            rule whitespace() = quiet!{[' ' | '\n' | '\t']+}
        }
    }

    parser::root(input, &mut rules).unwrap();
    rules
}

fn find_containers<'a>(rr: &'a Rules, start: &'a BagSpec) 
        -> (Vec<BagSpec<'a>>, usize) {
    let mut res = vec![*start];
    let mut count = 0;
    let containers = rr.get_vec(start);
    
    // for c in containers.iter().flat_map(|&v|v) {
    for c  in containers.into_iter().flatten().collect::<Vec<&(usize, (&str, &str))>>() {
        let mut b = find_containers(rr, &(c.1));
        res.append(&mut b.0);
        count += c.0 * b.1;
    }
    (res, count + 1)
}

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

    #[test]
    fn test_phase2() {
        let input = r"shiny gold bags contain 2 dark red bags.
dark red bags contain 2 dark orange bags.
dark orange bags contain 2 dark yellow bags.
dark yellow bags contain 2 dark green bags.
dark green bags contain 2 dark blue bags.
dark blue bags contain 2 dark violet bags.
dark violet bags contain no other bags.";
        let rules = parse_rules(input);
        let shiny_gold_spec = ("shiny", "gold"); 
        let shiny_gold_contained = find_containers(&rules, &shiny_gold_spec);
        assert_eq!(shiny_gold_contained.1 - 1, 126);
    }
}