defmodule NSQ.Consumer.RDY do
  @moduledoc """
  Consumers have a lot of logic around calculating and distributing RDY. This
  is where that goes!
  """

  alias NSQ.ConnInfo
  alias NSQ.Consumer, as: C
  alias NSQ.Consumer.Connections
  import NSQ.Consumer.Helpers

  @doc """
  Initialized from NSQ.Consumer.Supervisor, sends the consumer a message on a
  fixed interval.
  """
  @spec redistribute_loop(pid) :: any
  def redistribute_loop(cons) do
    cons_state = C.get_state(cons)
    GenServer.call(cons, :redistribute_rdy)
    delay = cons_state.config.rdy_redistribute_interval
    :timer.sleep(delay)
    redistribute_loop(cons)
  end

  @doc """
  If we're not in backoff mode and we've hit a "trigger point" to update RDY,
  then go ahead and update RDY. Not for external use.
  """
  @spec maybe_update(pid, C.connection(), C.state()) :: {:ok, C.state()}
  def maybe_update(cons, conn, cons_state) do
    if cons_state.backoff_counter > 0 || cons_state.backoff_duration > 0 do
      # In backoff mode, we only let `start_stop_continue_backoff/3` handle
      # this case.
      NSQ.Logger.debug("""
        (#{inspect(conn)}) skip sending RDY in_backoff:#{cons_state.backoff_counter} || in_backoff_timeout:#{cons_state.backoff_duration}
      """)

      {:ok, cons_state}
    else
      [remain, last_rdy] =
        ConnInfo.fetch(
          cons_state,
          ConnInfo.conn_id(conn),
          [:rdy_count, :last_rdy]
        )

      desired_rdy = per_conn_max_in_flight(cons_state)

      if remain <= 1 || remain < last_rdy / 4 || (desired_rdy > 0 && desired_rdy < remain) do
        NSQ.Logger.debug("""
          (#{inspect(conn)}) sending RDY #{desired_rdy} \
          (#{remain} remain from last RDY #{last_rdy})
        """)

        {:ok, _cons_state} = update(cons, conn, desired_rdy, cons_state)
      else
        NSQ.Logger.debug("""
          (#{inspect(conn)}) skip sending RDY #{desired_rdy} \
          (#{remain} remain out of last RDY #{last_rdy})
        """)

        {:ok, cons_state}
      end
    end
  end

  def maybe_update!(cons, conn, cons_state) do
    {:ok, cons_state} = maybe_update(cons, conn, cons_state)
    cons_state
  end

  @doc """
  Try to update RDY for a given connection, taking configuration and the
  current state into account. Not for external use.
  """
  @spec update(pid, C.connection(), integer, C.state()) :: {:ok, C.state()}
  def update(cons, conn, new_rdy, cons_state) do
    conn_info = ConnInfo.fetch(cons_state, ConnInfo.conn_id(conn))

    cancel_outstanding_retry(cons_state, conn)

    # Cap the given RDY based on the connection config.
    new_rdy = [new_rdy, conn_info.max_rdy] |> Enum.min() |> round

    # Cap the given RDY based on how much we can actually assign. Unless it's
    # 0, in which case we'll be retrying.
    max_possible_rdy = calc_max_possible(cons_state, conn_info)

    new_rdy =
      if max_possible_rdy > 0 do
        [new_rdy, max_possible_rdy] |> Enum.min() |> round
      else
        new_rdy
      end

    {:ok, cons_state} =
      if max_possible_rdy <= 0 && new_rdy > 0 do
        if conn_info.rdy_count == 0 do
          # Schedule RDY.update(consumer, conn, new_rdy) for this connection
          # again in 5 seconds. This is to prevent eternal starvation.
          retry(cons, conn, new_rdy, cons_state)
        else
          {:ok, cons_state}
        end
      else
        transmit(conn, new_rdy, cons_state)
      end

    {:ok, cons_state}
  end

  def update!(cons, conn, new_rdy, cons_state) do
    {:ok, cons_state} = update(cons, conn, new_rdy, cons_state)
    cons_state
  end

  @doc """
  Delay for a configured interval, then call RDY.update.
  """
  @spec retry(pid, C.connection(), integer, C.state()) :: {:ok, C.state()}
  def retry(cons, conn, count, cons_state) do
    delay = cons_state.config.rdy_retry_delay
    NSQ.Logger.debug("(#{inspect(conn)}) retry RDY in #{delay / 1000} seconds")

    {:ok, retry_pid} =
      Task.start_link(fn ->
        :timer.sleep(delay)
        GenServer.call(cons, {:update_rdy, conn, count})
      end)

    ConnInfo.update(cons_state, ConnInfo.conn_id(conn), %{retry_rdy_pid: retry_pid})

    {:ok, cons_state}
  end

  @doc """
  Send a RDY command for the given connection.
  """
  @spec transmit(C.connection(), integer, C.state()) :: {:ok, C.state()}
  def transmit({_id, pid} = conn, count, cons_state) do
    [last_rdy] = ConnInfo.fetch(cons_state, ConnInfo.conn_id(conn), [:last_rdy])

    if count == 0 && last_rdy == 0 do
      {:ok, cons_state}
    else
      # We intentionally don't match this GenServer.call. If the socket isn't
      # set up or is erroring out, we don't want to propagate that connection
      # error to the consumer.
      NSQ.Connection.cmd_noresponse(pid, {:rdy, count})
      {:ok, cons_state}
    end
  end

  @doc """
  This will only be triggered in odd cases where we're in backoff or when there
  are more connections than max in flight. It will randomly change RDY on
  some connections to 0 and 1 so that they're all guaranteed to eventually
  process messages.
  """
  @spec redistribute(pid, C.state()) :: {:ok, C.state()}
  def redistribute(cons, cons_state) do
    if should_redistribute?(cons_state) do
      conns = Connections.get(cons_state)
      conn_count = length(conns)

      if conn_count > cons_state.max_in_flight do
        NSQ.Logger.debug("""
          redistributing RDY state
          (#{conn_count} conns > #{cons_state.max_in_flight} max_in_flight)
        """)
      end

      if cons_state.backoff_counter > 0 && conn_count > 1 do
        NSQ.Logger.debug("""
          redistributing RDY state (in backoff and #{conn_count} conns > 1)
        """)
      end

      # Free up any connections that are RDY but not processing messages.
      Connections.idle_with_rdy(cons_state)
      |> Enum.map(fn conn ->
        NSQ.Logger.debug("(#{inspect(conn)}) idle connection, giving up RDY")
        {:ok, _cons_state} = update(cons, conn, 0, cons_state)
      end)

      # Determine how much RDY we can distribute. This needs to happen before
      # we give up RDY, or max_in_flight will end up equalling RDY.
      available_max_in_flight = get_available_max_in_flight(cons_state)

      # Distribute it!
      {sorted_conns, cons_state} = sort_conns_for_round_robin(conns, cons_state)
      distribute(cons, sorted_conns, available_max_in_flight, cons_state)
    else
      # Nothing to do. This is the usual path!
      {:ok, cons_state}
    end
  end

  def redistribute!(cons, cons_state) do
    {:ok, cons_state} = redistribute(cons, cons_state)
    cons_state
  end

  # Helper for redistribute; we set RDY to 1 for _some_ connections that
  # were halted, until there's no more RDY left to assign. We assume that the
  # list of connections has already been sorted in the order that we should
  # distribute.
  @spec distribute(pid, [C.connection()], integer, C.state()) ::
          {:ok, C.state()}
  defp distribute(cons, possible_conns, available_max_in_flight, cons_state) do
    if length(possible_conns) == 0 || available_max_in_flight <= 0 do
      {:ok, cons_state}
    else
      [conn | rest] = possible_conns
      NSQ.Logger.debug("(#{inspect(conn)}) redistributing RDY")
      {:ok, cons_state} = update(cons, conn, 1, cons_state)
      distribute(cons, rest, available_max_in_flight - 1, cons_state)
    end
  end

  @spec sort_conns_for_round_robin([C.connection()], C.state()) :: {[C.connection()], C.state()}
  defp sort_conns_for_round_robin(conns, cons_state) do
    # We sort to ensure consistency of start_index across runs.
    sorted_conns = conns |> Enum.sort_by(fn {conn_id, _} -> conn_id end)
    start_index = rem(cons_state.distribution_counter, length(sorted_conns))

    # We want to start distributing from a specific index. This reorders the
    # list of connections so that the target index becomes first, followed by
    # elements in their natural sorted order. Once we hit the end of the list,
    # it wraps back to the beginning, still in order.
    sorted_conns =
      sorted_conns
      |> Enum.split(start_index)
      |> Tuple.to_list()
      |> Enum.reverse()
      |> Enum.concat()

    # By increasing distribution count, we ensure that the start_index will be
    # different next time this runs. If the number of connections does not
    # change, we will shift the distribution by 1 with each run.
    {
      sorted_conns,
      %{cons_state | distribution_counter: cons_state.distribution_counter + 1}
    }
  end

  @spec cancel_outstanding_retry(C.state(), C.connection()) :: any
  defp cancel_outstanding_retry(cons_state, conn) do
    conn_info = ConnInfo.fetch(cons_state, ConnInfo.conn_id(conn))

    # If this is for a connection that's retrying, kill the timer and clean up.
    if retry_pid = conn_info.retry_rdy_pid do
      if Process.alive?(retry_pid) do
        NSQ.Logger.debug(
          "(#{inspect(conn)}) rdy retry pid #{inspect(retry_pid)} detected, killing"
        )

        Process.exit(retry_pid, :normal)
      end

      ConnInfo.update(cons_state, ConnInfo.conn_id(conn), %{retry_rdy_pid: nil})
    end
  end

  @spec calc_max_possible(C.state(), map) :: integer
  defp calc_max_possible(cons_state, conn_info) do
    rdy_count = conn_info.rdy_count
    max_in_flight = cons_state.max_in_flight
    total_rdy = total_rdy_count(cons_state)
    max_in_flight - total_rdy + rdy_count
  end

  @spec should_redistribute?(C.state()) :: boolean
  defp should_redistribute?(cons_state) do
    conn_count = Connections.count(cons_state)
    in_backoff = cons_state.backoff_counter > 0
    in_backoff_timeout = cons_state.backoff_duration > 0

    !in_backoff_timeout &&
      conn_count > 0 &&
      (conn_count > cons_state.max_in_flight ||
         (in_backoff && conn_count > 1) ||
         cons_state.need_rdy_redistributed)
  end

  # Cap available max in flight based on current RDY/backoff status.
  defp get_available_max_in_flight(cons_state) do
    total_rdy = total_rdy_count(cons_state)

    if cons_state.backoff_counter > 0 do
      # In backoff mode, we only ever want RDY=1 for the whole consumer. This
      # makes sure that available is only 1 if total_rdy is 0.
      1 - total_rdy
    else
      cons_state.max_in_flight - total_rdy
    end
  end
end
