defmodule MonkeyUtil do
  def parse_op(z, s) do
    "Operation: new = " <> z, s -> parse_op(z, s)
    "old " <> z, s -> parse_op(z, [:old | s])
    "old" <> z, s -> parse_op(z, [:old | s])
    "+ " <> z, s -> parse_op(z, [:plus | s])
    "- " <> z, s -> parse_op(z, [:minus | s])
    "* " <> z, s -> parse_op(z, [:mul | s])
    "/ " <> z, s -> parse_op(z, [:div | s])
    "", _, s -> s |> Enum.reverse()
    n, s ->
      [a | b] = String.split( n, " ")
      a = String.trim(a) |> String.to_integer()
      parse_op(b |> Enum.join(" "), [a | s])
  end

  def subs(b, a), do: if b == :old, do: a, else: b
  def resolve_stack(x, y) do
    (fn
    old, [a, :plus, b] -> subs(a, old) + subs(b, old)
    old, [a, :minus, b] -> subs(a, old) - subs(b, old)
    old, [a, :mul, b] -> subs(a, old) * subs(b, old)
    old, [a, :div, b] -> subs(a, old) / subs(b, old) end).(x, y)
  end
end

defmodule Monkey do
  defstruct id: nil, items: nil, op: nil, mod: nil, if_true: nil, if_false: nil, inspections: 0
end

defmodule MonkeyMachine do
  def load(instructions) do
    {barrel, _} = for instruction <- instructions, reduce: {%{}, nil} do
      {barrel, monkey} ->
        (fn
          {:break} ->
          {barrel |> Map.put(monkey.id, monkey), nil}
          {:add_monkey, id} -> {barrel, %Monkey{id: id}}
          {:starting_items, i} -> {barrel, %{monkey | items: i }}
          {a, o} -> {barrel, %{monkey | a => o }}
        end).(instruction)
    end

    barrel
  end

  def run(barrel, :part1) do
    run(barrel, 0, (Map.keys(barrel) |> Enum.count()) - 1)
  end

  def run(barrel, :part2) do
    run(barrel, 0, (Map.keys(barrel) |> Enum.count()) - 1, Map.values(barrel) |> Stream.map(&(&1.mod)) |> Enum.product())
  end

  defp run(barrel, i, n) when i > n, do: barrel
  defp run(barrel, id, n) do
    monkey = Map.get(barrel, id)
    barrel = for item <- monkey.items, reduce: barrel do
      b ->
        item = div(monkey.op.(item), 3)
        dest = if rem(item, monkey.mod) == 0, do: monkey.if_true, else: monkey.if_false

        b
          |> Map.update!(dest, fn d -> %{d | items: [item | d.items]} end)
          |> Map.update!(id, fn m -> Map.update!(m, :inspections, &(&1 + 1)) end)
        end |> Map.update!(id, fn m -> %{m | items: []} end)

    run(barrel, id + 1, n)
  end

  defp run(barrel, i, n, _) when i > n, do: barrel
  defp run(barrel, id, n, m) do
    monkey = Map.get(barrel, id)
    barrel = for item <- monkey.items, reduce: barrel do
      b ->
        r = rem(monkey.op.(item), m)
        item = if r == 0, do: m, else: r
        dest = if rem(r, monkey.mod) == 0, do: monkey.if_true, else: monkey.if_false

        b
          |> Map.update!(dest, fn d -> %{d | items: [item | d.items]} end)
          |> Map.update!(id, fn m -> Map.update!(m, :inspections, &(&1 + 1)) end)
        end |> Map.update!(id, fn m -> %{m | items: []} end)

    run(barrel, id + 1, n, m)
  end
end

defmodule MonkeySolver do
  # part := :part1 || :part2
  def solve(input, part) do
    barrel = input
    |> String.split("\n")
    |> Stream.map(&String.trim/1)
    |> Enum.map(
      fn
        "Monkey " <> z -> {:add_monkey, z |> String.split(":") |> Enum.at(0) |> String.to_integer()}
        "Starting items: " <> i -> {:starting_items, i |> String.split(",", trim: true) |> Enum.map(&String.to_integer(String.trim(&1)))}
        "Operation: " <> _ = u -> {:op, fn o -> MonkeyUtil.resolve_stack(o, MonkeyUtil.parse_op(u, [])) end}
        "Test: divisible by " <> n -> {:mod, String.to_integer(n)}
        "If true: throw to monkey " <> n -> {:if_true, String.to_integer(n)}
        "If false: throw to monkey " <> n -> {:if_false, String.to_integer(n)}
        "" -> {:break}
        u -> u
      end
    )
    |> MonkeyMachine.load()

    iterations = if part == :part1, do: 20, else: 10_000 end

    for _ <- 1..iterations, reduce: barrel do
      b -> MonkeyMachine.run(b, part)
    end
    |> Map.values()
    |> Stream.map(&(&1.inspections))
    |> Enum.sort(:desc)
    |> Enum.take(2)
    |> Enum.product()
  end
end