import std/[strutils, tables]
import lib/[loader, util]


type
    Bag = object
        adj: string
        color: string

    BagCount = object
        bag: Bag
        count: int

    Rule = object
        outer: Bag
        contents: seq[BagCount]


proc `==`(a, b: Bag): bool =
    a.adj == b.adj and a.color == b.color


func parseBag(spec: string): Bag =
    [adj, color] <- spec.split()
    result = Bag(adj: adj, color: color)


func parseRule(spec: string): Rule =
    [outer, inner] <- spec.split(" bags contain ")
    let outerBag = parseBag(outer)

    var bagCounts: seq[BagCount]
    if inner != "no other bags.":
        # make sure to strip "bags." off the end of the contents
        for bc_spec in inner[0..^6].split(", "):
            [count, bag_spec] <- bc_spec.split(maxsplit = 1)
            let bc = BagCount(
                bag: parseBag(bag_spec),
                count: parseInt(count)
            )
            bagCounts.add(bc)

    result = Rule(
        outer: parseBag(outer),
        contents: bagCounts
    )


proc loadRules(): Table[Bag, Rule] =
    for line in loadStrings(7):
        let rule = parseRule(line)
        result[rule.outer] = rule


proc partOne(rules: Table[Bag, Rule]): int =
    # inner function so we can capture the "rules" var I guess?
    proc find(rule: Rule, target: Bag): bool =
        for bagCount in rule.contents:
            if bagCount.bag == target:
                return true
            elif find(rules[bagCount.bag], target):
                return true
        return false

    let b = Bag(adj: "shiny", color: "gold")
    for _, r in rules:
        if find(r, b):
            result += 1


proc partTwo(rules: Table[Bag, Rule]): int =
    # ok so that worked, let's do it again
    proc sumContents(bag: Bag): int =
        for bagCount in rules[bag].contents:
            result += bagCount.count # for the inner bag itself
            result += bagCount.count * sumContents(bagCount.bag)
        # "x contains no other bag" is a no-op, so returns default 0

    let b = Bag(adj: "shiny", color: "gold")
    result = sumContents(b)


when isMainModule:
    let rules = loadRules()
    echo "One: ", partOne(rules)
    echo "Two: ", partTwo(rules)