package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

func main() {
	partOne()
	partTwo()
}

// Bag is a bag that contains other bags
type Bag struct {
	name          string
	contains      map[string]*Bag
	containCounts map[string]int
	isContainedBy map[string]*Bag
	seen          bool
}

func parseLine(s string) (string, map[string]int) {
	bags := strings.Split(s, " bags contain ")
	if bags[1] == "no other bags." {
		return bags[0], map[string]int{}
	}
	contained := strings.Split(bags[1], ", ")
	result := map[string]int{}
	for _, b := range contained {
		// strip off the bags?\.? and split on spaces
		innerBag := strings.Split(strings.TrimRight(b, "bags."), " ")
		// parse the count
		i, _ := strconv.Atoi(innerBag[0])
		result[innerBag[1]+" "+innerBag[2]] = i
	}
	return bags[0], result
}

func countUnseenParents(b *Bag) int {
	sum := 0
	for _, p := range b.isContainedBy {
		if !p.seen {
			p.seen = true
			sum = sum + countUnseenParents(p) + 1
		}
	}
	return sum
}

func countContainedBags(b *Bag) int {
	sum := 0
	for n, p := range b.contains {
		// make sure you count the bag itself when multiplying
		sum = sum + ((countContainedBags(p) + 1) * b.containCounts[n])
	}
	return sum
}

// You have a shiny gold bag.
// If you wanted to carry it in at least one other bag, how many different bag colors would be valid for the outermost bag?
// (In other words: how many colors can, eventually, contain at least one shiny gold bag?)

func partOne() {
	f, _ := os.Open("input")
	reader := bufio.NewReader(f)
	scanner := bufio.NewScanner(reader)

	bags := map[string]*Bag{}

	for scanner.Scan() {
		line := scanner.Text()
		bag, contains := parseLine(line)

		// create bag if it doesn't already exist
		if _, ok := bags[bag]; !ok {
			bags[bag] = &Bag{
				name:          bag,
				isContainedBy: map[string]*Bag{},
				contains:      map[string]*Bag{},
			}
		}
		// for the bags we now know it can contain
		for c := range contains {
			// if those bags already exist, update them
			if _, ok := bags[c]; ok {
				bags[c].isContainedBy[bag] = bags[bag]
			} else {
				// otherwise, create new bag with pointer to parent
				bags[c] = &Bag{
					name: c,
					isContainedBy: map[string]*Bag{
						bag: bags[bag],
					},
					contains: map[string]*Bag{},
				}
			}
			// finally, update our current bag to point to the one it contains
			bags[bag].contains[c] = bags[c]
		}
	}

	// actually do the walk up the tree
	fmt.Println(countUnseenParents(bags["shiny gold"]))
}

// Consider again your shiny gold bag and the rules from the above example:
// faded blue bags contain 0 other bags.
// dotted black bags contain 0 other bags.
// vibrant plum bags contain 11 other bags: 5 faded blue bags and 6 dotted black bags.
// dark olive bags contain 7 other bags: 3 faded blue bags and 4 dotted black bags.
// So, a single shiny gold bag must contain 1 dark olive bag (and the 7 bags within it) plus 2 vibrant plum bags (and the 11 bags within each of those): 1 + 1*7 + 2 + 2*11 = 32 bags!

// Of course, the actual rules have a small chance of going several levels deeper than this example;
// be sure to count all of the bags, even if the nesting becomes topologically impractical!

// Here's another example:
// 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.
// In this example, a single shiny gold bag must contain 126 other bags.

func partTwo() {
	f, _ := os.Open("input")
	reader := bufio.NewReader(f)
	scanner := bufio.NewScanner(reader)

	bags := map[string]*Bag{}

	for scanner.Scan() {
		line := scanner.Text()
		bag, contains := parseLine(line)

		// create bag if it doesn't already exist
		if _, ok := bags[bag]; !ok {
			bags[bag] = &Bag{
				name:          bag,
				isContainedBy: map[string]*Bag{},
				contains:      map[string]*Bag{},
				containCounts: map[string]int{},
			}
		}
		// for the bags we now know it can contain
		for c, i := range contains {
			// if those bags already exist, update them
			if _, ok := bags[c]; ok {
				bags[c].isContainedBy[bag] = bags[bag]
			} else {
				// otherwise, create new bag with pointer to parent
				bags[c] = &Bag{
					name: c,
					isContainedBy: map[string]*Bag{
						bag: bags[bag],
					},
					contains:      map[string]*Bag{},
					containCounts: map[string]int{},
				}
			}
			// finally, update our current bag to point to the one it contains
			// and how many it must contain of that bag
			bags[bag].contains[c] = bags[c]
			bags[bag].containCounts[c] = i
		}
	}

	// actually do the walk down the tree
	fmt.Println(countContainedBags(bags["shiny gold"]))
}