defmodule Membrane.Transcoder.Video do
  @moduledoc false

  import Membrane.ChildrenSpec
  require Membrane.Logger
  alias Membrane.{ChildrenSpec, H264, H265, RawVideo, RemoteStream, VP8, VP9}
  alias Membrane.FFmpeg.SWScale

  @type video_stream_format :: VP8.t() | VP9.t() | H264.t() | H265.t() | RawVideo.t()

  defguardp is_raw_video_format(format)
            when is_struct(format) and format.__struct__ == RawVideo

  defguardp is_h26x_format(format)
            when is_struct(format) and
                   (format.__struct__ in [H264, H265] or
                      (format.__struct__ == RemoteStream and
                         format.content_format in [H264, H265]))

  defguardp is_vpx_format(format)
            when is_struct(format) and
                   (format.__struct__ in [VP8, VP9] or
                      (format.__struct__ == RemoteStream and
                         format.content_format in [VP8, VP9] and
                         format.type == :packetized))

  defguard is_video_format(format)
           when is_h26x_format(format) or
                  is_vpx_format(format) or
                  is_raw_video_format(format)

  @spec plug_video_transcoding(
          ChildrenSpec.builder(),
          video_stream_format(),
          video_stream_format(),
          :always | :if_needed | :never,
          boolean()
        ) :: ChildrenSpec.builder()
  def plug_video_transcoding(
        builder,
        input_format,
        output_format,
        transcoding_policy,
        use_hardware_acceleration?
      )
      when is_video_format(input_format) and is_video_format(output_format) do
    do_plug_video_transcoding(
      builder,
      input_format,
      output_format,
      transcoding_policy,
      use_hardware_acceleration?
    )
  end

  defp do_plug_video_transcoding(
         builder,
         %RemoteStream{content_format: h26x},
         output_format,
         transcoding_policy,
         use_hardware_acceleration?
       )
       when h26x in [H264, H265] do
    do_plug_video_transcoding(
      builder,
      struct!(h26x),
      output_format,
      transcoding_policy,
      use_hardware_acceleration?
    )
  end

  defp do_plug_video_transcoding(
         builder,
         %H264{},
         %H264{} = output_format,
         transcoding_policy,
         _use_hardware_acceleration?
       )
       when transcoding_policy in [:if_needed, :never] do
    builder
    |> child(:h264_parser, %H264.Parser{
      output_stream_structure: stream_structure_type(output_format),
      output_alignment: output_format.alignment
    })
  end

  defp do_plug_video_transcoding(
         builder,
         %H265{},
         %H265{} = output_format,
         transcoding_policy,
         _use_hardware_acceleration?
       )
       when transcoding_policy in [:if_needed, :never] do
    builder
    |> child(:h265_parser, %H265.Parser{
      output_stream_structure: stream_structure_type(output_format),
      output_alignment: output_format.alignment
    })
  end

  defp do_plug_video_transcoding(
         builder,
         %RawVideo{} = input_format,
         %RawVideo{} = output_format,
         _transcoding_policy,
         true
       ) do
    builder
    |> maybe_plug_swscale_converter_vulkan(input_format, output_format)
  end

  defp do_plug_video_transcoding(
         builder,
         %RawVideo{} = input_format,
         %RawVideo{} = output_format,
         _transcoding_policy,
         false
       ) do
    builder
    |> maybe_plug_swscale_converter(input_format, output_format)
  end

  defp do_plug_video_transcoding(
         builder,
         %format_module{},
         %format_module{},
         transcoding_policy,
         _use_hardware_acceleration?
       )
       when transcoding_policy in [:if_needed, :never] do
    Membrane.Logger.debug("""
    This bin will only forward buffers, as the input stream format is the same type as the output stream format.
    """)

    builder
  end

  defp do_plug_video_transcoding(
         _builder,
         input_format,
         output_format,
         :never,
         _use_hardware_acceleration?
       ),
       do:
         raise("""
         Cannot convert input format #{inspect(input_format)} to output format #{inspect(output_format)} \
         with :transcoding_policy option set to :never.
         """)

  defp do_plug_video_transcoding(builder, input_format, output_format, _transcoding_policy, true) do
    builder
    |> maybe_plug_parser_and_decoder_vulkan(input_format)
    |> maybe_plug_swscale_converter_vulkan(input_format, output_format)
    |> maybe_plug_encoder_and_parser_vulkan(output_format)
  end

  defp do_plug_video_transcoding(builder, input_format, output_format, _transcoding_policy, false) do
    builder
    |> maybe_plug_parser_and_decoder(input_format)
    |> maybe_plug_swscale_converter(input_format, output_format)
    |> maybe_plug_encoder_and_parser(output_format)
  end

  # VK-specific decoder: child name :vk_h264_decoder distinguishes it from FFmpeg's :h264_decoder
  defp maybe_plug_parser_and_decoder_vulkan(builder, %H264{}) do
    builder
    |> child(:h264_input_parser, %H264.Parser{
      output_stream_structure: :annexb,
      output_alignment: :au
    })
    |> child(:vk_h264_decoder, Membrane.VKVideo.Decoder)
  end

  defp maybe_plug_parser_and_decoder_vulkan(builder, format),
    do: maybe_plug_parser_and_decoder(builder, format)

  defp maybe_plug_parser_and_decoder(builder, %H264{}) do
    builder
    |> child(:h264_input_parser, %H264.Parser{
      output_stream_structure: :annexb,
      output_alignment: :au
    })
    |> child(:h264_decoder, %H264.FFmpeg.Decoder{})
  end

  defp maybe_plug_parser_and_decoder(builder, %H265{}) do
    builder
    |> child(:h265_input_parser, %H265.Parser{
      output_stream_structure: :annexb,
      output_alignment: :au
    })
    |> child(:h265_decoder, %H265.FFmpeg.Decoder{})
  end

  defp maybe_plug_parser_and_decoder(builder, %vpx{}) when vpx in [VP8, VP9] do
    decoder_module = Module.concat(vpx, Decoder)
    builder |> child(:vp8_decoder, decoder_module)
  end

  defp maybe_plug_parser_and_decoder(builder, %RemoteStream{
         content_format: vpx,
         type: :packetized
       })
       when vpx in [VP8, VP9] do
    decoder_module = Module.concat(vpx, Decoder)
    builder |> child(:vp8_decoder, decoder_module)
  end

  defp maybe_plug_parser_and_decoder(builder, %RawVideo{}), do: builder

  # VK decoder outputs NV12; these clauses handle the NV12 <-> other-format conversion
  # when mixing VK decode/encode with different pixel formats.

  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{pixel_format: nil}),
    do: builder

  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{pixel_format: :NV12}),
    do: builder

  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{} = output_format) do
    builder |> child(:raw_video_converter, %SWScale.Converter{format: output_format.pixel_format})
  end

  defp maybe_plug_swscale_converter_vulkan(builder, %RawVideo{pixel_format: :NV12}, %H264{}),
    do: builder

  defp maybe_plug_swscale_converter_vulkan(builder, %RawVideo{}, %H264{}) do
    builder |> child(:raw_video_converter, %SWScale.Converter{format: :NV12})
  end

  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %H264{}), do: builder

  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, _output_format) do
    builder |> child(:raw_video_converter, %SWScale.Converter{format: :I420})
  end

  defp maybe_plug_swscale_converter_vulkan(builder, _input_format, %H264{}) do
    builder |> child(:raw_video_converter, %SWScale.Converter{format: :NV12})
  end

  defp maybe_plug_swscale_converter_vulkan(builder, input_format, output_format),
    do: maybe_plug_swscale_converter(builder, input_format, output_format)

  defp maybe_plug_swscale_converter(builder, input_format, %RawVideo{} = output_format) do
    case input_format do
      _any when output_format.pixel_format == nil ->
        builder

      %RawVideo{pixel_format: pixel_format} when pixel_format == output_format.pixel_format ->
        builder

      _input_format ->
        builder
        |> child(:raw_video_converter, %SWScale.Converter{format: output_format.pixel_format})
    end
  end

  defp maybe_plug_swscale_converter(builder, input_format, %h26x{}) when h26x in [H264, H265] do
    case input_format do
      %RawVideo{pixel_format: pixel_format} when pixel_format in [:I420, :I422] ->
        builder

      %h26x{} when h26x in [H264, H265] ->
        builder

      _input_format ->
        builder
        |> child(:raw_video_converter, %SWScale.Converter{format: :I420})
    end
  end

  defp maybe_plug_swscale_converter(builder, _input_format, _output_format), do: builder

  defp maybe_plug_encoder_and_parser_vulkan(builder, %H264{} = h264) do
    builder
    |> child(:vk_h264_encoder, Membrane.VKVideo.Encoder)
    |> child(:h264_output_parser, %H264.Parser{
      output_stream_structure: stream_structure_type(h264),
      output_alignment: h264.alignment
    })
  end

  defp maybe_plug_encoder_and_parser_vulkan(builder, format),
    do: maybe_plug_encoder_and_parser(builder, format)

  defp maybe_plug_encoder_and_parser(builder, %H264{} = h264) do
    builder
    |> child(:h264_encoder, %H264.FFmpeg.Encoder{preset: :ultrafast})
    |> child(:h264_output_parser, %H264.Parser{
      output_stream_structure: stream_structure_type(h264),
      output_alignment: h264.alignment
    })
  end

  defp maybe_plug_encoder_and_parser(builder, %H265{} = h265) do
    builder
    |> child(:h265_encoder, %H265.FFmpeg.Encoder{preset: :ultrafast})
    |> child(:h265_output_parser, %H265.Parser{
      output_stream_structure: stream_structure_type(h265),
      output_alignment: h265.alignment
    })
  end

  defp maybe_plug_encoder_and_parser(builder, %VP8{}) do
    cpu_quota = :erlang.system_info(:cpu_quota)

    number_of_threads =
      if cpu_quota != :unknown,
        do: cpu_quota,
        else: :erlang.system_info(:logical_processors_available)

    builder |> child(:vp8_encoder, %VP8.Encoder{g_threads: number_of_threads, cpu_used: 15})
  end

  defp maybe_plug_encoder_and_parser(builder, %VP9{}) do
    cpu_quota = :erlang.system_info(:cpu_quota)

    number_of_threads =
      if cpu_quota != :unknown,
        do: cpu_quota,
        else: :erlang.system_info(:logical_processors_available)

    builder |> child(:vp8_encoder, %VP9.Encoder{g_threads: number_of_threads, cpu_used: 15})
  end

  defp maybe_plug_encoder_and_parser(builder, %RawVideo{}), do: builder

  defp stream_structure_type(%h26x{stream_structure: stream_structure})
       when h26x in [H264, H265] do
    case stream_structure do
      type when type in [:annexb, :avc1, :avc3, :hvc1, :hev1] -> type
      {type, _dcr} when type in [:avc1, :avc3, :hvc1, :hev1] -> type
    end
  end
end