defmodule AOC.Day15 do
  use AOC.Day, day: 15

  def parse_input(lines) do
    lines
    |> hd()
    |> String.split(",")
  end

  def part1(init_seq) do
    init_seq
    |> Enum.map(&String.to_charlist/1)
    |> Enum.map(&hash/1)
    |> Enum.sum()
  end

  def part2(init_seq) do
    init_seq
    |> Enum.map(fn step ->
      if String.contains?(step, "=") do
        [label, focal_length] = String.split(step, "=")
        {:add, String.to_charlist(label), String.to_integer(focal_length)}
      else
        label = step |> String.trim_trailing("-") |> String.to_charlist()
        {:remove, label}
      end
    end)
    |> Enum.reduce(%{}, fn
      {:remove, label}, acc ->
        box_number = hash(label)

        case Map.fetch(acc, box_number) do
          {:ok, lenses} ->
            %{acc | box_number => Enum.reject(lenses, &(elem(&1, 0) == label))}

          :error ->
            acc
        end

      {:add, label, focal_length}, acc ->
        box_number = hash(label)

        case Map.fetch(acc, box_number) do
          {:ok, lenses} ->
            case Enum.find_index(lenses, &(elem(&1, 0) == label)) do
              nil ->
                %{acc | box_number => [{label, focal_length} | lenses]}

              index ->
                %{acc | box_number => List.replace_at(lenses, index, {label, focal_length})}
            end

          :error ->
            Map.put(acc, box_number, [{label, focal_length}])
        end
    end)
    |> Enum.flat_map(fn {box_number, lenses} ->
      lenses
      |> Enum.reverse()
      |> Enum.with_index()
      |> Enum.map(fn {{_label, focal_length}, index} ->
        (box_number + 1) * (index + 1) * focal_length
      end)
    end)
    |> Enum.sum()
  end

  def hash(step), do: hash(step, 0)
  def hash([], acc), do: acc
  def hash([hd | tl], acc), do: hash(tl, rem(17 * (hd + acc), 256))
end