defmodule Absinthe.Middleware.Batch do
  @moduledoc """
  Batch the resolution of multiple fields.

  ## Motivation
  Consider the following graphql query:
  ```
  {
    posts {
      author {
        name
      }
    }
  }
  ```

  `posts` returns a list of `post` objects, which has an associated `author` field.
  If the `author` field makes a call to the database we have the classic N + 1 problem.
  What we want is a way to load all authors for all posts in one database request.

  This plugin provides this, without any eager loading at the parent level. That is,
  the code for the `posts` field does not need to do anything to facilitate the
  efficient loading of its children.

  ## Example Usage
  The API for this plugin is a little on the verbose side because it is not specific
  to any particular batching mechanism. That is, this API is just as useful for an Ecto
  based DB as it is for talking to S3 or the File System. Thus we anticipate people
  (including ourselves) will be creating additional functions more tailored to each
  of those specific use cases.

  Here is an example using the `Absinthe.Resolution.Helpers.batch/3` helper.
  ```elixir
  object :post do
    field :name, :string
    field :author, :user do
      resolve fn post, _, _ ->
        batch({__MODULE__, :users_by_id}, post.author_id, fn batch_results ->
          {:ok, Map.get(batch_results, post.author_id)}
        end)
      end
    end
  end

  def users_by_id(_, user_ids) do
    users = Repo.all from u in User, where: u.id in ^user_ids
    Map.new(users, fn user -> {user.id, user} end)
  end
  ```

  Let's look at this piece by piece:
  - `{__MODULE__, :users_by_id}`: is the batching function which will be used. It must
  be a 2 arity function. For details see the `batch_fun` typedoc.
  - `post.author_id`: This is the information to be aggregated. The aggregated values
  are the second argument to the batching function.
  - `fn batch_results`: This function takes the results from the batching function.
  it should return one of the resolution function values.
  """

  @behaviour Absinthe.Middleware
  @behaviour Absinthe.Plugin

  @typedoc """
  The function to be called with the aggregate batch information.

  It comes in both a 2 tuple and 3 tuple form. The first two elements are the module
  and function name. The third element is an arbitrary parameter that is passed
  as the first argument to the batch function.

  For example, one could parameterize the `users_by_id` function from the moduledoc
  to make it more generic. Instead of doing `{__MODULE__, :users_by_id}` you could do
  `{__MODULE__, :by_id, User}`. Then the function would be:

  ```elixir
  def by_id(model, ids) do
    model
    |> where([m], m.id in ^ids)
    |> Repo.all()
    |> Map.new(&{&1.id, &1})
  end
  ```
  It could also be used to set options unique to the execution of a particular
  batching function.
  """
  @type batch_fun :: {module, atom} | {module, atom, term}

  @type post_batch_fun :: (term -> Absinthe.Type.Field.result())

  def before_resolution(exec) do
    case exec.acc do
      %{__MODULE__ => _} ->
        put_in(exec.acc[__MODULE__][:input], [])

      _ ->
        put_in(exec.acc[__MODULE__], %{input: [], output: %{}})
    end
  end

  def call(%{state: :unresolved} = res, {batch_key, field_data, post_batch_fun, batch_opts}) do
    acc = res.acc

    acc =
      update_in(acc[__MODULE__][:input], fn
        nil -> [{{batch_key, batch_opts}, field_data}]
        data -> [{{batch_key, batch_opts}, field_data} | data]
      end)

    %{
      res
      | state: :suspended,
        middleware: [{__MODULE__, {batch_key, post_batch_fun}} | res.middleware],
        acc: acc
    }
  end

  def call(%{state: :suspended} = res, {batch_key, post_batch_fun}) do
    batch_data_for_fun =
      res.acc
      |> Map.fetch!(__MODULE__)
      |> Map.fetch!(:output)
      |> Map.fetch!(batch_key)

    res
    |> Absinthe.Resolution.put_result(post_batch_fun.(batch_data_for_fun))
  end

  def after_resolution(exec) do
    output = do_batching(exec.acc[__MODULE__][:input])
    put_in(exec.acc[__MODULE__][:output], output)
  end

  defp do_batching(input) do
    input
    |> Enum.group_by(&elem(&1, 0), &elem(&1, 1))
    |> Enum.map(fn {{batch_fun, batch_opts}, batch_data} ->
      system_time = System.system_time()
      start_time_mono = System.monotonic_time()

      task =
        Task.async(fn ->
          {batch_fun, call_batch_fun(batch_fun, batch_data)}
        end)

      metadata = emit_start_event(system_time, batch_fun, batch_opts, batch_data)

      {batch_opts, task, start_time_mono, metadata}
    end)
    |> Map.new(fn {batch_opts, task, start_time_mono, metadata} ->
      timeout = Keyword.get(batch_opts, :timeout, 5_000)
      result = Task.await(task, timeout)

      duration = System.monotonic_time() - start_time_mono
      emit_stop_event(duration, metadata, result)

      result
    end)
  end

  @batch_start [:absinthe, :middleware, :batch, :start]
  @batch_stop [:absinthe, :middleware, :batch, :stop]
  defp emit_start_event(system_time, batch_fun, batch_opts, batch_data) do
    id = :erlang.unique_integer()

    metadata = %{
      id: id,
      telemetry_span_context: id,
      batch_fun: batch_fun,
      batch_opts: batch_opts,
      batch_data: batch_data
    }

    :telemetry.execute(
      @batch_start,
      %{system_time: system_time},
      metadata
    )

    metadata
  end

  defp emit_stop_event(duration, metadata, result) do
    :telemetry.execute(
      @batch_stop,
      %{duration: duration},
      Map.put(metadata, :result, result)
    )
  end

  defp call_batch_fun({module, fun}, batch_data) do
    call_batch_fun({module, fun, []}, batch_data)
  end

  defp call_batch_fun({module, fun, config}, batch_data) do
    apply(module, fun, [config, batch_data])
  end

  # If the flag is set we need to do another resolution phase.
  # otherwise, we do not
  def pipeline(pipeline, exec) do
    case exec.acc[__MODULE__][:input] do
      [_ | _] ->
        [Absinthe.Phase.Document.Execution.Resolution | pipeline]

      _ ->
        pipeline
    end
  end
end