membrane_mp3_lame_plugin

defmodule Membrane.MP3.Lame.Encoder do
  @moduledoc """
  Element encoding raw audio into MPEG-1, layer 3 format
  """
  use Membrane.Filter

  require Membrane.Logger

  alias __MODULE__.{Native, RateControl}
  alias Membrane.Buffer
  alias Membrane.MPEGAudio
  alias Membrane.RawAudio

  @samples_per_frame 1152
  @channels 2
  @sample_size 4

  @typedoc """
  VBR configuration for the `rate_control` option.

  For details, see Lame's header file:
  https://sourceforge.net/p/lame/svn/HEAD/tree/trunk/lame/include/lame.h

  Available options:
  * `:mode` (required) - Sets the algorithm via `lame_set_VBR`.
  * `:mean_bitrate` (required for `:abr`, not allowed otherwise) -
    target average bitrate in kbps. Calls  `lame_set_VBR_mean_bitrate_kbps`.
  * `:quality` (only for `:rh`, `:mt`, `:mtrh`, `:default`) -
    VBR quality level as a number in `[0, 10)` (0 = best, 9 = worst).
    Calls `lame_set_VBR_quality`, which accepts fractional values
    for finer-grained control. If omitted, LAME's default is used.
  * `:min_bitrate` - minimum allowed bitrate in kbps. See
    `lame_set_VBR_min_bitrate_kbps` for reference.
  * `:max_bitrate` - maximum allowed bitrate in kbps. Calls
    `lame_set_VBR_max_bitrate_kbps`.
  * `:hard_min` (boolean) - strictly enforce `:min_bitrate` even
    for silence. Calls `lame_set_VBR_hard_min`.
  """
  @type vbr_config :: [
          mode: :mt | :rh | :abr | :mtrh | :default,
          mean_bitrate: pos_integer(),
          quality: number(),
          min_bitrate: pos_integer(),
          max_bitrate: pos_integer(),
          hard_min: boolean()
        ]

  def_output_pad :output,
    accepted_format: %MPEGAudio{channels: 2, sample_rate: 44_100, layer: :layer3, version: :v1}

  def_input_pad :input,
    accepted_format:
      any_of(
        %RawAudio{sample_format: :s32le, sample_rate: 44_100, channels: 2},
        Membrane.RemoteStream
      )

  def_options gapless_flush: [
                spec: boolean(),
                default: true,
                description: """
                When this option is set to true, encoder will be flushed without
                outputting any tags and allowing to play such file gaplessly
                if concatenated with another file.
                """
              ],
              bitrate: [
                spec: pos_integer() | nil,
                default: nil,
                description: """
                Deprecated: pass the bitrate via the `rate_control` option instead,
                e.g. `rate_control: {:cbr, bitrate: 192}`.

                Output bitrate of encoded stream in kbit/sec, defaults to 192.
                """
              ],
              quality: [
                spec: non_neg_integer(),
                default: 5,
                description: """
                Value of this parameter affects quality by selecting one of the algorithms
                for encoding: `0` being best (and very slow) and `9` being worst.

                Recommended values:
                  * `2` - near-best quality, not too slow
                  * `5` - good quality, fast
                  * `7` - ok quality, really fast
                """
              ],
              disable_reservoir: [
                spec: boolean(),
                default: false,
                description: """
                When set to true, disables the LAME bit reservoir. This makes each
                MP3 frame self-contained (main_data_begin is always 0), so frames
                can be cleanly spliced at any boundary without decoder artifacts.

                Useful for live streaming with ad insertion where audio from different
                sources is concatenated.
                """
              ],
              rate_control: [
                spec: :cbr | {:cbr, bitrate: pos_integer()} | {:vbr, vbr_config()},
                default: :cbr,
                description: """
                Rate control mode for the encoder:
                  * `:cbr` - Constant Bitrate at the default bitrate (192 kbps).
                  * `{:cbr, bitrate: kbps}` - Constant Bitrate at the given bitrate.
                  * `{:vbr, config}` - Variable Bitrate. See `t:vbr_config/0` for details.
                """
              ]

  @impl true
  def handle_init(_ctx, options) do
    {[],
     %{
       native: nil,
       queue: <<>>,
       options: options,
       raw_frame_size: MPEGAudio.samples_per_frame(:v1, :layer3) * @sample_size * @channels,
       next_frame_pts: nil
     }}
  end

  @impl true
  def handle_setup(_ctx, state) do
    quality_val = state.options.quality |> map_quality_to_value!

    rate_control =
      state.options.rate_control
      |> apply_legacy_bitrate(state.options.bitrate)
      |> RateControl.parse!()

    with {:ok, native} <-
           Native.create(
             @channels,
             quality_val,
             state.options.disable_reservoir,
             rate_control
           ) do
      {[], %{state | native: native}}
    else
      {:error, reason} ->
        raise "Error initializing mpeg: #{inspect(reason)}"
    end
  end

  defp apply_legacy_bitrate(rate_control, nil), do: rate_control

  defp apply_legacy_bitrate(rate_control, legacy_bitrate) do
    Membrane.Logger.warning("""
    The `bitrate` option is deprecated. Pass the bitrate via `rate_control: \
    {:cbr, bitrate: #{legacy_bitrate}}` instead.
    """)

    if rate_control == :cbr, do: {:cbr, bitrate: legacy_bitrate}, else: rate_control
  end

  @impl true
  def handle_playing(_ctx, state) do
    stream_format = %MPEGAudio{
      channels: 2,
      sample_rate: 44_100,
      version: :v1,
      layer: :layer3,
      bitrate: 192
    }

    {[stream_format: {:output, stream_format}], state}
  end

  @impl true
  def handle_stream_format(:input, _stream_format, _ctx, state) do
    {[], state}
  end

  @impl true
  def handle_end_of_stream(:input, __ctx, %{queue: ""} = state) do
    {[notify_parent: {:end_of_stream, :input}, end_of_stream: :output], state}
  end

  def handle_end_of_stream(:input, _ctx, state) do
    %{native: native, queue: queue, options: options, next_frame_pts: pts} = state

    buffers = encode_last_frame!(native, queue, pts, options.gapless_flush)
    actions = [end_of_stream: :output, notify_parent: {:end_of_stream, :input}]
    {buffers ++ actions, %{state | queue: <<>>}}
  end

  @impl true
  def handle_buffer(:input, %Buffer{payload: data, pts: pts}, _ctx, state) do
    %{native: native, queue: queue, next_frame_pts: maybe_next_frame_pts} = state
    next_frame_pts = if queue == <<>>, do: pts, else: maybe_next_frame_pts

    to_encode = queue <> data

    case encode_buffer(native, to_encode, next_frame_pts, pts) do
      {:ok, {[], 0}} ->
        {[], %{state | queue: to_encode, next_frame_pts: next_frame_pts}}

      {:ok, {encoded_bufs, bytes_used}} ->
        <<_handled::binary-size(bytes_used), rest::binary>> = to_encode
        {[buffer: {:output, encoded_bufs}], %{state | queue: rest, next_frame_pts: pts}}

      {:error, reason} ->
        raise "Error #{inspect(reason)}"
    end
  end

  # init
  defp encode_buffer(native, buffer, next_frame_pts, rest_frames_pts) do
    raw_frame_size = @samples_per_frame * @sample_size * @channels

    encode_buffer(native, buffer, next_frame_pts, rest_frames_pts, [], 0, raw_frame_size)
  end

  # handle single frame
  defp encode_buffer(native, buffer, this_frame_pts, rest_pts, acc, bytes_used, raw_frame_size)
       when byte_size(buffer) >= raw_frame_size do
    <<raw_frame::binary-size(raw_frame_size), rest::binary>> = buffer

    case Native.encode_frame(raw_frame, native) do
      {:ok, encoded_frame} ->
        encoded_buffer = %Buffer{payload: encoded_frame, pts: this_frame_pts}

        encode_buffer(
          native,
          rest,
          rest_pts,
          rest_pts,
          [encoded_buffer | acc],
          bytes_used + raw_frame_size,
          raw_frame_size
        )

      {:error, reason} ->
        Membrane.Logger.error("Terminating stream because of malformed frame. Reason: #{reason}")
        {:error, reason}
    end
  end

  # Not enough samples for a frame
  defp encode_buffer(_native, _partial_buffer, _pts, _next_pts, acc, bytes_used, _raw_frame_size) do
    {:ok, {acc |> Enum.reverse(), bytes_used}}
  end

  defp encode_last_frame!(native, queue, pts, gapless?) do
    with {:ok, encoded_frame} <- Native.encode_frame(queue, native),
         {:ok, flushed_frame} <- Native.flush(gapless?, native) do
      bufs =
        [encoded_frame, flushed_frame]
        |> Enum.flat_map(fn
          "" -> []
          frame -> [%Buffer{payload: frame, pts: pts}]
        end)

      [buffer: {:output, bufs}]
    else
      {:error, reason} ->
        Membrane.Logger.error(
          "Terminating stream because of malformed last frame. Reason: #{reason}"
        )

        raise "Error #{inspect(reason)}"
    end
  end

  defp map_quality_to_value!(quality) when quality in 0..9, do: quality

  defp map_quality_to_value!(value),
    do: raise("Error parsing quality argument: #{inspect(value)}")
end