defmodule NSQ.Consumer.Connections do
  @moduledoc """
  Functions for connecting, disconnecting, managing connections, etc.
  """

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

  @doc """
  Initialized from NSQ.Consumer.Supervisor, sends the consumer a message on a
  fixed interval.
  """
  @spec discovery_loop(pid) :: any
  def discovery_loop(cons) do
    cons_state = C.get_state(cons)

    %NSQ.Config{
      lookupd_poll_interval: poll_interval,
      lookupd_poll_jitter: poll_jitter
    } = cons_state.config

    delay = poll_interval + round(poll_interval * poll_jitter * :rand.uniform())
    :timer.sleep(delay)

    GenServer.call(cons, :discover_nsqds)
    discovery_loop(cons)
  end

  def close(cons_state) do
    NSQ.Logger.info("Closing connections for consumer #{inspect(self())}")
    connections = get(cons_state)

    Enum.each(connections, fn {_, conn_pid} ->
      Task.start_link(NSQ.Connection, :close, [conn_pid])
    end)

    {:ok, %{cons_state | stop_flag: true}}
  end

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

  def refresh(cons_state) do
    {:ok, cons_state} = delete_dead(cons_state)
    {:ok, cons_state} = reconnect_failed(cons_state)
    {:ok, cons_state} = discover_nsqds_and_connect(self(), cons_state)
    {:ok, cons_state}
  end

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

  @doc """
  Finds and updates list of live NSQDs using either NSQ.Config.nsqlookupds or
  NSQ.Config.nsqds, depending on what's configured. Preference is given to
  nsqlookupd. Not for external use.
  """
  @spec discover_nsqds_and_connect(pid, C.state()) :: {:ok, C.state()}
  def discover_nsqds_and_connect(cons, cons_state) do
    nsqds =
      cond do
        length(cons_state.config.nsqlookupds) > 0 ->
          NSQ.Logger.debug(
            "(#{inspect(self())}) Discovering nsqds via nsqlookupds #{inspect(cons_state.config.nsqlookupds)}"
          )

          cons_state.config.nsqlookupds
          |> NSQ.Lookupd.nsqds_with_topic(cons_state.topic)

        length(cons_state.config.nsqds) > 0 ->
          NSQ.Logger.debug(
            "(#{inspect(self())}) Using configured nsqds #{inspect(cons_state.config.nsqds)}"
          )

          cons_state.config.nsqds

        true ->
          raise "No nsqds or nsqlookupds are configured"
      end

    {:ok, _cons_state} = update(nsqds, cons, cons_state)
  end

  @doc """
  Any inactive connections will be killed and any newly discovered connections
  will be added. Existing connections with no change are left alone. Not for
  external use.
  """
  @spec update([C.host_with_port()], pid, C.state()) :: {:ok, C.state()}
  def update(discovered_nsqds, cons, cons_state) do
    dead_conns = dead_connections(discovered_nsqds, cons, cons_state)
    {:ok, cons_state} = stop_connections(dead_conns, cons, cons_state)

    nsqds_to_connect = new_nsqds(discovered_nsqds, cons, cons_state)
    {:ok, cons_state} = connect_to_nsqds(nsqds_to_connect, cons, cons_state)

    {:ok, cons_state}
  end

  @doc """
  Given a list of NSQD hosts, open a connection for each.
  """
  @spec connect_to_nsqds([C.host_with_port()], pid, C.state()) :: {:ok, C.state()}
  def connect_to_nsqds(nsqds, cons, cons_state \\ nil) do
    if length(nsqds) > 0 do
      NSQ.Logger.info("Connecting to nsqds #{inspect(nsqds)}")
    end

    cons_state =
      Enum.reduce(nsqds, cons_state, fn nsqd, last_state ->
        {:ok, new_state} = connect_to_nsqd(nsqd, cons, last_state)
        new_state
      end)

    {:ok, cons_state}
  end

  @doc """
  Create a connection to NSQD and add it to the consumer's supervised list.
  Not for external use.
  """
  @spec connect_to_nsqd(C.host_with_port(), pid, C.state()) :: {:ok, C.state()}
  def connect_to_nsqd(nsqd, cons, cons_state) do
    Process.flag(:trap_exit, true)

    try do
      {:ok, _pid} =
        NSQ.Connection.Supervisor.start_child(
          cons,
          nsqd,
          cons_state
        )

      # We normally set RDY to 1, but if we're spawning more connections than
      # max_in_flight, we don't want to break our contract. In that case, the
      # `RDY.redistribute` loop will take care of getting this connection some
      # messages later.
      remaining_rdy = cons_state.max_in_flight - total_rdy_count(cons_state)

      cons_state =
        if remaining_rdy > 0 do
          conn = conn_from_nsqd(cons, nsqd, cons_state)
          {:ok, cons_state} = RDY.transmit(conn, 1, cons_state)
          cons_state
        else
          cons_state
        end

      {:ok, cons_state}
    catch
      :error, _ ->
        NSQ.Logger.error("#{inspect(cons)}: Error connecting to #{inspect(nsqd)}")
        conn_id = ConnInfo.conn_id(cons, nsqd)
        ConnInfo.delete(cons_state, conn_id)
        {:ok, cons_state}
    after
      Process.flag(:trap_exit, false)
    end
  end

  @doc """
  Given a list of connections, force them to stop. Return the new state without
  those connections.
  """
  @spec stop_connections([C.connection()], pid, C.state()) :: {:ok, C.state()}
  def stop_connections(dead_conns, cons, cons_state) do
    if length(dead_conns) > 0 do
      NSQ.Logger.info("Stopping connections #{inspect(dead_conns)}")
    end

    cons_state =
      Enum.reduce(dead_conns, cons_state, fn {nsqd, _pid}, last_state ->
        {:ok, new_state} = stop_connection(cons, nsqd, last_state)
        new_state
      end)

    {:ok, cons_state}
  end

  @doc """
  Given a single connection, immediately terminate its process (and all
  descendant processes, such as message handlers) and remove its info from the
  ConnInfo agent. Not for external use.
  """
  @spec stop_connection(pid, C.host_with_port(), C.state()) :: {:ok, C.state()}
  def stop_connection(cons, conn_id, cons_state) do
    # Terminate the connection for real.
    # TODO: Change this method to `kill_connection` and make `stop_connection`
    # graceful.
    Supervisor.terminate_child(cons_state.conn_sup_pid, conn_id)
    {:ok, cons_state} = cleanup_connection(cons, conn_id, cons_state)

    {:ok, cons_state}
  end

  @doc """
  When a connection is terminated or dies, we must do some extra cleanup.
  First, a terminated process isn't necessarily removed from the supervisor's
  list; therefore we call `Supervisor.delete_child/2`. And info about this
  connection like RDY must be removed so it doesn't contribute to `total_rdy`.
  Not for external use.
  """
  @spec cleanup_connection(pid, C.host_with_port(), C.state()) :: {:ok, C.state()}
  def cleanup_connection(_cons, conn_id, cons_state) do
    # If a connection is terminated normally or non-normally, it will still be
    # listed in the supervision tree. Let's remove it when we clean up.
    Supervisor.delete_child(cons_state.conn_sup_pid, conn_id)

    # Delete the connection info from the shared map so we don't use it to
    # perform calculations.
    ConnInfo.delete(cons_state, conn_id)

    {:ok, cons_state}
  end

  @doc """
  We may have open connections which nsqlookupd stops reporting. This function
  tells us which connections we have stored in state but not in nsqlookupd.
  Not for external use.
  """
  @spec dead_connections([C.host_with_port()], pid, C.state()) :: [C.connection()]
  def dead_connections(discovered_nsqds, cons, cons_state) do
    Enum.reject(get(cons_state), fn conn ->
      conn_already_discovered?(cons, conn, discovered_nsqds)
    end)
  end

  @doc """
  When nsqlookupd reports available producers, there are some that may not
  already be in our connection list. This function reports which ones are new
  so we can connect to them.
  """
  @spec new_nsqds([C.host_with_port()], pid, C.state()) :: [C.host_with_port()]
  def new_nsqds(discovered_nsqds, cons, cons_state) do
    Enum.reject(discovered_nsqds, fn nsqd ->
      nsqd_already_has_connection?(nsqd, cons, cons_state)
    end)
  end

  @doc """
  Frequently, when testing, we publish a message then immediately want a
  consumer to process it, but this doesn't work if the consumer doesn't
  discover the nsqd first. Only meant for testing.
  """
  @spec discover_nsqds(pid) :: :ok
  def discover_nsqds(sup_pid) do
    cons = C.get(sup_pid)
    GenServer.call(cons, :discover_nsqds)
    :ok
  end

  @doc """
  Iterate over all listed connections and delete the ones that are dead. This
  exists because it is difficult to reliably clean up a connection immediately
  after it is terminated (it might still be running). This function runs in the
  discovery loop to provide consistency.
  """
  @spec delete_dead(C.state()) :: {:ok, C.state()}
  def delete_dead(state) do
    Enum.each(get(state), fn {conn_id, pid} ->
      unless Process.alive?(pid) do
        Supervisor.delete_child(state.conn_sup_pid, conn_id)
      end
    end)

    {:ok, state}
  end

  def delete_dead!(state) do
    {:ok, state} = delete_dead(state)
    state
  end

  def reconnect_failed(state) do
    Enum.each(get(state), fn {_, pid} ->
      if Process.alive?(pid), do: GenServer.cast(pid, :reconnect)
    end)

    {:ok, state}
  end

  @spec count(C.state()) :: integer
  def count(cons_state) do
    %{active: active} = Supervisor.count_children(cons_state.conn_sup_pid)

    if is_integer(active) do
      active
    else
      NSQ.Logger.warn(
        "(#{inspect(self())}) non-integer #{inspect(active)} returned counting connections, returning 0 instead"
      )

      0
    end
  end

  @doc """
  Returns all live connections for a consumer. This function, which takes
  a consumer's entire state as an argument, is for convenience. Not for
  external use.
  """
  @spec get(C.state()) :: [C.connection()]
  def get(%{conn_sup_pid: conn_sup_pid}) do
    children = Supervisor.which_children(conn_sup_pid)
    Enum.map(children, fn {child_id, pid, _, _} -> {child_id, pid} end)
  end

  @doc """
  Returns all live connections for a consumer. Used in tests. Not for external
  use.
  """
  @spec get(pid, C.state()) :: [C.connection()]
  def get(cons, cons_state \\ nil) when is_pid(cons) do
    cons_state = cons_state || C.get_state(cons)
    children = Supervisor.which_children(cons_state.conn_sup_pid)
    Enum.map(children, fn {child_id, pid, _, _} -> {child_id, pid} end)
  end

  @spec idle_with_rdy(C.state()) :: [C.connection()]
  def idle_with_rdy(cons_state) do
    conns = get(cons_state)

    Enum.filter(conns, fn conn ->
      conn_id = ConnInfo.conn_id(conn)

      [last_msg_t, rdy_count] =
        ConnInfo.fetch(
          cons_state,
          conn_id,
          [:last_msg_timestamp, :rdy_count]
        )

      sec_since_last_msg = now() - last_msg_t
      ms_since_last_msg = sec_since_last_msg * 1000

      NSQ.Logger.debug(
        "(#{inspect(conn)}) rdy: #{rdy_count} (last message received #{sec_since_last_msg} seconds ago)"
      )

      ms_since_last_msg > cons_state.config.low_rdy_idle_timeout && rdy_count > 0
    end)
  end

  @spec conn_already_discovered?(pid, C.connection(), [C.host_with_port()]) :: boolean
  defp conn_already_discovered?(cons, {conn_id, _}, discovered_nsqds) do
    Enum.any?(discovered_nsqds, fn nsqd ->
      ConnInfo.conn_id(cons, nsqd) == conn_id
    end)
  end

  @spec nsqd_already_has_connection?(C.host_with_port(), pid, C.state()) :: boolean
  defp nsqd_already_has_connection?(nsqd, cons, cons_state) do
    needle = ConnInfo.conn_id(cons, nsqd)

    Enum.any?(get(cons_state), fn {conn_id, _} ->
      conn_id == needle
    end)
  end
end
