kadabra

defmodule Kadabra.Stream do
  @moduledoc false

  defstruct id: nil,
            body: "",
            client: nil,
            connection: nil,
            socket: nil,
            encoder: nil,
            decoder: nil,
            flow: nil,
            uri: nil,
            headers: [],
            on_response: nil

  require Logger

  alias Kadabra.{Encodable, Frame, Hpack, Packetizer, Socket, Stream, Tasks}

  alias Kadabra.Frame.{
    Continuation,
    Data,
    Headers,
    PushPromise,
    RstStream,
    WindowUpdate
  }

  alias Kadabra.Stream.Response

  @type t :: %__MODULE__{
          id: pos_integer,
          client: pid,
          connection: pid,
          uri: URI.t(),
          flow: Kadabra.Stream.FlowControl.t(),
          headers: [...],
          body: binary
        }

  @closed :closed
  @hc_local :half_closed_local
  @hc_remote :half_closed_remote
  @idle :idle
  @open :open
  # @reserved_local :reserved_local
  @reserved_remote :reserved_remote

  def new(config, stream_id, initial_window_size, max_frame_size) do
    flow_opts = [
      window: initial_window_size,
      max_frame_size: max_frame_size
    ]

    %__MODULE__{
      id: stream_id,
      client: config.client,
      uri: config.uri,
      socket: config.socket,
      encoder: config.encoder,
      decoder: config.decoder,
      connection: self(),
      flow: Stream.FlowControl.new(flow_opts)
    }
  end

  def start_link(%Stream{} = stream) do
    :gen_statem.start_link(__MODULE__, stream, [])
  end

  def close(pid) do
    call_recv(pid, :close)
  end

  def call_recv(pid, frame) do
    :gen_statem.call(pid, {:recv, frame})
  end

  def cast_send(pid, frame) do
    :gen_statem.cast(pid, {:send, frame})
  end

  # recv

  def recv(from, :close, _state, _stream) do
    {:stop, :normal, [{:reply, from, :ok}]}
  end

  def recv(from, %Data{end_stream: true, data: data}, state, %Stream{} = stream)
      when state in [@hc_local] do
    :gen_statem.reply(from, :ok)
    stream = %Stream{stream | body: stream.body <> data}
    {:next_state, @closed, stream}
  end

  def recv(
        from,
        %Data{end_stream: true, data: data},
        _state,
        %Stream{} = stream
      ) do
    :gen_statem.reply(from, :ok)
    stream = %Stream{stream | body: stream.body <> data}
    {:next_state, @hc_remote, stream}
  end

  def recv(
        from,
        %Data{end_stream: false, data: data},
        _state,
        %Stream{} = stream
      ) do
    :gen_statem.reply(from, :ok)
    stream = %Stream{stream | body: stream.body <> data}
    {:keep_state, stream}
  end

  def recv(
        from,
        %Headers{end_stream: end_stream?} = frame,
        _state,
        %Stream{} = stream
      ) do
    case Hpack.decode(stream.decoder, frame.header_block_fragment) do
      {:ok, headers} ->
        :gen_statem.reply(from, :ok)

        stream = %Stream{stream | headers: stream.headers ++ headers}

        if end_stream?,
          do: {:next_state, @hc_remote, stream},
          else: {:keep_state, stream}

      _error ->
        :gen_statem.reply(from, {:connection_error, :COMPRESSION_ERROR})
        {:stop, :normal}
    end
  end

  def recv(from, %RstStream{} = _frame, state, stream)
      when state in [@open, @hc_local, @hc_remote, @closed] do
    # IO.inspect(frame, label: "Got RST_STREAM")
    {:next_state, :closed, stream, [{:reply, from, :ok}]}
  end

  def recv(from, %PushPromise{} = frame, state, %Stream{} = stream)
      when state in [@idle] do
    {:ok, headers} = Hpack.decode(stream.decoder, frame.header_block_fragment)

    stream = %Stream{stream | headers: stream.headers ++ headers}

    :gen_statem.reply(from, :ok)

    response = Response.new(stream.id, stream.headers, stream.body)
    send(stream.connection, {:push_promise, response})
    {:next_state, @reserved_remote, stream}
  end

  def recv(from, %WindowUpdate{window_size_increment: inc}, _state, stream) do
    :gen_statem.reply(from, :ok)

    flow =
      stream.flow
      |> Stream.FlowControl.increment_window(inc)
      |> Stream.FlowControl.process()
      |> send_data_frames(stream.socket, stream.id)

    {:keep_state, %{stream | flow: flow}}
  end

  def recv(from, %Continuation{} = frame, _state, %Stream{} = stream) do
    {:ok, headers} = Hpack.decode(stream.decoder, frame.header_block_fragment)

    :gen_statem.reply(from, :ok)

    stream = %Stream{stream | headers: stream.headers ++ headers}
    {:keep_state, stream}
  end

  def recv(frame, state, stream) do
    """
    Unknown RECV on stream #{stream.id}
    Frame: #{inspect(frame)}
    State: #{inspect(state)}
    """
    |> Logger.info()

    {:keep_state, stream}
  end

  # Enter Events

  def handle_event(:enter, _old, @hc_remote, %{socket: socket} = stream) do
    bin = stream.id |> RstStream.new() |> Encodable.to_bin()
    Socket.send(socket, bin)

    :gen_statem.cast(self(), :close)
    {:keep_state, stream}
  end

  def handle_event(:enter, _old, @closed, stream) do
    response = Response.new(stream.id, stream.headers, stream.body)
    Tasks.run(stream.on_response, response)
    send(stream.client, {:end_stream, response})

    {:stop, {:shutdown, {:finished, stream.id}}}
  end

  def handle_event(:enter, _old, _new, stream), do: {:keep_state, stream}

  # For SETTINGS initial_window_size and max_frame_size changes
  def handle_event(:info, {:settings_change, window, max_frame}, _, stream) do
    flow =
      stream.flow
      |> Stream.FlowControl.increment_window(window)
      |> Stream.FlowControl.set_max_frame_size(max_frame)

    {:keep_state, %{stream | flow: flow}}
  end

  # Casts

  def handle_event(:cast, :close, _state, stream) do
    {:next_state, @closed, stream}
  end

  def handle_event(:cast, {:recv, frame}, state, stream) do
    recv(frame, state, stream)
  end

  def handle_event(:cast, msg, state, stream) do
    """
    === Unknown cast ===
    #{inspect(msg)}
    State: #{inspect(state)}
    Stream: #{inspect(stream)}
    """
    |> Logger.info()

    {:keep_state, stream}
  end

  # Calls

  def handle_event({:call, from}, {:recv, frame}, state, stream) do
    recv(from, frame, state, stream)
  end

  def handle_event({:call, from}, {:send_headers, request}, _state, stream) do
    %{headers: headers, body: payload, on_response: on_resp} = request

    headers_payload = encode_headers(stream.encoder, headers, stream.uri)
    max_size = stream.flow.max_frame_size
    send_headers(stream.socket, stream.id, headers_payload, payload, max_size)

    # Reply early for better performance
    :gen_statem.reply(from, :ok)

    stream =
      stream
      |> process_payload_if_needed(payload)
      |> Map.put(:on_response, on_resp)

    {:next_state, @open, stream}
  end

  defp encode_headers(pid, headers, uri) do
    headers = add_headers(headers, uri)
    {:ok, encoded} = Hpack.encode(pid, headers)
    :erlang.iolist_to_binary(encoded)
  end

  def add_headers(headers, %{scheme: scheme, authority: auth}) do
    h = headers ++ [{":scheme", scheme}, {":authority", auth}]

    # sorting headers to have pseudo headers first.
    Enum.sort(h, fn {a, _b}, {c, _d} -> a < c end)
  end

  defp send_headers(socket, stream_id, headers_payload, payload, max_size) do
    bin =
      stream_id
      |> Packetizer.headers(headers_payload, max_size, is_nil(payload))
      |> encode_and_flatten()

    Socket.send(socket, bin)
    # Logger.info("Sending, Stream ID: #{stream.id}, size: #{byte_size(h)}")
  end

  defp encode_and_flatten(frames) do
    Enum.reduce(frames, <<>>, &(&2 <> Encodable.to_bin(&1)))
  end

  @spec process_payload_if_needed(Stream.t(), binary | nil) :: Stream.t()
  defp process_payload_if_needed(stream, nil), do: stream

  defp process_payload_if_needed(stream, payload) do
    flow =
      stream.flow
      |> Stream.FlowControl.add(payload)
      |> Stream.FlowControl.process()
      |> send_data_frames(stream.socket, stream.id)

    %{stream | flow: flow}
  end

  def send_data_frames(flow_control, socket, stream_id) do
    bin =
      flow_control.out_queue
      |> :queue.to_list()
      |> Enum.map(fn {data, end_stream?} ->
        %Frame.Data{stream_id: stream_id, end_stream: end_stream?, data: data}
      end)
      |> encode_and_flatten()

    Socket.send(socket, bin)

    %{flow_control | out_queue: :queue.new()}
  end

  # Other Callbacks

  def init(stream), do: {:ok, @idle, stream}

  def callback_mode, do: [:handle_event_function, :state_enter]

  def terminate(_reason, _state, _stream), do: :void

  def code_change(_vsn, state, data, _extra), do: {:ok, state, data}
end