defmodule Load do

  require Logger

  def scale(count, address \\ :all) when is_integer(count) and count > 0 and (address == :all or is_binary(address)) do
    DynamicSupervisor.which_children(Load.Connection.Supervisor)
    |> Enum.each(fn {:undefined, pid, :worker, [Load.WSClient]} ->
      GenServer.cast(pid, {:ws_send, address, %{command: "scale", count: count}})
      end)
  end

  def connect(addresses \\ ["localhost"]) when is_list(addresses) do
    DynamicSupervisor.which_children(Load.Connection.Supervisor)
    |> Enum.reduce(MapSet.new(addresses), fn {:undefined, pid, :worker, [Load.WSClient]}, acc ->
      address = GenServer.call(pid, :get_address)
      if MapSet.member?(acc, address) do
        MapSet.delete(acc, address)
      else
        GenServer.cast(pid, {:ws_send, address, %{command: "terminate"}})
        acc
      end
    end)
    |> Enum.each(fn address ->
      DynamicSupervisor.start_child(Load.Connection.Supervisor, {Load.WSClient, address: address})
    end)
  end

  def q, do: Stats.get()

  def h, do:
    IO.puts(
      "+-----------------------------------------------\n"
    <>"| Commands available:\n"
    <>"| Load.scale  (count, :all | address) - scale to count workers on selected nodes\n"
    <>"| Load.connect(addresses)             - connect to addresses\n"
    <>"| Load:q() - print current stats\n"
    <>"+-----------------------------------------------"
    )

end