defmodule GenStage.PartitionDispatcher do
  @moduledoc """
  A dispatcher that sends events according to partitions.

  This dispatcher assumes that partitions are *evenly distributed*.
  See the ["Even distribution"](#module-even-distribution) section for
  more information.

  When multiple consumers subscribe to one partition, the producer
  behaves like a `GenStage.DemandDispatcher` *within that partition*.

  ## Options

  The partition dispatcher accepts the following options
  on initialization:

    * `:partitions` - the number of partitions to dispatch to. It may be
      an integer with a total number of partitions, where each partition
      is named from 0 up to `integer - 1`. For example, `partitions: 4`
      will contain four partitions named `0`, `1`, `2` and `3`.

      It may also be an *enumerable* that specifies the name of each partition.
      For instance, `partitions: [:odd, :even]` will build two partitions,
      named `:odd` and `:even`.

    * `:hash` - the hashing algorithm. It's a function of type
      `t:hash_function/0`, which receives the event and returns a tuple with two
      elements: the event to be dispatched and the partition to dispatch it to.
      The function can also return `:none`, in which case the event
      is discarded. The partition must be one of the partitions specified in
      `:partitions` above. The default uses:

          fn event -> {event, :erlang.phash2(event, Enum.count(partitions))} end

  ### Examples

  To start a producer with four partitions named `0`, `1`, `2`, and `3`:

      {:producer, state, dispatcher: {GenStage.PartitionDispatcher, partitions: 0..3}}

  To start a producer with two partitions named `:odd` and `:even`:

      {:producer, state, dispatcher: {GenStage.PartitionDispatcher, partitions: [:odd, :even]}}

  ## Subscribe options

  When subscribing to a `GenStage` with a partition dispatcher the following
  option is required:

    * `:partition` - the name of the partition. The partition must be one of
      the partitions specified in `:partitions` above.

  ### Examples

  The partition function can be given either on `init`'s subscribe_to:

      {:consumer, :ok, subscribe_to: [{producer, partition: 0}]}

  Or when calling `sync_subscribe`:

      GenStage.sync_subscribe(consumer, to: producer, partition: 0)

  ## Even distribution

  This dispatcher assumes that partitions are *evenly distributed*.
  If the data is uneven for long periods of time, then you may
  buffer excessive data from busy partitions for long periods of
  time. This happens because the producer is unable to distinguish
  from which particular consumer/partition demand arrives.

  Let's see an example. Imagine you have three consumers, each
  for one partition: `A`, `B`, and `C`.

  Let's assume 60% of the data goes to `A`, 20% to `B`, and 20% to
  `C`. Let's also say that `max_demand` is `10` and `min_demand` is
  `5`. When the consumers initially request data (`10` events each),
  the producer receives a total demand of `30`. A will receive `18` of
  those (60%), while `B` and `C` receive `6` each (20%). After
  processing `5` events (the `min_demand`), each consumer requests
  additional `5` events, for a total of `15` additional events. At
  this point, that will be `9` additional elements for A, and 3
  additional elements for B and C. At the end of these two rounds, we
  will have:

      A = 18 - 5 + 9 = 22 events
      B = 6 - 5 + 3 = 4 events
      C = 6 - 5 + 3 = 4 events

  Furthermore, as B and C request more items, A will only go further
  behind. This behaviour is fine for spikes that should quickly
  resolve, but it can be problematic if the data is consistently uneven.
  """

  @typedoc """
  The type used for the function passed to the `:hash` option.
  """
  @typedoc since: "1.2.0"
  @type hash_function :: (event :: any -> {event :: any, partition :: any} | :none)

  @behaviour GenStage.Dispatcher
  @init {nil, nil, 0}

  require Logger

  @doc false
  def init(opts) do
    partitions =
      case Keyword.get(opts, :partitions) do
        nil ->
          raise ArgumentError,
                "the enumerable of :partitions is required when using the partition dispatcher"

        partitions when is_integer(partitions) ->
          0..(partitions - 1)

        partitions ->
          partitions
      end

    hash_present? = Keyword.has_key?(opts, :hash)

    partitions =
      for partition <- partitions, into: %{} do
        if not hash_present? and not is_integer(partition) do
          raise ArgumentError,
                "when :partitions contains partitions that are not integers, you have to pass " <>
                  "in the :hash option as well"
        end

        Process.put(partition, [])
        {partition, @init}
      end

    size = map_size(partitions)
    hash = Keyword.get(opts, :hash, &hash(&1, size))
    {:ok, {make_ref(), hash, 0, 0, partitions, %{}, %{}}}
  end

  defp hash(event, range) do
    {event, :erlang.phash2(event, range)}
  end

  @doc false
  def info(msg, {tag, hash, waiting, pending, partitions, references, infos}) do
    info = make_ref()

    {partitions, queued} =
      Enum.reduce(partitions, {partitions, []}, fn
        {partition, {pid, ref, queue}}, {partitions, queued} when not is_integer(queue) ->
          {Map.put(partitions, partition, {pid, ref, :queue.in({tag, info}, queue)}),
           [partition | queued]}

        _, {partitions, queued} ->
          {partitions, queued}
      end)

    infos =
      case queued do
        [] ->
          send(self(), msg)
          infos

        _ ->
          Map.put(infos, info, {msg, queued})
      end

    {:ok, {tag, hash, waiting, pending, partitions, references, infos}}
  end

  @doc false
  def subscribe(opts, {pid, ref}, {tag, hash, waiting, pending, partitions, references, infos}) do
    partition = Keyword.get(opts, :partition)

    case partitions do
      %{^partition => {nil, nil, demand_or_queue}} ->
        partitions = Map.put(partitions, partition, {pid, ref, demand_or_queue})
        references = Map.put(references, ref, partition)
        {:ok, 0, {tag, hash, waiting, pending, partitions, references, infos}}

      %{^partition => {pid, _, _}} ->
        raise ArgumentError, "the partition #{partition} is already taken by #{inspect(pid)}"

      _ when is_nil(partition) ->
        raise ArgumentError,
              "the :partition option is required when subscribing to a producer with partition dispatcher"

      _ ->
        keys = Map.keys(partitions)
        raise ArgumentError, ":partition must be one of #{inspect(keys)}, got: #{partition}"
    end
  end

  @doc false
  def cancel({_, ref}, {tag, hash, waiting, pending, partitions, references, infos}) do
    {partition, references} = Map.pop(references, ref)
    {_pid, _ref, demand_or_queue} = Map.get(partitions, partition)
    partitions = Map.put(partitions, partition, @init)

    case demand_or_queue do
      demand when is_integer(demand) ->
        {:ok, 0, {tag, hash, waiting, pending + demand, partitions, references, infos}}

      queue ->
        {length, infos} = clear_queue(queue, tag, partition, 0, infos)
        {:ok, length, {tag, hash, waiting + length, pending, partitions, references, infos}}
    end
  end

  @doc false
  def ask(counter, {_, ref}, {tag, hash, waiting, pending, partitions, references, infos}) do
    partition = Map.fetch!(references, ref)
    {pid, ref, demand_or_queue} = Map.fetch!(partitions, partition)

    {demand_or_queue, infos} =
      case demand_or_queue do
        demand when is_integer(demand) ->
          {demand + counter, infos}

        queue ->
          send_from_queue(queue, tag, pid, ref, partition, counter, [], infos)
      end

    partitions = Map.put(partitions, partition, {pid, ref, demand_or_queue})
    already_sent = min(pending, counter)
    demand = counter - already_sent
    pending = pending - already_sent
    {:ok, demand, {tag, hash, waiting + demand, pending, partitions, references, infos}}
  end

  defp send_from_queue(queue, _tag, pid, ref, _partition, 0, acc, infos) do
    maybe_send(acc, pid, ref)
    {queue, infos}
  end

  defp send_from_queue(queue, tag, pid, ref, partition, counter, acc, infos) do
    case :queue.out(queue) do
      {{:value, {^tag, info}}, queue} ->
        maybe_send(acc, pid, ref)
        infos = maybe_info(infos, info, partition)
        send_from_queue(queue, tag, pid, ref, partition, counter, [], infos)

      {{:value, event}, queue} ->
        send_from_queue(queue, tag, pid, ref, partition, counter - 1, [event | acc], infos)

      {:empty, _queue} ->
        maybe_send(acc, pid, ref)
        {counter, infos}
    end
  end

  defp clear_queue(queue, tag, partition, counter, infos) do
    case :queue.out(queue) do
      {{:value, {^tag, info}}, queue} ->
        clear_queue(queue, tag, partition, counter, maybe_info(infos, info, partition))

      {{:value, _}, queue} ->
        clear_queue(queue, tag, partition, counter + 1, infos)

      {:empty, _queue} ->
        {counter, infos}
    end
  end

  # Important: events must be in reverse order
  defp maybe_send([], _pid, _ref), do: :ok

  defp maybe_send(events, pid, ref),
    do: Process.send(pid, {:"$gen_consumer", {self(), ref}, :lists.reverse(events)}, [:noconnect])

  defp maybe_info(infos, info, partition) do
    case infos do
      %{^info => {msg, [^partition]}} ->
        send(self(), msg)
        Map.delete(infos, info)

      %{^info => {msg, partitions}} ->
        Map.put(infos, info, {msg, List.delete(partitions, partition)})
    end
  end

  @doc false
  def dispatch(events, _length, {tag, hash, waiting, pending, partitions, references, infos}) do
    {deliver_later, waiting} = split_events(events, waiting, hash, partitions)

    partitions =
      partitions
      |> :maps.to_list()
      |> dispatch_per_partition()
      |> :maps.from_list()

    {:ok, deliver_later, {tag, hash, waiting, pending, partitions, references, infos}}
  end

  defp split_events(events, 0, _hash, _partitions), do: {events, 0}
  defp split_events([], counter, _hash, _partitions), do: {[], counter}

  defp split_events([event | events], counter, hash, partitions) do
    case hash.(event) do
      {event, partition} ->
        case :erlang.get(partition) do
          :undefined ->
            raise "unknown partition #{inspect(partition)} computed for GenStage event " <>
                    "#{inspect(event)}. The known partitions are #{inspect(Map.keys(partitions))}. " <>
                    "See the :partitions option to set your own. This event has been discarded."

          current ->
            Process.put(partition, [event | current])
            split_events(events, counter - 1, hash, partitions)
        end

      :none ->
        split_events(events, counter, hash, partitions)

      other ->
        raise "the :hash function should return {event, partition}, got: #{inspect(other)}"
    end
  end

  defp dispatch_per_partition([{partition, {pid, ref, demand_or_queue} = value} | rest]) do
    case Process.put(partition, []) do
      [] ->
        [{partition, value} | dispatch_per_partition(rest)]

      events ->
        events = :lists.reverse(events)

        {events, demand_or_queue} =
          case demand_or_queue do
            demand when is_integer(demand) ->
              split_into_queue(events, demand, [])

            queue ->
              {[], put_into_queue(events, queue)}
          end

        maybe_send(events, pid, ref)
        [{partition, {pid, ref, demand_or_queue}} | dispatch_per_partition(rest)]
    end
  end

  defp dispatch_per_partition([]) do
    []
  end

  defp split_into_queue(events, 0, acc), do: {acc, put_into_queue(events, :queue.new())}
  defp split_into_queue([], counter, acc), do: {acc, counter}

  defp split_into_queue([event | events], counter, acc),
    do: split_into_queue(events, counter - 1, [event | acc])

  defp put_into_queue(events, queue) do
    Enum.reduce(events, queue, &:queue.in/2)
  end
end
