defmodule NSQ.Connection.Initializer do
  alias NSQ.Connection, as: C
  alias NSQ.Connection.MessageHandling
  alias NSQ.Connection.Buffer
  alias NSQ.ConnInfo
  import NSQ.Protocol

  @socket_opts [as: :binary, mode: :passive, packet: :raw]

  @project ElixirNsq.Mixfile.project()
  @user_agent "#{@project[:app]}/#{@project[:version]}"
  @ssl_versions [:sslv3, :tlsv1, :"tlsv1.1", :"tlsv1.2"] |> Enum.with_index()

  @spec connect(%{nsqd: C.host_with_port()}) :: {:ok, C.state()} | {:error, String.t()}
  def connect(%{nsqd: {host, port}} = state) do
    if should_connect?(state) do
      socket_opts =
        @socket_opts
        |> Keyword.merge(
          send: [{:timeout, state.config.write_timeout}],
          timeout: state.config.dial_timeout
        )

      case Socket.TCP.connect(host, port, socket_opts) do
        {:ok, socket} ->
          state.reader |> Buffer.setup_socket(socket, state.config.read_timeout)
          state.writer |> Buffer.setup_socket(socket, state.config.read_timeout)

          state =
            %{state | socket: socket}
            |> do_handshake!
            |> start_receiving_messages!
            |> reset_connects

          {:ok, %{state | connected: true}}

        {:error, reason} ->
          if length(state.config.nsqlookupds) > 0 do
            NSQ.Logger.warn(
              "(#{inspect(self())}) connect failed; #{reason}; discovery loop should respawn"
            )

            {{:error, reason}, %{state | connect_attempts: state.connect_attempts + 1}}
          else
            if state.config.max_reconnect_attempts > 0 do
              NSQ.Logger.warn(
                "(#{inspect(self())}) connect failed; #{reason}; discovery loop should respawn"
              )

              {{:error, reason}, %{state | connect_attempts: state.connect_attempts + 1}}
            else
              NSQ.Logger.error(
                "(#{inspect(self())}) connect failed; #{reason}; reconnect turned off; terminating connection"
              )

              Process.exit(self(), :connect_failed)
            end
          end
      end
    else
      NSQ.Logger.error("#{inspect(self())}: Failed to connect; terminating connection")
      Process.exit(self(), :connect_failed)
    end
  end

  @doc """
  Immediately after connecting to the NSQ socket, both consumers and producers
  follow this protocol.
  """
  @spec do_handshake(C.state()) :: {:ok, C.state()}
  def do_handshake(conn_state) do
    conn_state |> send_magic_v2

    {:ok, conn_state} = identify(conn_state)

    # Producers don't have a channel, so they won't do this.
    if conn_state.channel do
      subscribe(conn_state)
    end

    {:ok, conn_state}
  end

  def do_handshake!(conn_state) do
    {:ok, conn_state} = do_handshake(conn_state)
    conn_state
  end

  @spec send_magic_v2(C.state()) :: :ok
  defp send_magic_v2(conn_state) do
    NSQ.Logger.debug("(#{inspect(self())}) sending magic v2...")
    conn_state |> Buffer.send!(encode(:magic_v2))
  end

  @spec identify(C.state()) :: {:ok, binary}
  defp identify(conn_state) do
    NSQ.Logger.debug("(#{inspect(self())}) identifying...")
    identify_obj = encode({:identify, identify_props(conn_state)})
    conn_state |> Buffer.send!(identify_obj)
    {:response, json} = recv_nsq_response(conn_state)
    {:ok, _conn_state} = update_from_identify_response(conn_state, json)
  end

  @spec identify_props(C.state()) :: map
  defp identify_props(%{nsqd: {host, port}, config: config} = conn_state) do
    %{
      client_id: "#{host}:#{port} (#{inspect(conn_state.parent)})",
      hostname: to_string(:net_adm.localhost()),
      feature_negotiation: true,
      heartbeat_interval: config.heartbeat_interval,
      output_buffer: config.output_buffer_size,
      output_buffer_timeout: config.output_buffer_timeout,
      tls_v1: config.tls_v1,
      snappy: false,
      deflate: config.deflate,
      deflate_level: config.deflate_level,
      sample_rate: 0,
      user_agent: config.user_agent || @user_agent,
      msg_timeout: config.msg_timeout
    }
  end

  def inflate(data) do
    z = :zlib.open()
    :ok = z |> :zlib.inflateInit(-15)
    inflated = z |> :zlib.inflateChunk(data)
    NSQ.Logger.warn("inflated chunk?")
    NSQ.Logger.warn(inspect(inflated))
    :ok = z |> :zlib.inflateEnd()
    :ok = z |> :zlib.close()
    inflated
  end

  @spec update_from_identify_response(C.state(), binary) :: {:ok, C.state()}
  defp update_from_identify_response(conn_state, json) do
    {:ok, parsed} = Jason.decode(json)

    # respect negotiated max_rdy_count
    if parsed["max_rdy_count"] do
      ConnInfo.update(conn_state, %{max_rdy: parsed["max_rdy_count"]})
    end

    # respect negotiated msg_timeout
    timeout = parsed["msg_timeout"] || conn_state.config.msg_timeout
    conn_state = %{conn_state | msg_timeout: timeout}

    # wrap our socket with SSL if TLS is enabled
    conn_state =
      if parsed["tls_v1"] == true do
        NSQ.Logger.debug("Upgrading to TLS...")

        socket =
          Socket.SSL.connect!(conn_state.socket,
            cacertfile: conn_state.config.tls_cert,
            keyfile: conn_state.config.tls_key,
            versions: ssl_versions(conn_state.config.tls_min_version),
            verify: ssl_verify_atom(conn_state.config)
          )

        conn_state = %{conn_state | socket: socket}
        conn_state.reader |> Buffer.setup_socket(socket, conn_state.config.read_timeout)
        conn_state.writer |> Buffer.setup_socket(socket, conn_state.config.read_timeout)
        conn_state |> wait_for_ok!
        conn_state
      else
        conn_state
      end

    # If compression is enabled, we expect to receive a compressed "OK"
    # immediately.
    conn_state.reader |> Buffer.setup_compression(parsed, conn_state.config)
    conn_state.writer |> Buffer.setup_compression(parsed, conn_state.config)

    if parsed["deflate"] == true || parsed["snappy"] == true do
      conn_state |> wait_for_ok!
    end

    if parsed["auth_required"] == true do
      NSQ.Logger.debug("sending AUTH")
      auth_cmd = encode({:auth, conn_state.config.auth_secret})
      conn_state |> Buffer.send!(auth_cmd)
      {:response, json} = recv_nsq_response(conn_state)
      NSQ.Logger.debug(json)
    end

    {:ok, conn_state}
  end

  defp ssl_verify_atom(config) do
    if config.tls_insecure_skip_verify == true do
      :verify_none
    else
      :verify_peer
    end
  end

  @spec subscribe(C.state()) :: {:ok, binary}
  defp subscribe(%{topic: topic, channel: channel} = conn_state) do
    NSQ.Logger.debug("(#{inspect(self())}) subscribe to #{topic} #{channel}")
    conn_state |> Buffer.send!(encode({:sub, topic, channel}))

    NSQ.Logger.debug("(#{inspect(self())}) wait for subscription acknowledgment")
    conn_state |> wait_for_ok!
  end

  @spec recv_nsq_response(C.state()) :: {:response, binary}
  defp recv_nsq_response(conn_state) do
    <<msg_size::size(32)>> = conn_state |> Buffer.recv!(4)
    raw_msg_data = conn_state |> Buffer.recv!(msg_size)
    {:response, _response} = decode(raw_msg_data)
  end

  defp wait_for_ok!(state) do
    expected = ok_msg()
    ^expected = state |> Buffer.recv!(byte_size(expected))
  end

  @spec ssl_versions(NSQ.Config.t()) :: [atom]
  def ssl_versions(tls_min_version) do
    if tls_min_version do
      min_index = @ssl_versions[tls_min_version]

      @ssl_versions
      |> Enum.drop_while(fn {_, index} -> index < min_index end)
      |> Enum.map(fn {version, _} -> version end)
      |> Enum.reverse()
    else
      @ssl_versions
      |> Enum.map(fn {version, _} -> version end)
      |> Enum.reverse()
    end
  end

  @spec should_connect?(C.state()) :: boolean
  defp should_connect?(state) do
    state.connect_attempts == 0 ||
      state.connect_attempts <= state.config.max_reconnect_attempts
  end

  @spec start_receiving_messages(C.state()) :: {:ok, C.state()}
  defp start_receiving_messages(state) do
    reader_pid = spawn_link(MessageHandling, :recv_nsq_messages, [state, self()])
    state = %{state | reader_pid: reader_pid}
    GenServer.cast(self(), :flush_cmd_queue)
    {:ok, state}
  end

  defp start_receiving_messages!(state) do
    {:ok, state} = start_receiving_messages(state)
    state
  end

  @spec reset_connects(C.state()) :: C.state()
  defp reset_connects(state), do: %{state | connect_attempts: 0}
end