use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(7, 2020), part_1, part_2).with_expected(169, 82372)
}

const TARGET: &str = "shiny gold";

fn parse_bags(input: &str) -> Vec<(&str, Vec<(usize, &str)>)> {
  let mut bags = vec![];

  for line in input.lines() {
    let main_bag = &line[0..line.find(" bags").unwrap()];
    let other_bags = line[line.find("contain ").unwrap() + "contain ".len()..]
      .split(',')
      .map(str::trim)
      .map(|bag| {
        let count = bag[0..bag.find(' ').unwrap()]
          .parse::<usize>()
          .unwrap_or_default();
        let bag = &bag[bag.find(' ').unwrap() + 1..bag.find(" bag").unwrap()];
        (count, bag)
      })
      .collect::<Vec<(usize, &str)>>();

    bags.push((main_bag, other_bags));
  }

  bags
}

fn can_hold_target_bag<'a>(
  target: &'a str,
  bags: &[(&'a str, Vec<(usize, &'a str)>)],
) -> Vec<&'a str> {
  let mut can_hold = vec![];
  for (bag, children) in bags {
    if children.iter().any(|(_, child)| *child == target) {
      can_hold.push(*bag);
      can_hold.append(&mut can_hold_target_bag(bag, bags));
    }
  }

  can_hold
}

fn target_bag_holds_amount<'a>(
  target: &'a str,
  bags: &[(&'a str, Vec<(usize, &'a str)>)],
) -> usize {
  let mut count = 0;

  for (bag, children) in bags {
    if bag == &target {
      for (child_count, child_bag) in children {
        if child_count != &0 {
          count += child_count;
          count += child_count * target_bag_holds_amount(child_bag, bags);
        }
      }
    }
  }

  count
}

fn part_1(input: &str) -> Result<String> {
  Ok(
    can_hold_target_bag(TARGET, &parse_bags(input))
      .into_iter()
      .collect::<HashSet<&str>>()
      .len()
      .to_string(),
  )
}

fn part_2(input: &str) -> Result<String> {
  Ok(target_bag_holds_amount("shiny gold", &parse_bags(input)).to_string())
}