Handy Haversacks

defmodule Aoc2020.Day7 do
  @moduledoc "Handy Haversacks"

  def run(), do: parse_input() |> part_2()

  def part_1(rules) do
    find_parents(rules, :"shiny gold") |> length
  end

  def part_2(rules) do
    count_children(rules, :"shiny gold")
    |> List.flatten()
    |> Enum.sum()
  end

  def parse_input() do
    File.read!("priv/inputs/2020/day7.txt")
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(&parse_rule/1)
  end

  def parse_rule(line) do
    [colour, children] = String.split(line, " contain ")
    {String.to_atom(String.replace(colour, " bags", "")), parse_children(children)}
  end

  def parse_children(str) do
    str
    |> String.split(", ")
    |> Enum.map(&String.replace(&1, ~r( bags?.?), ""))
    |> Enum.map(&parse_count/1)
    |> Enum.filter(& &1)
  end

  def parse_count("no other") do
    nil
  end

  def parse_count(str) do
    [count, colour] = String.split(str, " ", parts: 2)
    {String.to_atom(colour), String.to_integer(count)}
  end

  # part 1
  def find_parents(rules, colour) do
    bag_containers(rules, colour)
    |> List.flatten()
    |> Enum.uniq()
  end

  def bag_containers(rules, colour) do
    has_child(rules, colour)
    |> Enum.map(fn {c, _} -> [c, bag_containers(rules, c)] end)
  end

  def has_child(rules, child_colour) do
    Enum.filter(rules, fn {col, children} -> Keyword.has_key?(children, child_colour) end)
  end

  # part 2
  def count_children(rules, colour, multiplier \\ 1) do
    {_, children} = Enum.find(rules, fn {col, children} -> col == colour end)

    Enum.map(children, fn {k, count} ->
      [count * multiplier, count_children(rules, k, count * multiplier)]
    end)
  end
end