soulless

defmodule Soulless.Websocket.Client do
  @behaviour :gen_statem
  require Logger

  @default_timeout :timer.seconds(30)
  @default_heartbeat_interval :timer.seconds(30)
  @default_reconnect_interval :timer.minutes(5)

  defstruct [
    # transport layer
    :conn,
    :stream_ref,
    # user-provided callback handler
    :handler,
    # handler and its private mutable state
    :implementation,
    :implementation_state,
    :init_implementation_state,
    # task supervisor for notice dispatch
    :task_sup,
    # sequence counter for marking requests
    request_id: 0,
    # %{id => gen_statem from()} - callers awaiting a matching response
    pending: %{},
    connection_attempt: 0,
    heartbeat_interval: @default_heartbeat_interval,
    reconnect_interval: @default_reconnect_interval
  ]

  # Public API

  @spec start(keyword()) :: :gen_statem.start_ret()
  def start(opts) do
    {name, init_opts} = Keyword.pop(opts, :name)

    case name do
      nil -> :gen_statem.start(__MODULE__, init_opts, [])
      name -> :gen_statem.start(name, __MODULE__, init_opts, [])
    end
  end

  @spec start_link(keyword()) :: :gen_statem.start_ret()
  def start_link(opts) do
    {name, init_opts} = Keyword.pop(opts, :name)

    case name do
      nil -> :gen_statem.start_link(__MODULE__, init_opts, [])
      name -> :gen_statem.start_link(name, __MODULE__, init_opts, [])
    end
  end

  def child_spec(args) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [args]},
      restart: :permanent,
      shutdown: 5000,
      type: :worker
    }
  end

  @doc """
  Send `payload` and block until the matching response arrives.
  Returns `{:ok, response}` or `{:error, reason}`.
  """
  @spec fetch(:gen_statem.server_ref(), term(), timeout() | nil) ::
          {:ok, term()} | {:error, term()}
  def fetch(client, payload, timeout \\ nil) do
    timeout = timeout || @default_timeout
    :gen_statem.call(client, {:fetch, payload}, timeout)
  end

  @doc """
  Retrieve value from the handler_state by path.
  """
  @spec fetch_state(:gen_statem.server_ref(), [atom()]) :: term() | nil
  def fetch_state(client, path) do
    :gen_statem.call(client, {:fetch_state, path}, @default_timeout)
  end

  # gen_statem callbacks

  @impl :gen_statem
  def callback_mode, do: :handle_event_function

  @impl :gen_statem
  def init(opts) do
    heartbeat_interval = Keyword.get(opts, :heartbeat_interval, @default_heartbeat_interval)
    reconnect_interval = Keyword.get(opts, :reconnect_interval, @default_reconnect_interval)
    handler = Keyword.get(opts, :handler, Soulless.Handler.Default)
    implementation = Keyword.fetch!(opts, :implementation)
    implementation_state = implementation.init(opts)

    {:ok, task_sup} = Task.Supervisor.start_link()

    data = %__MODULE__{
      handler: handler,
      implementation: implementation,
      implementation_state: implementation_state,
      init_implementation_state: implementation_state,
      heartbeat_interval: heartbeat_interval,
      reconnect_interval: reconnect_interval,
      task_sup: task_sup
    }

    {:ok, :disconnected, data, [initial_event(data)]}
  end

  @impl :gen_statem
  def handle_event({:timeout, :resolve_endpoint}, :resolve_endpoint, :disconnected, _data) do
    {:keep_state_and_data, [{:next_event, :internal, :resolve_endpoint}]}
  end

  def handle_event(:internal, :resolve_endpoint, :disconnected, data) do
    Logger.debug("[#{__MODULE__}] Resolving endpoint")

    case data.implementation.resolve_endpoint(data.implementation_state) do
      {:ok, new_handler_state} ->
        data = %{data | implementation_state: new_handler_state}
        {:next_state, :disconnected, data, [{:next_event, :internal, :connect}]}

      {:error, :maintenance} ->
        Logger.warning(
          "[#{__MODULE__}] Server is under maintenance. Retrying in #{data.reconnect_interval / 1000}s"
        )

        {:keep_state, next_connection_attempt(data),
         [{{:timeout, :resolve_endpoint}, data.reconnect_interval, :resolve_endpoint}]}

      {:error, error} ->
        Logger.error("[#{__MODULE__}] Connection failed: #{inspect(error)}")

        {:keep_state, next_connection_attempt(data),
         [{{:timeout, :resolve_endpoint}, backoff(data.connection_attempt), :resolve_endpoint}]}
    end
  end

  def handle_event({:timeout, :connect}, :connect, :disconnected, data) do
    {:keep_state_and_data, [initial_event(data)]}
  end

  def handle_event(:internal, :connect, :disconnected, data) do
    endpoint_uri = data.implementation.endpoint_uri(data.implementation_state)
    {host, _path, port, opts} = Soulless.HTTP.gun_args(endpoint_uri)

    Logger.debug("[#{__MODULE__}] Connecting to #{endpoint_uri}")

    case :gun.open(host, port, opts) do
      {:ok, conn} ->
        {:next_state, :connecting, %{data | conn: conn}}

      {:error, reason} ->
        Logger.error(
          "[#{__MODULE__}] Connection failed: #{inspect(reason)}. Retrying in: #{data.reconnect_interval / 1000}s"
        )

        {:keep_state, next_connection_attempt(data),
         [{{:timeout, :connect}, backoff(data.connection_attempt), :connect}]}
    end
  end

  # send ping frames periodically to maintain the connection
  def handle_event({:timeout, :heartbeat}, :heartbeat, state, data)
      when state in [:connected, :authenticated] do
    :gun.ws_send(data.conn, data.stream_ref, :ping)
    {:keep_state_and_data, [{{:timeout, :heartbeat}, data.heartbeat_interval, :heartbeat}]}
  end

  def handle_event({:timeout, :heartbeat}, :heartbeat, _state, _data) do
    :keep_state_and_data
  end

  def handle_event(:info, {:gun_up, conn, _protocol}, :connecting, %{conn: conn} = data) do
    Logger.debug("[#{__MODULE__}] TCP up; upgrading to WebSocket")
    endpoint_uri = data.implementation.endpoint_uri(data.implementation_state)
    {_body, path, _port, _opts} = Soulless.HTTP.gun_args(endpoint_uri)
    stream_ref = :gun.ws_upgrade(conn, path)
    {:next_state, :upgrading, %{data | stream_ref: stream_ref}}
  end

  def handle_event(
        :info,
        {:gun_upgrade, conn, ref, ["websocket"], _headers},
        :upgrading,
        %{conn: conn, stream_ref: ref} = data
      ) do
    Logger.info("[#{__MODULE__}] WebSocket ready")
    on_ready(data)
  end

  def handle_event(
        :info,
        {:gun_upgrade, conn, ref, _other, headers},
        :upgrading,
        %{conn: conn, stream_ref: ref} = _data
      ) do
    Logger.error("[#{__MODULE__}] WebSocket upgrade failed: #{inspect(headers)}")
    {:stop, :upgrade_failed}
  end

  def handle_event(
        :info,
        {:gun_response, conn, ref, _, status, headers},
        :upgrading,
        %{conn: conn, stream_ref: ref} = _data
      ) do
    Logger.error("[#{__MODULE__}] Unexpected HTTP #{status}: #{inspect(headers)}")
    {:stop, {:unexpected_http_response, status}}
  end

  # respond to ping frames
  def handle_event(:info, {:gun_ws, conn, ref, :ping}, _state, %{conn: conn, stream_ref: ref}) do
    :gun.ws_send(conn, ref, :ping)
    :keep_state_and_data
  end

  def handle_event(:info, {:gun_ws, conn, ref, {:ping, _}}, _state, %{conn: conn, stream_ref: ref}) do
    :gun.ws_send(conn, ref, :ping)
    :keep_state_and_data
  end

  # response handling
  def handle_event(
        :info,
        {:gun_ws, conn, ref, {frame_type, raw}},
        state,
        %{conn: conn, stream_ref: ref} = data
      )
      when frame_type in [:text, :binary] do
    case data.implementation.handle_decode(frame_type, raw, data.implementation_state) do
      {:ok, {msg, new_handler_state}} ->
        on_message(msg, state, %{data | implementation_state: new_handler_state})

      {:error, reason} ->
        Logger.error("[#{__MODULE__}] Decode error: #{inspect(reason)}")
        :keep_state_and_data
    end
  end

  # server-initiated websocket close
  def handle_event(
        :info,
        {:gun_ws, conn, ref, {:close, code, reason}},
        _state,
        %{conn: conn, stream_ref: ref} = data
      ) do
    on_disconnect({:websocket_closed, code, reason}, data)
  end

  # transport-level connection loss
  def handle_event(
        :info,
        {:gun_down, conn, _protocol, reason, _killed},
        _state,
        %{conn: conn} = data
      ) do
    on_disconnect({:connection_down, reason}, data)
  end

  # gun-level error
  def handle_event(:info, {:gun_error, conn, _ref, reason}, _state, %{conn: conn} = data) do
    on_disconnect({:gun_error, reason}, data)
  end

  # request handling
  def handle_event({:call, from}, {:fetch, payload}, state, data)
      when state in [:connected, :authenticated] do
    data = send_request_frame(payload, from, data)
    {:keep_state, data}
  end

  def handle_event({:call, from}, {:fetch, _payload}, _state, _data) do
    :gen_statem.reply(from, {:error, :not_ready})
    :keep_state_and_data
  end

  def handle_event({:call, from}, {:fetch_state, path}, _state, data) do
    result = data.implementation.fetch_state(data.implementation_state, path)
    :gen_statem.reply(from, result)
    :keep_state_and_data
  end

  # exit states sent by handle_notice
  def handle_event(:info, {:DOWN, _ref, :process, _pid, :normal}, _state, _data) do
    :keep_state_and_data
  end

  def handle_event(:info, {:DOWN, _ref, :process, _pid, reason}, _state, _data) do
    Logger.warning("[#{__MODULE__}] Notice handler task crashed: #{inspect(reason)}")
    :keep_state_and_data
  end

  @impl :gen_statem
  def terminate(reason, _state, data) do
    error = {:error, {:client_terminated, reason}}

    Enum.each(data.pending, fn {_id, from} ->
      :gen_statem.reply(from, error)
    end)

    if data.conn, do: :gun.close(data.conn)
    :ok
  end

  # private state machine helpers

  defp initial_event(data) do
    next_event =
      case data.implementation.endpoint_uri(data.implementation_state) do
        nil -> :resolve_endpoint
        _ -> :connect
      end

    {:next_event, :internal, next_event}
  end

  # called once the websocket upgrade succeeds
  # delegates the choice of the next step to the handler
  defp on_ready(data) do
    case data.implementation.handle_ready(self(), data.implementation_state) do
      {:ok, new_handler_state} ->
        data = %{data | implementation_state: new_handler_state}

        data.handler.handle_ready(self(), :connected)

        {:next_state, :connected, data,
         [{{:timeout, :heartbeat}, data.heartbeat_interval, :heartbeat}]}

      {:authenticate, new_handler_state} ->
        data = %{data | implementation_state: new_handler_state}
        start_auth(data)
    end
  end

  # initiate auth via handler
  defp start_auth(data) do
    case data.implementation.handle_auth_ready(data.implementation_state) do
      {:send, payload, new_handler_state} ->
        Logger.debug("[#{__MODULE__}] Auth started")
        data = send_auth_frame(payload, %{data | implementation_state: new_handler_state})
        {:next_state, :authenticating, data}

      {:error, reason} ->
        Logger.error("[#{__MODULE__}] Auth init failed: #{inspect(reason)}")
        {:stop, {:auth_failed, reason}}
    end
  end

  # messages received while authenticating are assumed to be handle_auth_message's responsibility
  defp on_message(message, :authenticating, data) do
    case data.implementation.handle_auth_message(message, data.implementation_state) do
      {:done, new_handler_state} ->
        Logger.info("[#{__MODULE__}] Auth complete")
        data = %{data | implementation_state: new_handler_state}

        data.handler.handle_ready(self(), :authenticated)

        {:next_state, :authenticated, data,
         [{{:timeout, :heartbeat}, data.heartbeat_interval, :heartbeat}]}

      {:send, payload, new_handler_state} ->
        data = send_auth_frame(payload, %{data | implementation_state: new_handler_state})
        {:keep_state, data}

      {:error, reason} ->
        Logger.error("[#{__MODULE__}] Auth failed: #{inspect(reason)}")
        {:stop, {:auth_failed, reason}}
    end
  end

  # response and server sent message handling
  defp on_message(message, :authenticated, data) do
    case data.implementation.packet_request_id(message) do
      # messages without an ID are assumed to not be a response
      nil ->
        body = data.implementation.packet_body(message)
        dispatch_notice(body, data)
        :keep_state_and_data

      # messages with an ID are a response we need to send back to the caller
      request_id ->
        case Map.pop(data.pending, request_id) do
          {nil, _pending} ->
            Logger.warning(
              "[#{__MODULE__}] Unexpected response for request_id=#{request_id}: #{inspect(message)}"
            )

            :keep_state_and_data

          {from, pending} ->
            body = data.implementation.packet_body(message)
            {:keep_state, %{data | pending: pending}, [{:reply, from, {:ok, body}}]}
        end
    end
  end

  defp on_message(message, state, _data) do
    Logger.warning("[#{__MODULE__}] Unhandled message in state #{state}: #{inspect(message)}")
    :keep_state_and_data
  end

  # lets handler decide what to do after disconnection
  defp on_disconnect(reason, data) do
    Logger.warning("[#{__MODULE__}] Disconnected: #{inspect(reason)}")

    case data.implementation.handle_disconnect(reason, self(), data.implementation_state) do
      {:reconnect, new_handler_state} ->
        reconnect(reason, %{data | implementation_state: new_handler_state})

      {:ok, new_handler_state} ->
        shutdown(reason, %{data | implementation_state: new_handler_state})
    end
  end

  # fail in-flight requests, restart the connection
  defp reconnect(reason, data) do
    Logger.info("[#{__MODULE__}] Reconnecting")
    error = {:error, {:disconnected, reason}}
    pending_replies = Enum.map(data.pending, fn {_id, from} -> {:reply, from, error} end)

    old_conn = data.conn

    data = %{
      data
      | conn: nil,
        stream_ref: nil,
        implementation_state: data.init_implementation_state,
        connection_attempt: 0,
        pending: %{}
    }

    if old_conn, do: :gun.close(old_conn)

    {:next_state, :disconnected, data, [initial_event(data) | pending_replies]}
  end

  # fail all callers and stop
  defp shutdown(reason, data) do
    error = {:error, {:disconnected, reason}}
    pending_replies = Enum.map(data.pending, fn {_id, from} -> {:reply, from, error} end)

    if data.conn, do: :gun.close(data.conn)

    {:stop_and_reply, {:shutdown, reason}, pending_replies, data}
  end

  # every other notice_handler runs in a task, so that it can call the socket client
  defp dispatch_notice(notice, data) do
    handler = data.handler
    client_pid = self()

    Task.Supervisor.start_child(data.task_sup, fn ->
      try do
        handler.handle_notice(notice, client_pid)
      rescue
        e ->
          Logger.error(
            "[#{__MODULE__}] #{inspect(handler)}.handle_notice/2 raised: " <>
              Exception.format(:error, e, __STACKTRACE__)
          )
      end
    end)
  end

  # send a caller request frame and register the caller in pending
  defp send_request_frame(payload, from, data) do
    {request_id, data} = next_request_id(data)

    {:ok, {frame_type, encoded, new_handler_state}} =
      data.implementation.handle_encode(request_id, payload, data.implementation_state)

    :ok = :gun.ws_send(data.conn, data.stream_ref, {frame_type, encoded})

    if from do
      pending = Map.put(data.pending, request_id, from)
      %{data | implementation_state: new_handler_state, pending: pending}
    else
      %{data | implementation_state: new_handler_state}
    end
  end

  # auth request are sent from within the socket process, so we opt-out of called tracking
  defp send_auth_frame(payload, data) do
    send_request_frame(payload, nil, data)
  end

  defp next_request_id(data) do
    id = data.request_id + 1
    {id, %{data | request_id: id}}
  end

  defp next_connection_attempt(data) do
    %{data | connection_attempt: data.connection_attempt + 1}
  end

  defp backoff(attempt) do
    cond do
      attempt < 3 -> :timer.seconds(attempt * 5)
      attempt < 10 -> :timer.seconds(attempt * 30)
      true -> @default_reconnect_interval
    end
  end
end