defmodule Aoc2023.Day12 do
  use Aoc2023

  defmodule Cache do
    use GenServer

    def start do
      {:ok, pid} = GenServer.start_link(Cache, nil)
      pid
    end

    def memoize(cache, key, fun) do
      got = GenServer.call(cache, {:get, key})
      if got == nil do
        val = fun.(key)
        GenServer.call(cache, {:set, key, val})
        val
      else
        got
      end
    end

    @impl true
    def init(_) do
      {:ok, %{}}
    end

    @impl true
    def handle_call({:get, key}, _, state) do
      case state do
        %{^key => value} ->
          {:reply, value, state}

        _ ->
          {:reply, nil, state}
      end
    end

    @impl true
    def handle_call({:set, key, value}, _, state) do
      {:reply, nil, Map.put_new(state, key, value)}
    end
  end

  def parse(input) do
    input
    |> lines
    |> Enum.map(
      pipe(
        String.split()
        |> then(fn [rec, blocks] ->
          {rec |> String.to_charlist(),
           blocks |> String.split(",") |> Enum.map(&String.to_integer/1)}
        end)
      )
    )
  end

  defp possibilities({rec, blocks}) do
    loop = fn recur, arg ->
      case arg do
        {rec, []} -> if Enum.all?(rec, &(&1 in [??, ?.])) do 1 else 0 end
        {rec, [block | blocks]} ->
          # Expect rec to start with any number of ./?s up to the max allowing all remaining groups
          # i.e. their lengths + an undamaged seperator
          0..(length(rec) - Enum.sum(blocks) - length(blocks) - block)
          |> Enum.map(fn i -> repeat('.', i) ++ repeat('#', block) ++ '.' end)
          |> Enum.filter(pipe(Enum.zip_with(rec, &(&1 == &2 || &2 == ??)) |> Enum.all?))
          |> Enum.map(&(recur.(recur, {Enum.drop(rec, length(&1)), blocks})))
          |> Enum.sum
      end
    end
    cache = Cache.start()
    memoized = fn recur, arg -> Cache.memoize(cache, arg, fn arg -> loop.(recur, arg) end) end
    memoized.(memoized, {rec, blocks})
  end

  def part1(input) do
    input
    |> Enum.map(pipe(possibilities()))
    |> Enum.sum()
  end

  def part2(input) do
    input
    |> Enum.map(fn {rec, blocks} -> {
        [rec] |> repeat(5) |> Enum.intersperse('?') |> Enum.flat_map(&id/1),
        repeat(blocks, 5)
      } end)
    |> Enum.map(&possibilities/1)
    |> Enum.sum()
  end
end