defmodule Membrane.RTP.Muxer do
  @moduledoc """
  Element that combines multiple streams into a single RTP stream.

  Each new input stream is assigned a unique SSRC that the packets transporting this stream will have.
  The demuxer will try to resolve what `payload_type` and `clock_rate` should be assumed based on the
  incoming stream format and passed pad options.
  Timestamps are calculated based on assumed `clock_rate`.
  """
  use Membrane.Filter

  require Membrane.Pad

  alias Membrane.{Pad, RemoteStream, RTP}

  @max_ssrc Bitwise.bsl(1, 32) - 1
  @max_sequence_number Bitwise.bsl(1, 16) - 1
  @max_timestamp Bitwise.bsl(1, 32) - 1

  @payload_format_to_encoding_name %{
    Membrane.H264 => :H264,
    Membrane.H265 => :H265,
    Membrane.VP8 => :VP8,
    Membrane.AAC => :AAC,
    Membrane.Opus => :opus,
    Membrane.MPEGAudio => :MPA
  }

  def_input_pad :input,
    accepted_format: RTP,
    availability: :on_request,
    options: [
      ssrc: [
        spec: RTP.ssrc() | :random,
        default: :random,
        description: """
        SSRC that this stream will be assigned. If not provided, a random free value will be assigned.
        """
      ],
      initial_sequence_number: [
        spec: ExRTP.Packet.uint16() | :random,
        default: :random,
        description: """
        Sequence number that will be assigned to the first packet of this stream. Sequence numbers of 
        subsequent packets will be determined by incrementing this value. This value should be generated 
        randomly to prevent known payload attacks.
        """
      ],
      initial_timestamp: [
        spec: ExRTP.Packet.uint16() | :random,
        default: :random,
        description: """
        RTP Timestamp that will be assigned to the first packet of this stream. Timestamps of subsequent 
        packets will be calculated starting from this value. This value should be generated 
        randomly to prevent known payload attacks.
        """
      ],
      payload_type: [
        spec: RTP.payload_type() | nil,
        default: nil,
        description: """
        Payload type of the stream. If not provided, determined from the resolved encoding name.
        """
      ],
      encoding: [
        spec: RTP.encoding_name() | nil,
        default: nil,
        description: """
        Encoding name of the stream that will be used if the muxer fails to infer it from incoming stream format.
        Used for determining payload_type, if it wasn't provided.
        """
      ],
      clock_rate: [
        spec: non_neg_integer() | nil,
        default: nil,
        description: """
        Clock rate to use. If not provided, determined from resolved payload type.
        """
      ]
    ]

  def_output_pad :output, accepted_format: %RemoteStream{type: :packetized, content_format: RTP}

  def_options payload_type_mapping: [
                spec: RTP.PayloadFormat.payload_type_mapping(),
                default: %{},
                description: "Mapping of the custom RTP payload types ( > 95)."
              ],
              srtp: [
                spec: false | [ExLibSRTP.Policy.t()],
                default: false,
                description: """
                Specifies whether to use SRTP to encrypt the output stream. Requires adding [srtp](https://github.com/membraneframework/elixir_libsrtp)
                dependency to work. If true takes a list of SRTP policies to use for encrypting packets. See `t:ExLibSRTP.Policy.t/0` for details.
                """
              ]

  defmodule State do
    @moduledoc false

    use Bunch.Access

    defmodule StreamState do
      @moduledoc false
      use Bunch.Access

      alias Membrane.RTP

      @type t :: %__MODULE__{
              ssrc: RTP.ssrc(),
              sequence_number: ExRTP.Packet.uint16(),
              initial_timestamp: ExRTP.Packet.uint32(),
              clock_rate: RTP.clock_rate(),
              payload_type: RTP.payload_type(),
              end_of_stream: boolean()
            }

      @enforce_keys [:ssrc, :sequence_number, :initial_timestamp, :clock_rate, :payload_type]

      defstruct @enforce_keys ++ [end_of_stream: false]
    end

    @type t :: %__MODULE__{
            stream_states: %{Pad.ref() => StreamState.t()},
            payload_type_mapping: RTP.PayloadFormat.payload_type_mapping(),
            srtp: ExLibSRTP.t() | nil
          }

    @enforce_keys [:srtp, :payload_type_mapping]
    defstruct @enforce_keys ++ [stream_states: %{}]
  end

  @impl true
  def handle_init(_ctx, opts) do
    srtp =
      case opts.srtp do
        false ->
          nil

        policies ->
          if not Code.ensure_loaded?(ExLibSRTP) do
            raise "Optional dependency :ex_libsrtp is required for SRTP"
          end

          srtp = apply(ExLibSRTP, :new, [])
          Enum.each(policies, &apply(ExLibSRTP, :add_stream, [srtp, &1]))
          srtp
      end

    {[], %State{srtp: srtp, payload_type_mapping: opts.payload_type_mapping}}
  end

  @impl true
  def handle_playing(_ctx, state) do
    {[stream_format: {:output, %RemoteStream{type: :packetized, content_format: RTP}}], state}
  end

  @impl true
  def handle_stream_format(pad, stream_format, ctx, state) do
    pad_options = ctx.pads[pad].options

    ssrc = get_stream_ssrc(pad_options.ssrc, state)

    encoding_name =
      @payload_format_to_encoding_name[stream_format.payload_format] ||
        pad_options.encoding

    %{payload_type: payload_type, clock_rate: clock_rate} =
      RTP.PayloadFormat.resolve(
        encoding_name: encoding_name,
        payload_type: pad_options.payload_type,
        clock_rate: pad_options.clock_rate,
        payload_type_mapping: state.payload_type_mapping
      )

    if payload_type == nil do
      raise "Could not resolve payload type, information provided via pad options and stream format not sufficient"
    end

    if clock_rate == nil do
      raise "Could not resolve clock rate, information provided via pad options and stream format not sufficient"
    end

    new_stream_state = %State.StreamState{
      ssrc: ssrc,
      sequence_number:
        if(pad_options.initial_sequence_number == :random,
          do: Enum.random(0..@max_sequence_number),
          else: pad_options.initial_sequence_number
        ),
      initial_timestamp:
        if(pad_options.initial_timestamp == :random,
          do: Enum.random(0..@max_timestamp),
          else: pad_options.initial_timestamp
        ),
      clock_rate: clock_rate,
      payload_type: payload_type
    }

    state = put_in(state.stream_states[pad], new_stream_state)

    {[], state}
  end

  @impl true
  def handle_buffer(Pad.ref(:input, _ref) = pad_ref, %Membrane.Buffer{} = buffer, _ctx, state) do
    rtp_metadata = Map.get(buffer.metadata, :rtp, %{})
    stream_state = state.stream_states[pad_ref]

    rtp_offset =
      buffer.pts
      |> Membrane.Time.as_seconds()
      |> Numbers.mult(stream_state.clock_rate)
      |> Ratio.trunc()

    timestamp = rem(stream_state.initial_timestamp + rtp_offset, @max_timestamp + 1)

    packet =
      ExRTP.Packet.new(buffer.payload,
        payload_type: stream_state.payload_type,
        sequence_number: stream_state.sequence_number,
        timestamp: timestamp,
        ssrc: stream_state.ssrc,
        csrc: Map.get(rtp_metadata, :csrcs, []),
        marker: Map.get(rtp_metadata, :marker, false)
      )

    state =
      update_in(
        state.stream_states[pad_ref].sequence_number,
        &rem(&1 + 1, @max_sequence_number + 1)
      )

    buffer_action =
      packet
      |> ExRTP.Packet.encode()
      |> protect_packet(state.srtp)
      |> case do
        nil ->
          []

        raw_packet ->
          buffer = %Membrane.Buffer{
            buffer
            | payload: raw_packet,
              metadata: Map.put(buffer.metadata, :rtp, %{packet | payload: <<>>})
          }

          [buffer: {:output, buffer}]
      end

    {buffer_action, state}
  end

  @impl true
  def handle_end_of_stream(Pad.ref(:input, _ref) = pad_ref, ctx, state) do
    state =
      if is_map_key(state.stream_states, pad_ref) do
        state |> put_in([:stream_states, pad_ref, :end_of_stream], true)
      else
        state
      end

    all_input_pad_refs =
      ctx.pads
      |> Enum.flat_map(fn {pad_ref, %{direction: direction}} ->
        if direction == :input, do: [pad_ref], else: []
      end)
      |> Enum.concat(Map.keys(state.stream_states))
      |> Enum.uniq()

    all_input_pad_refs
    |> Enum.all?(fn pad_ref ->
      (ctx.pads[pad_ref] != nil and ctx.pads[pad_ref].end_of_stream?) or
        (state.stream_states[pad_ref] != nil and state.stream_states[pad_ref].end_of_stream)
    end)
    |> if do
      {[end_of_stream: :output], state}
    else
      {[], state}
    end
  end

  @spec protect_packet(binary(), ExLibSRTP.t() | nil) :: binary() | nil
  defp protect_packet(rtp_packet, nil) do
    rtp_packet
  end

  defp protect_packet(rtp_packet, srtp) do
    case apply(ExLibSRTP, :protect, [srtp, rtp_packet]) do
      {:ok, protected_rtp_packet} ->
        protected_rtp_packet

      {:error, reason} when reason in [:replay_fail, :replay_old] ->
        Membrane.Logger.warning("Ignoring packet due to `#{reason}`")
        nil

      {:error, reason} ->
        raise "Failed to unprotect packet due to `#{reason}`"
    end
  end

  @spec get_stream_ssrc(RTP.ssrc() | :random, State.t()) :: RTP.ssrc()
  defp get_stream_ssrc(ssrc, state) do
    assigned_ssrcs = Enum.map(state.stream_states, fn {_pad_ref, %{ssrc: ssrc}} -> ssrc end)

    case ssrc do
      :random ->
        Stream.repeatedly(fn -> Enum.random(0..@max_ssrc) end)
        |> Enum.find(&(&1 not in assigned_ssrcs))

      provided_ssrc ->
        if provided_ssrc in assigned_ssrcs do
          raise("SSRC #{provided_ssrc} already assigned to a different stream")
        end

        provided_ssrc
    end
  end
end