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

const INPUT: &str = include_str!("../../data/day7");

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
struct BagId(usize);

#[derive(Clone, Debug)]
struct Bag {
	name: String,
	contained_by: HashMap<BagId, usize>,
	contains: HashMap<BagId, usize>,
	contained_bags: Option<u64>,
}

impl Bag {
	fn new(bag: &str) -> Self {
		Self {
			name: bag.into(),
			contained_by: HashMap::new(),
			contains: HashMap::new(),
			contained_bags: None,
		}
	}
}

#[derive(Clone, Debug)]
struct Bags {
	name_to_id: HashMap<String, BagId>,
	bags: Vec<Bag>,
}

impl Bags {
	fn get_bag_id(&mut self, bag: &str) -> BagId {
		let bags = &mut self.bags; // for access in closure
		*self.name_to_id.entry(bag.to_string()).or_insert_with(|| {
			let id = BagId(bags.len());
			bags.push(Bag::new(bag));
			id
		})
	}

	fn add_bag_spec(&mut self, bag: &str, contains: &[(&str, usize)]) {
		let bag_id = self.get_bag_id(bag);
		let contains: HashMap<BagId, usize> = contains.into_iter().map(|(name, amount)| {
			(self.get_bag_id(name), *amount)
		}).collect();
		for (&id, &amount) in &contains {
			assert!(amount > 0);
			self.bags[id.0].contained_by.insert(bag_id, amount);
		}
		let bag = &mut self.bags[bag_id.0];
		assert!(bag.contains.is_empty());
		bag.contains = contains;
	}

	fn add_bag_spec_str(&mut self, line: &str) {
		use std::convert::TryInto;

		let line = line.trim().strip_suffix(".").expect("need trailing dot");
		let [bag_name, specs]: [&str; 2] = line.splitn(2, " bags contain ").collect::<Vec<_>>().try_into().expect("invalid line");
		let specs = if specs != "no other bags" {
			specs.split(", ").map(|spec| {
				if let Some(spec) = spec.strip_suffix(" bag") {
					let spec = spec.strip_prefix("1 ").expect("singular bag requires amount 1");
					(spec, 1)
				} else {
					let spec = spec.strip_suffix(" bags").expect("either singular or plural required");
					let [amount, spec]: [&str; 2] = spec.splitn(2, " ").collect::<Vec<_>>().try_into().expect("invalid spec");
					(spec, amount.parse::<usize>().unwrap())
				}
			}).collect::<Vec<_>>()
		} else {
			Vec::new()
		};
		self.add_bag_spec(bag_name, &specs);
	}

	fn parse(data: &str) -> Self {
		let mut this = Self {
			name_to_id: HashMap::new(),
			bags: Vec::new(),
		};
		for line in data.lines() {
			this.add_bag_spec_str(line)
		}
		this
	}

	fn find_bag(&self, name: &str) -> Option<BagId> {
		self.name_to_id.get(name).cloned()
	}

	fn contained_bags(&mut self, bag: BagId) -> u64 {
		if let Some(n) = self.bags[bag.0].contained_bags {
			return n;
		}
		let contains = self.bags[bag.0].contains.clone();
		let n = contains.into_iter().map(|(inner_bag, amount)| {
			(amount as u64) * (self.contained_bags(inner_bag) + 1)
		}).sum::<u64>();
		self.bags[bag.0].contained_bags = Some(n);
		n
	}
}

fn main() {
	let mut bags = Bags::parse(INPUT);
	let shiny_gold = bags.find_bag("shiny gold").unwrap();

	let mut walk = VecDeque::new();
	let mut contain_shiny_gold = HashSet::new();
	walk.push_back(shiny_gold);
	while let Some(bag_id) = walk.pop_front() {
		if contain_shiny_gold.insert(bag_id) {
			for contains in bags.bags[bag_id.0].contained_by.keys() {
				walk.push_back(*contains);
			}
		}
	}
	println!("Bags containing shiny gold: {}", contain_shiny_gold.len() - 1);
	println!("shiny gold contains: {}", bags.contained_bags(shiny_gold));
}