defmodule Aoc2024.Day06 do
  def parse(input) do
    input = input |> String.trim() |> String.split("\n") |> Enum.map(&String.codepoints/1)
    w = Enum.count(Enum.at(input, 0))
    h = Enum.count(input)
    walk(input, w, h)
  end

  defp walk(t, w, h), do: walk(t, 0, 0, w, h, MapSet.new(), nil)
  defp walk(_t, i, _j, w, h, walls, start) when i==h, do: {w, h, walls, start}
  defp walk(t, i, j, w, h, walls, start) when j==w, do: walk(t, i+1, 0, w, h, walls, start)
  defp walk(t, i, j, w, h, walls, start) do
    case t |> Enum.at(i) |> Enum.at(j) do
      "^" -> walk(t, i, j+1, w, h, walls, {i, j})
      "#" -> walk(t, i, j+1, w, h, MapSet.put(walls, {i, j}), start)
      _ -> walk(t, i, j+1, w, h, walls, start)
    end
  end

  def part1({w, h, walls, {i, j}}) do
    guard(i, j, w, h, walls)
  end
  defp guard(i, j, w, h, walls), do: guard(i, j, -1, 0, w, h, walls, MapSet.new())
  defp guard(i, j, _, _, w, h, _, path) when i<0 or j<0 or i==h or j==w, do: path |> Enum.count()
  defp guard(i, j, di, dj, w, h, walls, path) do 
    if MapSet.member?(walls, {i, j}) do
      i = i-di
      j = j-dj
      {di, dj} = {dj, di} 
      dj = -dj
      guard(i, j, di, dj, w, h, walls, path)
    else 
      guard(i+di, j+dj, di, dj, w, h, walls, MapSet.put(path, {i, j}))
    end
  end

  def part2({w, h, walls, {i, j}}) do
    for wi <- 0..h-1, wj <- 0..w-1 do
      Task.async(fn -> guard2(i, j, w, h, walls |> MapSet.put({wi, wj})) end)
    end
      |> Enum.map(&Task.await/1)
      |> Enum.sum()
  end
  defp guard2(i, j, w, h, walls), do: guard2(i, j, -1, 0, w, h, walls, MapSet.new())
  defp guard2(i, j, _, _, w, h, _, _) when i<0 or j<0 or i==h or j==w, do: 0
  defp guard2(i, j, di, dj, w, h, walls, path) do 
    cond do
      MapSet.member?(path, {i, j, di, dj}) -> 
        1
      MapSet.member?(walls, {i, j}) ->
        i = i-di
        j = j-dj
        {di, dj} = {dj, di} 
        dj = -dj
        guard2(i, j, di, dj, w, h, walls, path)
      true ->
        guard2(i+di, j+dj, di, dj, w, h, walls, MapSet.put(path, {i, j, di, dj}))
    end
  end
end