#!/usr/bin/env lua
--[[
    Representation in both ways:
    The input mentions what bags a bag can contain,
    we store both what bags a bag can contain and
    what bags can contain a bag and in which amount.
    Part 2 requires the former and part 1 requires the latter.
--]]
local containers = {}
local contains = {}

for line in io.lines() do
    local bag_end, contain_start = line:find(" bags contain ")
    local bag_color = line:sub(1, bag_end - 1)
    contains[bag_color] = {}

    for count, color in line:sub(contain_start + 1):gmatch("(%d+) (%w+ %w+) bag") do
        if not containers[color] then
            containers[color] = {}
        end
        containers[color][bag_color] = count
        contains[bag_color][color] = count
    end
end

--[[
    For part 1, we build a set of bags that can contain the shiny gold bag.
    We will iterate over and over on this set until we no longer find any new
    bag colors that can contain others bags that can contain the shiny gold bag.
--]]
local part1set, found = {}, 0
for color, _ in pairs(containers["shiny gold"] or {}) do
    part1set[color] = true
    found = found + 1
end

local found_in_iteration = 0
repeat
    found_in_iteration = 0
    for known_color, _ in pairs(part1set) do
        for color, _ in pairs(containers[known_color] or {}) do
            if not part1set[color] then
                found_in_iteration = found_in_iteration + 1
                found = found + 1
                part1set[color] = true
            end
        end
    end
until found_in_iteration < 1

print(found)

--[[
    For part 2, we use a recursive function to find how many bags a single bag can hold.
    We use a table to provide memoization, gotta go fast.
--]]
local part2memo = {}
local function count_bags(color)
    if part2memo[color] then return part2memo[color] end
    local total = 0
    for sub_color, count in pairs(contains[color]) do
        -- Count the bag itself and all the bags it contains
        total = total + count * (count_bags(sub_color) + 1)
    end
    part2memo[color] = total
    return total
end

print(count_bags("shiny gold"))