defmodule Day7 do
  @example """
  light red bags contain 1 bright white bag, 2 muted yellow bags.
  dark orange bags contain 3 bright white bags, 4 muted yellow bags.
  bright white bags contain 1 shiny gold bag.
  muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
  shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
  dark olive bags contain 3 faded blue bags, 4 dotted black bags.
  vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
  faded blue bags contain no other bags.
  dotted black bags contain no other bags.
  """

  @example2 """
  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.
  """

  def part1(example \\ false) do
    rules =
      if(example, do: @example, else: File.read!("lib/day7/input.txt"))
      |> parse_rules()

    Map.keys(rules)
    |> Enum.count(fn k ->
      can_contain?(k, "shiny gold", rules)
    end)
  end

  def part2(example \\ false) do
    rules =
      case example do
        false -> File.read!("lib/day7/input.txt")
        1 -> @example
        2 -> @example2
      end
      |> parse_rules()

    # don't count the outermost shiny gold bag itself
    count_contained(rules, "shiny gold") - 1
  end

  def parse_rules(input) do
    input
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(&parse_rule/1)
    |> Map.new()
  end

  def parse_rule(rule) do
    [_, rule_color, rest] = Regex.run(~r/^(\w+ \w+) bags contain (no other bags|.+)\.$/, rule)

    contains =
      case rest do
        "no other bags" ->
          []

        rest ->
          rest
          |> String.split(",")
          |> Enum.map(fn str ->
            str = String.trim(str)
            [_, qty, color] = Regex.run(~r/^(\d+) (\w+ \w+) bags?$/, str)
            qty = String.to_integer(qty)
            {qty, color}
          end)
      end

    {rule_color, contains}
  end

  def can_contain?(outer, inner, rules) do
    outer_rules = Map.fetch!(rules, outer)

    Enum.any?(outer_rules, fn {_, color} ->
      color == inner || can_contain?(color, inner, rules)
    end)
  end

  def count_contained(rules, color) do
    nested =
      Map.fetch!(rules, color)
      |> Enum.map(fn {qty, contained_color} ->
        qty * count_contained(rules, contained_color)
      end)
      |> Enum.sum()

    1 + nested
  end
end