defmodule AOC.Day16 do
  use AOC.Day, day: 16

  def parse_input(lines) do
    lines
    |> Enum.with_index()
    |> Enum.flat_map(fn {line, y} ->
      line
      |> String.to_charlist()
      |> Enum.with_index()
      |> Enum.map(fn {space, x} ->
        {{y, x}, space}
      end)
    end)
    |> Enum.into(%{})
  end

  def new_location({y, x}, :north), do: {y - 1, x}
  def new_location({y, x}, :east), do: {y, x + 1}
  def new_location({y, x}, :south), do: {y + 1, x}
  def new_location({y, x}, :west), do: {y, x - 1}

  def new_directions(:north, ?\\), do: [:west]
  def new_directions(:north, ?/), do: [:east]
  def new_directions(:north, ?-), do: [:east, :west]
  def new_directions(:east, ?\\), do: [:south]
  def new_directions(:east, ?/), do: [:north]
  def new_directions(:east, ?|), do: [:south, :north]
  def new_directions(:south, ?\\), do: [:east]
  def new_directions(:south, ?/), do: [:west]
  def new_directions(:south, ?-), do: [:east, :west]
  def new_directions(:west, ?\\), do: [:north]
  def new_directions(:west, ?/), do: [:south]
  def new_directions(:west, ?|), do: [:south, :north]
  def new_directions(direction, _), do: [direction]

  def simulate_laser(grid, location, direction) do
    case :ets.lookup(:laser_memo, {location, direction}) do
      [] ->
        :ets.insert_new(:laser_memo, {{location, direction}, true})
        new_location = new_location(location, direction)

        case Map.fetch(grid, new_location) do
          {:ok, symbol} ->
            new_directions(direction, symbol)
            |> Enum.each(&simulate_laser(grid, new_location, &1))

          :error ->
            :ok
        end

      _ ->
        :ok
    end
  end

  def count_energized(grid, start_location, direction) do
    :ets.new(:laser_memo, [:named_table])

    simulate_laser(grid, start_location, direction)

    count =
      :ets.tab2list(:laser_memo)
      |> Enum.map(&elem(&1, 0))
      |> Enum.map(&elem(&1, 0))
      |> Enum.uniq()
      |> Enum.count()

    :ets.delete(:laser_memo)

    count - 1
  end

  def part1(grid) do
    count_energized(grid, {0, -1}, :east)
  end

  def part2(grid) do
    max_x =
      grid
      |> Enum.map(&elem(&1, 0))
      |> Enum.map(&elem(&1, 1))
      |> Enum.max()

    max_y =
      grid
      |> Enum.map(&elem(&1, 0))
      |> Enum.map(&elem(&1, 0))
      |> Enum.max()

    [
      Enum.map(0..max_x, fn x -> {{-1, x}, :south} end),
      Enum.map(0..max_x, fn x -> {{max_y + 1, x}, :north} end),
      Enum.map(0..max_y, fn y -> {{y, -1}, :east} end),
      Enum.map(0..max_y, fn y -> {{y, max_x + 1}, :west} end)
    ]
    |> List.flatten()
    |> Enum.map(fn {position, direction} ->
      count_energized(grid, position, direction)
    end)
    |> Enum.max()
  end
end