/* 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/. */

#import <AVFoundation/AVFoundation.h>
#import <QuartzCore/QuartzCore.h>

#include "gtest/gtest.h"
#include "mozilla/ScopeExit.h"
#include "mozilla/TypedEnumBits.h"
#include "mozilla/gfx/MacIOSurface.h"
#include "mozilla/layers/NativeLayerCA.h"

using namespace mozilla;
using namespace mozilla::gfx;
using namespace mozilla::layers;

static const IntSize kVideoSize(64, 32);

// Properties of a presented frame.
enum class PresentFlags : uint8_t {
  None = 0,
  DRM = 1 << 0,
  HDR = 1 << 1,
};
MOZ_MAKE_ENUM_CLASS_BITWISE_OPERATORS(PresentFlags)

// A single surface-presentation layer committed into a detached root CALayer,
// either as the onscreen tree or as the offscreen tree used for snapshots.
class NativeLayerCAVideo : public ::testing::TestWithParam<bool> {
 protected:
  void SetUp() override {
    mRootCALayer = [[CALayer layer] retain];
    mOffscreenRootCALayer = [[CALayer layer] retain];
    mRoot = NativeLayerRootCA::CreateForCALayer(mRootCALayer);
    mLayer = mRoot->CreateLayerForSurfacePresentation(kVideoSize, true);
    mLayer->SetDisplayRect(IntRect(IntPoint(), kVideoSize));
    mRoot->AppendLayer(mLayer);
  }

  void TearDown() override {
    mRoot->RemoveLayer(mLayer);
    mLayer = nullptr;
    mRoot = nullptr;
    [mOffscreenRootCALayer release];
    [mRootCALayer release];
  }

  bool IsOffscreen() const { return GetParam(); }

  void Present(const RefPtr<MacIOSurface>& aSurface, PresentFlags aFlags) {
    ASSERT_TRUE(aSurface);
    IntSize size = kVideoSize;
    mLayer->SetSurfaceToPresent(aSurface->GetIOSurfaceRef(), size,
                                bool(aFlags & PresentFlags::DRM),
                                bool(aFlags & PresentFlags::HDR));
    if (IsOffscreen()) {
      mRoot->CommitOffscreen(mOffscreenRootCALayer);
    } else {
      mRoot->CommitToScreen();
    }
  }

  AVSampleBufferDisplayLayer* VideoLayer() {
    return FindVideoLayer(IsOffscreen() ? mOffscreenRootCALayer : mRootCALayer);
  }

  static RefPtr<MacIOSurface> CreateNV12Surface() {
    return MacIOSurface::CreateBiPlanarSurface(
        kVideoSize, IntSize(kVideoSize.width / 2, kVideoSize.height / 2),
        ChromaSubsampling::HALF_WIDTH_AND_HEIGHT, YUVColorSpace::BT709,
        TransferFunction::BT709, ColorRange::LIMITED, ColorDepth::COLOR_8,
        MacIOSurface::AllowAlpha::No);
  }

  static RefPtr<MacIOSurface> CreateBGRASurface() {
    return MacIOSurface::CreateIOSurface(kVideoSize.width, kVideoSize.height,
                                         MacIOSurface::AllowAlpha::No);
  }

  static AVSampleBufferDisplayLayer* FindVideoLayer(CALayer* aLayer) {
    if ([aLayer isKindOfClass:[AVSampleBufferDisplayLayer class]]) {
      return static_cast<AVSampleBufferDisplayLayer*>(aLayer);
    }
    for (CALayer* sublayer in aLayer.sublayers) {
      if (AVSampleBufferDisplayLayer* videoLayer = FindVideoLayer(sublayer)) {
        return videoLayer;
      }
    }
    return nil;
  }

  CALayer* mRootCALayer = nil;
  CALayer* mOffscreenRootCALayer = nil;
  RefPtr<NativeLayerRootCA> mRoot;
  RefPtr<NativeLayerCA> mLayer;
};

TEST_P(NativeLayerCAVideo, DRMVideoLayerIsCaptureProtected) {
  // Present a DRM video frame; the video layer must block capture and hold
  // no frame in its contents.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::DRM));
  ASSERT_NE(VideoLayer(), nil);
  EXPECT_TRUE(VideoLayer().preventsCapture);
  EXPECT_EQ(VideoLayer().contents, nil);
}

TEST_P(NativeLayerCAVideo, RebuiltDRMVideoLayerIsCaptureProtected) {
  // Present a DRM video frame; DRM video always uses a video layer, so one is
  // built.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::DRM));
  ASSERT_NE(VideoLayer(), nil);

  // Present a non-video surface with DRM still set, so the video layer is torn
  // down by the layer type changing alone, without a DRM state change.
  ASSERT_NO_FATAL_FAILURE(Present(CreateBGRASurface(), PresentFlags::DRM));
  ASSERT_EQ(VideoLayer(), nil);

  // Present a DRM video frame again; the rebuilt video layer must block
  // capture and hold no frame in its contents.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::DRM));
  ASSERT_NE(VideoLayer(), nil);
  EXPECT_TRUE(VideoLayer().preventsCapture);
  EXPECT_EQ(VideoLayer().contents, nil);
}

TEST_P(NativeLayerCAVideo, NonDRMVideoLayerShowsFirstFrame) {
  // Present a non-DRM HDR frame, which also uses a video layer; the layer
  // must show the first frame through its contents and allow capture.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::HDR));
  ASSERT_NE(VideoLayer(), nil);
  EXPECT_FALSE(VideoLayer().preventsCapture);
  EXPECT_NE(VideoLayer().contents, nil);
}

TEST_P(NativeLayerCAVideo, SwitchingToDRMDropsNonDRMFrame) {
  // Present a non-DRM HDR frame; the video layer shows it through its
  // contents.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::HDR));
  ASSERT_NE(VideoLayer(), nil);
  ASSERT_NE(VideoLayer().contents, nil);

  // Present a DRM HDR frame on the same layer; the video layer must block
  // capture and hold no frame in its contents.
  ASSERT_NO_FATAL_FAILURE(
      Present(CreateNV12Surface(), PresentFlags::DRM | PresentFlags::HDR));
  ASSERT_NE(VideoLayer(), nil);
  EXPECT_TRUE(VideoLayer().preventsCapture);
  EXPECT_EQ(VideoLayer().contents, nil);
}

TEST_P(NativeLayerCAVideo, SwitchingFromDRMUsesNewLayer) {
  // Present a DRM HDR frame; the video layer blocks capture. Keep it alive to
  // inspect it after it is replaced.
  ASSERT_NO_FATAL_FAILURE(
      Present(CreateNV12Surface(), PresentFlags::DRM | PresentFlags::HDR));
  AVSampleBufferDisplayLayer* drmLayer = [VideoLayer() retain];
  auto releaseDRMLayer = MakeScopeExit([&] { [drmLayer release]; });
  ASSERT_NE(drmLayer, nil);
  ASSERT_TRUE(drmLayer.preventsCapture);

  // Present a non-DRM HDR frame on the same layer; it must go to a new video
  // layer that shows its first frame, while the layer that held DRM frames
  // stays capture-protected.
  ASSERT_NO_FATAL_FAILURE(Present(CreateNV12Surface(), PresentFlags::HDR));
  ASSERT_NE(VideoLayer(), nil);
  EXPECT_NE(VideoLayer(), drmLayer);
  EXPECT_FALSE(VideoLayer().preventsCapture);
  EXPECT_NE(VideoLayer().contents, nil);
  EXPECT_TRUE(drmLayer.preventsCapture);
}

INSTANTIATE_TEST_SUITE_P(, NativeLayerCAVideo, ::testing::Bool(),
                         [](const ::testing::TestParamInfo<bool>& aInfo) {
                           return aInfo.param ? "Offscreen" : "Onscreen";
                         });