defmodule AOC.Day10 do
  def find_diffs(xs), do: find_diffs(xs, %{})
  def find_diffs([_], acc), do: acc

  def find_diffs([x1, x2 | xs], acc) do
    acc = Map.update(acc, x2 - x1, 1, &(&1 + 1))

    find_diffs([x2 | xs], acc)
  end

  def count_walks([x | _] = xs) do
    memo =
      Enum.into(xs, %{}, &{&1, 0})
      |> Map.put(x, 1)

    count_walks(xs, memo)
  end

  def count_walks([x], memo), do: memo[x]

  def count_walks([x | xs], memo) do
    memo =
      Enum.take(xs, 3)
      |> Enum.reduce(memo, fn
        next, memo when next - x <= 3 ->
          Map.update!(memo, next, &(&1 + memo[x]))

        _, memo ->
          memo
      end)

    count_walks(xs, memo)
  end

  def parts do
    input =
      AOC.Util.input_integers(10, 1)
      |> Enum.sort()

    input = [0 | input] ++ [Enum.max(input) + 3]

    %{1 => ones, 3 => threes} = find_diffs(input)
    part1 = ones * threes
    part2 = count_walks(input)

    {part1, part2}
  end
end