defmodule AOC.Day14 do
  use AOC.Day, day: 14

  @compile {:parse_transform, :ms_transform}

  def parse_input(lines) do
    dish_table = :ets.new(:dish_table, [])

    lines
    |> Enum.with_index()
    |> Enum.each(fn {line, y} ->
      line
      |> String.to_charlist()
      |> Enum.with_index()
      |> Enum.each(fn
        {?O, x} -> :ets.insert(dish_table, {{y, x}, :round})
        {?#, x} -> :ets.insert(dish_table, {{y, x}, :cube})
        {?., _x} -> :ok
      end)
    end)

    max_y =
      :ets.tab2list(dish_table)
      |> Enum.map(&elem(&1, 0))
      |> Enum.map(&elem(&1, 0))
      |> Enum.max()

    max_x =
      :ets.tab2list(dish_table)
      |> Enum.map(&elem(&1, 0))
      |> Enum.map(&elem(&1, 1))
      |> Enum.max()

    {dish_table, {max_y, max_x}}
  end

  def get_round_rock_coords(dish_table) do
    :ets.select(dish_table, :ets.fun2ms(fn {coords, type} when type == :round -> coords end))
  end

  def perform_spin_cycle(dish_table, max_coords) do
    slide(dish_table, max_coords, fn {y1, _}, {y2, _} -> y1 <= y2 end, fn {y, x} ->
      {y - 1, x}
    end)

    slide(dish_table, max_coords, fn {_, x1}, {_, x2} -> x1 <= x2 end, fn {y, x} ->
      {y, x - 1}
    end)

    slide(dish_table, max_coords, fn {y1, _}, {y2, _} -> y1 >= y2 end, fn {y, x} ->
      {y + 1, x}
    end)

    slide(dish_table, max_coords, fn {_, x1}, {_, x2} -> x1 >= x2 end, fn {y, x} ->
      {y, x + 1}
    end)
  end

  def slide(dish_table, max_coords, sort_fun, slide_fun) do
    get_round_rock_coords(dish_table)
    |> Enum.sort(sort_fun)
    |> Enum.each(&slide_rock(dish_table, max_coords, slide_fun, &1))
  end

  def slide_rock(dish_table, {max_y, max_x} = max_coords, slide_fun, coords) do
    {slide_y, slide_x} = slide_coords = slide_fun.(coords)

    if not (slide_y < 0 or slide_x < 0 or slide_y > max_y or slide_x > max_x) do
      case :ets.lookup(dish_table, slide_coords) do
        [] ->
          :ets.delete(dish_table, coords)
          :ets.insert(dish_table, {slide_coords, :round})
          slide_rock(dish_table, max_coords, slide_fun, slide_coords)

        _ ->
          :ok
      end
    end
  end

  def calculate_total_load(dish_table, {max_y, _}) do
    get_round_rock_coords(dish_table)
    |> Enum.map(&elem(&1, 0))
    |> Enum.map(&(max_y + 1 - &1))
    |> Enum.sum()
  end

  def part1({dish_table, max_coords}) do
    slide(dish_table, max_coords, fn {y1, _}, {y2, _} -> y1 <= y2 end, fn {y, x} ->
      {y - 1, x}
    end)

    calculate_total_load(dish_table, max_coords)
  end

  def part2({dish_table, max_coords}) do
    memo_table = :ets.new(:memo_table, [])

    # Adjust according to calculations
    cycle_count = 109

    Enum.each(1..cycle_count, fn i ->
      perform_spin_cycle(dish_table, max_coords)
      memo = get_round_rock_coords(dish_table) |> Enum.sort()

      case :ets.lookup(memo_table, memo) do
        [] -> :ets.insert(memo_table, {memo, i})
        [{_, i_prev}] -> IO.puts("State already observed: #{i_prev} & #{i}")
      end
    end)

    calculate_total_load(dish_table, max_coords)
  end
end