/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "WMFMediaDataEncoder.h"

#include <comdef.h>

#include "ImageContainer.h"
#include "ImageConversion.h"
#include "MFTEncoder.h"
#include "PlatformEncoderModule.h"
#include "TimeUnits.h"
#include "WMFDataEncoderUtils.h"
#include "WMFUtils.h"
#include "mozilla/WindowsProcessMitigations.h"
#include "mozilla/dom/WebCodecsUtils.h"

namespace mozilla {

#define AUTO_MARKER(desc) AUTO_WEBCODECS_MARKER("WMFMediaDataEncoder", desc)

using InitPromise = MediaDataEncoder::InitPromise;
using EncodePromise = MediaDataEncoder::EncodePromise;
using ReconfigurationPromise = MediaDataEncoder::ReconfigurationPromise;

WMFMediaDataEncoder::WMFMediaDataEncoder(const EncoderConfig& aConfig,
                                         const RefPtr<TaskQueue>& aTaskQueue)
    : mConfig(aConfig),
      mTaskQueue(aTaskQueue),
      mHardwareNotAllowed(aConfig.mHardwarePreference ==
                          HardwarePreference::RequireSoftware) {
  WMF_ENC_LOGE("WMFMediaDataEncoder ctor: %s, (hw not allowed: %s)",
               aConfig.ToString().get(), mHardwareNotAllowed ? "yes" : "no");
  MOZ_ASSERT(mTaskQueue);
}

RefPtr<InitPromise> WMFMediaDataEncoder::Init() {
  return InvokeAsync(mTaskQueue, this, __func__,
                     &WMFMediaDataEncoder::ProcessInit);
}
RefPtr<EncodePromise> WMFMediaDataEncoder::Encode(const MediaData* aSample) {
  WMF_ENC_LOGD("Encode ts=%s", aSample->mTime.ToString().get());
  MOZ_ASSERT(aSample);

  RefPtr<const VideoData> sample(aSample->As<const VideoData>());

  return InvokeAsync<RefPtr<const VideoData>>(
      mTaskQueue, this, __func__, &WMFMediaDataEncoder::ProcessEncode,
      std::move(sample));
}
RefPtr<EncodePromise> WMFMediaDataEncoder::Encode(
    nsTArray<RefPtr<MediaData>>&& aSamples) {
  WMF_ENC_LOGD("Encode: num of samples=%zu", aSamples.Length());
  MOZ_ASSERT(!aSamples.IsEmpty());

  nsTArray<RefPtr<const VideoData>> videoSamples;
  for (auto& sample : aSamples) {
    videoSamples.AppendElement(sample->As<const VideoData>());
  }

  return InvokeAsync(mTaskQueue, this, __func__,
                     &WMFMediaDataEncoder::ProcessEncodeBatch,
                     std::move(videoSamples));
}
RefPtr<EncodePromise> WMFMediaDataEncoder::Drain() {
  WMF_ENC_LOGD("Drain");
  return InvokeAsync(mTaskQueue, this, __func__,
                     &WMFMediaDataEncoder::ProcessDrain);
}
RefPtr<ShutdownPromise> WMFMediaDataEncoder::Shutdown() {
  WMF_ENC_LOGD("Shutdown");
  return InvokeAsync(
      mTaskQueue, __func__, [self = RefPtr<WMFMediaDataEncoder>(this)]() {
        auto r = MediaResult(NS_ERROR_DOM_MEDIA_CANCELED,
                             "Canceled by WMFMediaDataEncoder::Shutdown");

        // Cancel encode in flight if any.
        self->mEncodeRequest.DisconnectIfExists();
        self->mEncodePromise.RejectIfExists(r, __func__);

        // Cancel drain in flight if any.
        self->mDrainRequest.DisconnectIfExists();
        self->mDrainPromise.RejectIfExists(r, __func__);

        if (self->mEncoder) {
          self->mEncoder->Destroy();
          self->mEncoder = nullptr;
        }
        return ShutdownPromise::CreateAndResolve(true, __func__);
      });
}
RefPtr<GenericPromise> WMFMediaDataEncoder::SetBitrate(uint32_t aBitsPerSec) {
  return InvokeAsync(
      mTaskQueue, __func__,
      [self = RefPtr<WMFMediaDataEncoder>(this), aBitsPerSec]() {
        MOZ_ASSERT(self->mEncoder);
        return SUCCEEDED(self->mEncoder->SetBitrate(aBitsPerSec))
                   ? GenericPromise::CreateAndResolve(true, __func__)
                   : GenericPromise::CreateAndReject(
                         NS_ERROR_DOM_MEDIA_NOT_SUPPORTED_ERR, __func__);
      });
}

RefPtr<ReconfigurationPromise> WMFMediaDataEncoder::Reconfigure(
    const RefPtr<const EncoderConfigurationChangeList>& aConfigurationChanges) {
  // General reconfiguration interface not implemented right now
  return MediaDataEncoder::ReconfigurationPromise::CreateAndReject(
      NS_ERROR_DOM_MEDIA_FATAL_ERR, __func__);
};

nsCString WMFMediaDataEncoder::GetDescriptionName() const {
  return MFTEncoder::GetFriendlyName(CodecToSubtype(mConfig.mCodec));
}

RefPtr<InitPromise> WMFMediaDataEncoder::ProcessInit() {
  AssertOnTaskQueue();

  WMF_ENC_LOGD("ProcessInit");

  MOZ_ASSERT(!mEncoder,
             "Should not initialize encoder again without shutting down");

  auto cleanup = MakeScopeExit([&] { mIsHardwareAccelerated = false; });

  if (!wmf::MediaFoundationInitializer::HasInitialized()) {
    return InitPromise::CreateAndReject(
        MediaResult(NS_ERROR_DOM_MEDIA_FATAL_ERR,
                    RESULT_DETAIL("Can't create the MFT encoder.")),
        __func__);
  }

  RefPtr<MFTEncoder> encoder = new MFTEncoder(
      mHardwareNotAllowed ? MFTEncoder::HWPreference::SoftwareOnly
                          : MFTEncoder::HWPreference::PreferHardware);
  HRESULT hr;
  mscom::EnsureMTA([&]() { hr = InitMFTEncoder(encoder); });

  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("init MFTEncoder: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return InitPromise::CreateAndReject(
        MediaResult(NS_ERROR_DOM_MEDIA_FATAL_ERR,
                    RESULT_DETAIL("Can't create the MFT encoder.")),
        __func__);
  }

  mEncoder = std::move(encoder);
  InitializeConfigData();
  mIsHardwareAccelerated = mEncoder->IsHardwareAccelerated();
  WMF_ENC_LOGD("HW accelerated: %s", mIsHardwareAccelerated ? "yes" : "no");
  cleanup.release();
  return InitPromise::CreateAndResolve(true, __func__);
}

HRESULT WMFMediaDataEncoder::InitMFTEncoder(RefPtr<MFTEncoder>& aEncoder) {
  HRESULT hr = aEncoder->Create(CodecToSubtype(mConfig.mCodec), mConfig.mSize,
                                mConfig.mCodecSpecific);
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("MFTEncoder::Create: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return hr;
  }

  hr = aEncoder->SetModes(mConfig);
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("MFTEncoder::SetMode: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return hr;
  }

  hr = SetMediaTypes(aEncoder, mConfig);
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("MFTEncoder::SetMediaType: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return hr;
  }

  return S_OK;
}

void WMFMediaDataEncoder::InitializeConfigData() {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);

  if (mConfig.mCodec != CodecType::H264) {
    return;
  }

  auto r = mEncoder->GetMPEGSequenceHeader();
  if (r.isErr()) {
    WMF_ENC_LOGE("GetMPEGSequenceHeader failed");
    return;
  }

  nsTArray<UINT8> header = r.unwrap();
  SetConfigData(header);
}

void WMFMediaDataEncoder::SetConfigData(const nsTArray<UINT8>& aHeader) {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);

  if (mConfig.mCodec != CodecType::H264) {
    return;
  }

  mConfigData =
      aHeader.Length() > 0 ? ParseH264Parameters(aHeader, IsAnnexB()) : nullptr;
  WMF_ENC_LOGD("ConfigData has been updated to %zu bytes",
               mConfigData ? mConfigData->Length() : 0);
}

RefPtr<EncodePromise> WMFMediaDataEncoder::ProcessEncode(
    RefPtr<const VideoData>&& aSample) {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);
  MOZ_ASSERT(aSample);
  MOZ_ASSERT(mEncodePromise.IsEmpty());
  MOZ_ASSERT(!mEncodeRequest.Exists());

  WMF_ENC_LOGD("ProcessEncode ts=%s duration=%s",
               aSample->mTime.ToString().get(),
               aSample->mDuration.ToString().get());

  RefPtr<IMFSample> nv12 = ConvertToNV12InputSample(std::move(aSample));
  if (!nv12) {
    WMF_ENC_LOGE("failed to convert sample into NV12 format");
    return EncodePromise::CreateAndReject(
        MediaResult(NS_ERROR_DOM_MEDIA_FATAL_ERR,
                    RESULT_DETAIL("Failed to convert sample into NV12 format")),
        __func__);
  }

  RefPtr<EncodePromise> p = mEncodePromise.Ensure(__func__);

  nsTArray<MFTEncoder::InputSample> inputs;
  inputs.AppendElement(MFTEncoder::InputSample{
      .mSample = nv12.forget(), .mKeyFrameRequested = aSample->mKeyframe});
  mEncoder->Encode(std::move(inputs))
      ->Then(
          GetCurrentSerialEventTarget(), __func__,
          [self = RefPtr<WMFMediaDataEncoder>(this)](
              MFTEncoder::EncodedData&& aOutput) {
            self->mEncodeRequest.Complete();
            self->mEncodePromise.Resolve(
                self->ProcessOutputSamples(std::move(aOutput)), __func__);
          },
          [self =
               RefPtr<WMFMediaDataEncoder>(this)](const MediaResult& aError) {
            WMF_ENC_SLOGE("Encode failed: %s", aError.Description().get());
            self->mEncodeRequest.Complete();
            self->mEncodePromise.Reject(aError, __func__);
          })
      ->Track(mEncodeRequest);

  return p;
}

RefPtr<EncodePromise> WMFMediaDataEncoder::ProcessEncodeBatch(
    nsTArray<RefPtr<const VideoData>>&& aSamples) {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);
  MOZ_ASSERT(!aSamples.IsEmpty());
  MOZ_ASSERT(mEncodePromise.IsEmpty());
  MOZ_ASSERT(!mEncodeRequest.Exists());

  WMF_ENC_LOGD("ProcessEncodeBatch: num of samples=%zu", aSamples.Length());

  nsTArray<MFTEncoder::InputSample> inputs;
  for (auto& sample : aSamples) {
    RefPtr<IMFSample> nv12 = ConvertToNV12InputSample(std::move(sample));
    if (!nv12) {
      WMF_ENC_LOGE(
          "failed to convert samples(ts=%s duration=%s) into NV12 format",
          sample->mTime.ToString().get(), sample->mDuration.ToString().get());
      return EncodePromise::CreateAndReject(
          MediaResult(
              NS_ERROR_DOM_MEDIA_FATAL_ERR,
              RESULT_DETAIL("Failed to convert sample into NV12 format")),
          __func__);
    }
    inputs.AppendElement(MFTEncoder::InputSample{
        .mSample = std::move(nv12), .mKeyFrameRequested = sample->mKeyframe});
  }

  RefPtr<EncodePromise> p = mEncodePromise.Ensure(__func__);

  mEncoder->Encode(std::move(inputs))
      ->Then(
          GetCurrentSerialEventTarget(), __func__,
          [self = RefPtr<WMFMediaDataEncoder>(this)](
              MFTEncoder::EncodedData&& aOutput) {
            self->mEncodeRequest.Complete();
            self->mEncodePromise.Resolve(
                self->ProcessOutputSamples(std::move(aOutput)), __func__);
          },
          [self =
               RefPtr<WMFMediaDataEncoder>(this)](const MediaResult& aError) {
            WMF_ENC_SLOGE("Encode failed: %s", aError.Description().get());
            self->mEncodeRequest.Complete();
            self->mEncodePromise.Reject(aError, __func__);
          })
      ->Track(mEncodeRequest);

  return p;
}

RefPtr<EncodePromise> WMFMediaDataEncoder::ProcessDrain() {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);
  MOZ_ASSERT(mDrainPromise.IsEmpty());
  MOZ_ASSERT(!mDrainRequest.Exists());

  WMF_ENC_LOGD("ProcessDrain");

  RefPtr<EncodePromise> p = mDrainPromise.Ensure(__func__);

  mEncoder->Drain()
      ->Then(
          GetCurrentSerialEventTarget(), __func__,
          [self = RefPtr<WMFMediaDataEncoder>(this)](
              MFTEncoder::EncodedData&& aOutput) {
            self->mDrainRequest.Complete();
            self->mDrainPromise.Resolve(
                self->ProcessOutputSamples(std::move(aOutput)), __func__);
          },
          [self =
               RefPtr<WMFMediaDataEncoder>(this)](const MediaResult& aError) {
            WMF_ENC_SLOGE("Drain failed: %s", aError.Description().get());
            self->mDrainRequest.Complete();
            self->mDrainPromise.Reject(aError, __func__);
          })
      ->Track(mDrainRequest);

  return p;
}

already_AddRefed<IMFSample> WMFMediaDataEncoder::ConvertToNV12InputSample(
    RefPtr<const VideoData>&& aData) {
  AssertOnTaskQueue();
  MOZ_ASSERT(mEncoder);

  AUTO_MARKER("::ConvertToNV12InputSample");

  size_t mBufferLength = 0;

  const int32_t ySrtride = mConfig.mSize.width;
  const int32_t uvStride = ySrtride;

  const int32_t yHeight = mConfig.mSize.height;
  const int32_t uvHeight = yHeight / 2 + (yHeight % 2);

  CheckedInt<size_t> yLength(ySrtride);
  yLength *= yHeight;
  if (!yLength.isValid()) {
    WMF_ENC_LOGE("dest yLength overflows");
    return nullptr;
  }

  CheckedInt<size_t> uvLength(uvStride);
  uvLength *= uvHeight;
  if (!uvLength.isValid()) {
    WMF_ENC_LOGE("dest uvLength overflows");
    return nullptr;
  }

  CheckedInt<size_t> length(yLength);
  length += uvLength;
  if (!length.isValid()) {
    WMF_ENC_LOGE("dest length overflows");
    return nullptr;
  }
  mBufferLength = length.value();

  RefPtr<IMFSample> input;
  HRESULT hr = mEncoder->CreateInputSample(&input, mBufferLength);
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("CreateInputSample: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return nullptr;
  }

  RefPtr<IMFMediaBuffer> buffer;
  hr = input->GetBufferByIndex(0, getter_AddRefs(buffer));
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("GetBufferByIndex: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return nullptr;
  }

  hr = buffer->SetCurrentLength(mBufferLength);
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("SetCurrentLength: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return nullptr;
  }

  LockBuffer lockBuffer(buffer);
  hr = lockBuffer.Result();
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("LockBuffer: error = 0x%lX, %ls", hr, error.ErrorMessage());
    return nullptr;
  }

  nsresult rv = ConvertToNV12(aData->mImage, lockBuffer.Data(), ySrtride,
                              lockBuffer.Data() + yLength.value(), uvStride,
                              mConfig.mSize);
  if (NS_FAILED(rv)) {
    WMF_ENC_LOGE("Failed to convert to NV12");
    return nullptr;
  }

  hr = input->SetSampleTime(UsecsToHNs(aData->mTime.ToMicroseconds()));
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("SetSampleTime: error = 0x%lX, %ls", hr, error.ErrorMessage());
    return nullptr;
  }

  hr = input->SetSampleDuration(UsecsToHNs(aData->mDuration.ToMicroseconds()));
  if (FAILED(hr)) {
    _com_error error(hr);
    WMF_ENC_LOGE("SetSampleDuration: error = 0x%lX, %ls", hr,
                 error.ErrorMessage());
    return nullptr;
  }

  return input.forget();
}

MediaDataEncoder::EncodedData WMFMediaDataEncoder::ProcessOutputSamples(
    nsTArray<MFTEncoder::OutputSample>&& aSamples) {
  EncodedData frames;

  WMF_ENC_LOGD("ProcessOutputSamples: %zu frames", aSamples.Length());

  for (MFTEncoder::OutputSample& sample : aSamples) {
    RefPtr<MediaRawData> frame = OutputSampleToMediaData(sample);
    if (frame) {
      frames.AppendElement(std::move(frame));
    } else {
      WMF_ENC_LOGE("failed to convert output frame");
    }
  }
  return frames;
}

already_AddRefed<MediaRawData> WMFMediaDataEncoder::OutputSampleToMediaData(
    MFTEncoder::OutputSample& aSample) {
  AssertOnTaskQueue();
  MOZ_ASSERT(aSample.mSample);

  RefPtr<IMFMediaBuffer> buffer;
  HRESULT hr = aSample.mSample->GetBufferByIndex(0, getter_AddRefs(buffer));
  NS_ENSURE_TRUE(SUCCEEDED(hr), nullptr);

  LockBuffer lockBuffer(buffer);
  NS_ENSURE_TRUE(SUCCEEDED(lockBuffer.Result()), nullptr);

  LONGLONG time = 0;
  hr = aSample.mSample->GetSampleTime(&time);
  NS_ENSURE_TRUE(SUCCEEDED(hr), nullptr);

  LONGLONG duration = 0;
  hr = aSample.mSample->GetSampleDuration(&duration);
  NS_ENSURE_TRUE(SUCCEEDED(hr), nullptr);

  bool isKeyframe = MFGetAttributeUINT32(aSample.mSample,
                                         MFSampleExtension_CleanPoint, false);

  auto frame = MakeRefPtr<MediaRawData>();

  if (mConfig.mCodec == CodecType::H264 &&
      mConfig.mScalabilityMode != ScalabilityMode::None) {
    auto maybeId =
        H264::ExtractSVCTemporalId(lockBuffer.Data(), lockBuffer.Length());
    frame->mTemporalLayerId = Some(maybeId.unwrapOr(0));
  }

  if (!aSample.mHeader.IsEmpty()) {
    SetConfigData(aSample.mHeader);
  }

  if (!WriteFrameData(frame, lockBuffer, isKeyframe)) {
    return nullptr;
  }

  frame->mTime = media::TimeUnit::FromMicroseconds(HNsToUsecs(time));
  frame->mDuration = media::TimeUnit::FromMicroseconds(HNsToUsecs(duration));
  frame->mKeyframe = isKeyframe;

  WMF_ENC_LOGD("converted MediaData: ts=%s", frame->mTime.ToString().get());

  return frame.forget();
}

bool WMFMediaDataEncoder::IsAnnexB() const {
  MOZ_ASSERT(mConfig.mCodec == CodecType::H264);
  return mConfig.mCodecSpecific.as<H264Specific>().mFormat ==
         H264BitStreamFormat::ANNEXB;
}

bool WMFMediaDataEncoder::WriteFrameData(RefPtr<MediaRawData>& aDest,
                                         LockBuffer& aSrc, bool aIsKeyframe) {
  // From raw encoded data, write in avCC or AnnexB format depending on the
  // config.

  if (mConfig.mCodec == CodecType::H264) {
    size_t prependLength = 0;
    RefPtr<MediaByteBuffer> avccHeader;
    if (aIsKeyframe && mConfigData) {
      if (IsAnnexB()) {
        const nsTArray<NAL_TYPES> aTypes = {H264_NAL_SPS, H264_NAL_PPS};
        if (!AnnexB::FindAllNalTypes(
                Span<const uint8_t>(aSrc.Data(), aSrc.Length()), aTypes)) {
          prependLength = mConfigData->Length();
        }
      } else {
        avccHeader = mConfigData;
      }
    }

    UniquePtr<MediaRawDataWriter> writer(aDest->CreateWriter());
    if (!writer->SetSize(prependLength + aSrc.Length())) {
      WMF_ENC_LOGE("fail to allocate output buffer");
      return false;
    }

    if (prependLength > 0) {
      PodCopy(writer->Data(), mConfigData->Elements(), prependLength);
    }
    PodCopy(writer->Data() + prependLength, aSrc.Data(), aSrc.Length());

    if (!IsAnnexB() && !AnnexB::ConvertSampleToAVCC(aDest, avccHeader)) {
      WMF_ENC_LOGE("fail to convert annex-b sample to AVCC");
      return false;
    }

    return true;
  }
  UniquePtr<MediaRawDataWriter> writer(aDest->CreateWriter());
  if (!writer->SetSize(aSrc.Length())) {
    WMF_ENC_LOGE("fail to allocate output buffer");
    return false;
  }

  PodCopy(writer->Data(), aSrc.Data(), aSrc.Length());
  return true;
}

#undef AUTO_MARKER

}  // namespace mozilla