// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "device/vr/openxr/windows/openxr_graphics_binding_d3d11.h"

#include <d3d11_4.h>
#include <wrl.h>

#include "base/check.h"
#include "base/feature_list.h"
#include "base/logging.h"
#include "base/notreached.h"
#include "base/win/scoped_handle.h"
#include "device/vr/openxr/openxr_api_wrapper.h"
#include "device/vr/openxr/openxr_platform.h"
#include "device/vr/openxr/openxr_util.h"
#include "device/vr/openxr/openxr_view_configuration.h"
#include "device/vr/openxr/windows/openxr_platform_helper_windows.h"
#include "device/vr/public/cpp/features.h"
#include "device/vr/windows/d3d11_texture_helper.h"
#include "gpu/GLES2/gl2extchromium.h"
#include "gpu/command_buffer/client/client_shared_image.h"
#include "gpu/command_buffer/client/shared_image_interface.h"
#include "gpu/command_buffer/common/shared_image_usage.h"
#include "third_party/openxr/src/include/openxr/openxr.h"
#include "ui/gfx/gpu_fence.h"
#include "ui/gfx/gpu_memory_buffer_handle.h"

namespace device {

namespace {

struct D3DLayerData : public OpenXrCompositionLayer::GraphicsBindingData {
  D3DLayerData() { type = kD3D; }
};

}  // namespace

// static
void OpenXrGraphicsBinding::GetRequiredExtensions(
    std::vector<const char*>& extensions) {
  extensions.push_back(XR_KHR_D3D11_ENABLE_EXTENSION_NAME);
}

OpenXrGraphicsBindingD3D11::OpenXrGraphicsBindingD3D11(
    base::WeakPtr<OpenXrPlatformHelperWindows> weak_platform_helper)
    : OpenXrGraphicsBinding(weak_platform_helper->GetExtensionEnumeration()),
      texture_helper_(std::make_unique<D3D11TextureHelper>()),
      weak_platform_helper_(weak_platform_helper) {}

OpenXrGraphicsBindingD3D11::~OpenXrGraphicsBindingD3D11() = default;

bool OpenXrGraphicsBindingD3D11::Initialize(XrInstance instance,
                                            XrSystemId system) {
  if (initialized_) {
    return true;
  }

  if (!texture_helper_) {
    DVLOG(1) << __func__ << " No TextureHelper";
    return false;
  }

  if (!weak_platform_helper_) {
    DVLOG(1) << __func__ << " WeakPtr failed to resolve";
    return false;
  }

  LUID luid;
  if (!weak_platform_helper_->TryGetLuid(&luid, system)) {
    DVLOG(1) << __func__ << " Did not get a luid";
    return false;
  }

  texture_helper_->SetUseBGRA(true);
  if (!texture_helper_->SetAdapterLUID(luid) ||
      !texture_helper_->EnsureInitialized()) {
    DVLOG(1) << __func__ << " Texture helper initialization failed";
    return false;
  }

  binding_.device = texture_helper_->GetDevice().Get();
  initialized_ = true;
  return true;
}

const void* OpenXrGraphicsBindingD3D11::GetSessionCreateInfo() const {
  CHECK(initialized_);
  return &binding_;
}

int64_t OpenXrGraphicsBindingD3D11::GetSwapchainFormat(
    XrSession session) const {
  // OpenXR's swapchain format expects to describe the texture content.
  // The result of a swapchain image created from OpenXR API always contains a
  // typeless texture. On the other hand, WebGL API uses CSS color convention
  // that's sRGB. The RGBA typeless texture from OpenXR swapchain image leads to
  // a linear format render target view (reference to function
  // D3D11TextureHelper::EnsureRenderTargetView in d3d11_texture_helper.cc).
  // Therefore, the content in this openxr swapchain image is in sRGB format.
  return DXGI_FORMAT_R8G8B8A8_UNORM_SRGB;
}

XrResult OpenXrGraphicsBindingD3D11::EnumerateSwapchainImages(
    OpenXrCompositionLayer& layer) {
  CHECK(layer.HasColorSwapchain());
  CHECK(layer.GetSwapchainImages().empty());

  uint32_t chain_length;
  RETURN_IF_XR_FAILED(xrEnumerateSwapchainImages(layer.color_swapchain(), 0,
                                                 &chain_length, nullptr));
  std::vector<XrSwapchainImageD3D11KHR> xr_color_swapchain_images(
      chain_length, {XR_TYPE_SWAPCHAIN_IMAGE_D3D11_KHR});

  RETURN_IF_XR_FAILED(xrEnumerateSwapchainImages(
      layer.color_swapchain(), xr_color_swapchain_images.size(), &chain_length,
      reinterpret_cast<XrSwapchainImageBaseHeader*>(
          xr_color_swapchain_images.data())));

  std::vector<OpenXrSwapchainInfo> color_swapchain_images;
  color_swapchain_images.reserve(xr_color_swapchain_images.size());
  for (const auto& swapchain_image : xr_color_swapchain_images) {
    color_swapchain_images.emplace_back(swapchain_image.texture);
  }
  layer.SetSwapchainImages(std::move(color_swapchain_images));

  return XR_SUCCESS;
}

bool OpenXrGraphicsBindingD3D11::CanUseSharedImages() const {
  // TODO(https://crbug.com/458430816): Investigate using shared images outside
  // of WebGPU sessions.
  return IsWebGPUSession();
}

void OpenXrGraphicsBindingD3D11::CreateSharedImages(
    OpenXrCompositionLayer& layer,
    gpu::SharedImageInterface* sii) {
  CHECK(sii);
  for (auto& swap_chain_info : layer.GetSwapchainImages()) {
    Microsoft::WRL::ComPtr<IDXGIResource1> dxgi_resource;
    HRESULT hr = swap_chain_info.d3d11_texture->QueryInterface(
        IID_PPV_ARGS(&dxgi_resource));
    CHECK_EQ(hr, S_OK);

    Microsoft::WRL::ComPtr<ID3D11Texture2D> d3d11_texture;
    hr = dxgi_resource.As(&d3d11_texture);
    if (FAILED(hr)) {
      DLOG(ERROR) << "QueryInterface for ID3D11Texture2D failed with error "
                  << std::hex << hr;
      return;
    }

    D3D11_TEXTURE2D_DESC texture2d_desc;
    d3d11_texture->GetDesc(&texture2d_desc);

    // Shared handle creation can fail if the OpenXR runtime allocates textures
    // without the necessary sharing flags (e.g., D3D11_RESOURCE_MISC_SHARED),
    // which can happen depending on the runtime driver or hardware
    // configuration. In such cases, we cannot wrap the runtime's texture
    // directly in a GPU SharedImage. To fallback gracefully, we create our own
    // shareable D3D11 texture with NT handle and keyed mutex flags, create a
    // shared handle from it, and we will copy the runtime's texture content
    // into this shareable texture before submission.
    HANDLE shared_handle;
    hr = dxgi_resource->CreateSharedHandle(
        nullptr, DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE,
        nullptr, &shared_handle);

    if (FAILED(hr)) {
      DLOG(WARNING) << "Unable to create shared handle for DXGIResource (0x"
                    << std::hex << hr << "). Creating a separate shareable "
                    << "texture instead.";

      D3D11_TEXTURE2D_DESC desc;
      desc.Width = texture2d_desc.Width;
      desc.Height = texture2d_desc.Height;
      desc.MipLevels = 1;
      desc.ArraySize = 1;
      desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
      desc.SampleDesc.Count = 1;
      desc.SampleDesc.Quality = 0;
      desc.Usage = D3D11_USAGE_DEFAULT;
      desc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_RENDER_TARGET;
      desc.CPUAccessFlags = 0;
      desc.MiscFlags = D3D11_RESOURCE_MISC_SHARED_NTHANDLE |
                       D3D11_RESOURCE_MISC_SHARED_KEYEDMUTEX;

      Microsoft::WRL::ComPtr<ID3D11Device> d3d11_device =
          texture_helper_->GetDevice();
      hr = d3d11_device->CreateTexture2D(&desc, nullptr,
                                         &swap_chain_info.d3d11_shared_texture);
      if (FAILED(hr)) {
        DLOG(ERROR) << "Failed to create shared texture (0x" << std::hex << hr
                    << ")";
        return;
      }

      hr = swap_chain_info.d3d11_shared_texture->QueryInterface(
          IID_PPV_ARGS(&dxgi_resource));
      CHECK_EQ(hr, S_OK);

      hr = dxgi_resource->CreateSharedHandle(
          nullptr, DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE,
          nullptr, &shared_handle);
      if (FAILED(hr)) {
        DLOG(ERROR)
            << "Unable to create shared handle for fallback DXGIResource (0x"
            << std::hex << hr << ").";
        swap_chain_info.d3d11_shared_texture = nullptr;
        return;
      }
    }

    gfx::GpuMemoryBufferHandle gpu_memory_buffer_handle{
        gfx::DXGIHandle(base::win::ScopedHandle(shared_handle))};

    // TODO(crbug.com/529457611): This size is the size of the texture
    // from the OpenXr runtime, which is fine but does not work properly if the
    // page requests any kind of framebuffer scaling, because then the image
    // size that the page uses would be different than this size, which can
    // cause errors in rendering.
    gfx::Size buffer_size =
        gfx::Size(texture2d_desc.Width, texture2d_desc.Height);

    // The SharedImages created here will eventually be transferred to other
    // processes to have their contents written by WebGL or WebGPU.
    // Readback by OpenXR is always via OpenGL.
    gpu::SharedImageUsageSet shared_image_usage =
        gpu::SHARED_IMAGE_USAGE_SCANOUT | gpu::SHARED_IMAGE_USAGE_DISPLAY_READ |
        gpu::SHARED_IMAGE_USAGE_GLES2_READ |
        gpu::SHARED_IMAGE_USAGE_GLES2_WRITE;

    if (layer.read_only_data().needs_raster_access) {
      shared_image_usage |= gpu::SHARED_IMAGE_USAGE_RASTER_READ |
                            gpu::SHARED_IMAGE_USAGE_RASTER_WRITE;
    }

    if (IsWebGPUSession()) {
      shared_image_usage |= gpu::SHARED_IMAGE_USAGE_WEBGPU_READ |
                            gpu::SHARED_IMAGE_USAGE_WEBGPU_WRITE;
    }

    swap_chain_info.shared_image = sii->CreateSharedImage(
        {viz::SinglePlaneFormat::kRGBA_8888, buffer_size,
         gfx::ColorSpace(gfx::ColorSpace::PrimaryID::BT709,
                         gfx::ColorSpace::TransferID::LINEAR),
         shared_image_usage, "OpenXrSwapChain"},
        std::move(gpu_memory_buffer_handle));
    CHECK(swap_chain_info.shared_image);
    swap_chain_info.sync_token = sii->GenVerifiedSyncToken();
  }
}

bool OpenXrGraphicsBindingD3D11::WaitOnFence(OpenXrCompositionLayer& layer,
                                             gfx::GpuFence& gpu_fence) {
  OpenXrSwapchainInfo* swapchain_info = layer.GetActiveSwapchainImage();
  if (!swapchain_info) {
    return false;
  }

  Microsoft::WRL::ComPtr<ID3D11Device> d3d11_device =
      texture_helper_->GetDevice();
  Microsoft::WRL::ComPtr<ID3D11Device5> d3d11_device5;
  HRESULT hr = d3d11_device.As(&d3d11_device5);
  if (FAILED(hr)) {
    DLOG(ERROR) << "Unable to retrieve ID3D11Device5 interface " << std::hex
                << hr;
    return false;
  }

  Microsoft::WRL::ComPtr<ID3D11Fence> d3d11_fence;
  hr = d3d11_device5->OpenSharedFence(gpu_fence.GetGpuFenceHandle().Peek(),
                                      IID_PPV_ARGS(&d3d11_fence));
  if (FAILED(hr)) {
    DLOG(ERROR) << "Unable to open a shared fence " << std::hex << hr;
    return false;
  }

  Microsoft::WRL::ComPtr<ID3D11DeviceContext> d3d11_device_context;
  d3d11_device5->GetImmediateContext(&d3d11_device_context);

  Microsoft::WRL::ComPtr<ID3D11DeviceContext4> d3d11_device_context4;
  hr = d3d11_device_context.As(&d3d11_device_context4);
  if (FAILED(hr)) {
    DLOG(ERROR) << "Unable to retrieve ID3D11DeviceContext4 interface "
                << std::hex << hr;
    return false;
  }

  hr = d3d11_device_context4->Wait(d3d11_fence.Get(), 1);
  if (FAILED(hr)) {
    DLOG(ERROR) << "Unable to Wait on D3D11 fence " << std::hex << hr;
    return false;
  }

  // In order for the fence to be respected by the system, it needs to stick
  // around until the next time the texture comes up for use.
  swapchain_info->d3d11_fence = std::move(d3d11_fence);

  return true;
}

bool OpenXrGraphicsBindingD3D11::RenderLayer(
    OpenXrCompositionLayer& layer,
    const scoped_refptr<viz::ContextProvider>& context_provider) {
  CHECK(texture_helper_);
  const OpenXrSwapchainInfo* swapchain_info = layer.GetActiveSwapchainImage();
  if (!swapchain_info) {
    return false;
  }
  if (swapchain_info->d3d11_shared_texture) {
    if (!texture_helper_->CopyToBackBuffer(
            context_provider, swapchain_info->d3d11_shared_texture)) {
      DLOG(ERROR) << "CopyToBackBuffer failed.";
      return false;
    }
  }

  // Even if a shared image was copied, always perform a composite to account
  // for any necessary overlays.
  return texture_helper_->UpdateBackbufferSizes() &&
         texture_helper_->CompositeToBackBuffer(context_provider);
}

void OpenXrGraphicsBindingD3D11::CleanupWithoutSubmit() {
  texture_helper_->CleanupNoSubmit();
}

bool OpenXrGraphicsBindingD3D11::ShouldFlipSubmittedImage(
    OpenXrCompositionLayer& layer) const {
  return layer.IsUsingSharedImages() && !IsWebGPUSession();
}

void OpenXrGraphicsBindingD3D11::OnSwapchainImageSizeChanged(
    OpenXrCompositionLayer& layer) {
  texture_helper_->SetDefaultSize(layer.GetSwapchainImageSize());
}

void OpenXrGraphicsBindingD3D11::OnSwapchainImageReady(
    OpenXrCompositionLayer& layer,
    gpu::SharedImageInterface* sii) {
  const OpenXrSwapchainInfo* swapchain_info = layer.GetActiveSwapchainImage();
  CHECK(swapchain_info);

  texture_helper_->SetBackbuffer(swapchain_info->d3d11_texture.get(),
                                 layer.IsUsingSharedImages(),
                                 ShouldFlipSubmittedImage(layer));
}

void OpenXrGraphicsBindingD3D11::OnSetOverlayAndWebXrVisibility() {
  texture_helper_->SetSourceAndOverlayVisible(webxr_visible_, overlay_visible_);
}

void OpenXrGraphicsBindingD3D11::SetWebXrTexture(
    mojo::PlatformHandle texture_handle,
    const gpu::SyncToken& sync_token,
    const gfx::RectF& left,
    const gfx::RectF& right) {
  base::win::ScopedHandle scoped_handle = texture_handle.is_valid()
                                              ? texture_handle.TakeHandle()
                                              : base::win::ScopedHandle();

  texture_helper_->SetSourceTexture(std::move(scoped_handle), sync_token, left,
                                    right);
}

bool OpenXrGraphicsBindingD3D11::SetOverlayTexture(
    gfx::GpuMemoryBufferHandle texture,
    const gpu::SyncToken& sync_token,
    const gfx::RectF& left,
    const gfx::RectF& right) {
  if (texture.is_null()) {
    return false;
  }

  return texture_helper_->SetOverlayTexture(
      std::move(texture).dxgi_handle().TakeBufferHandle(), sync_token, left,
      right);
}

gfx::Size OpenXrGraphicsBindingD3D11::GetMaxTextureSize() {
  return {D3D11_REQ_TEXTURE2D_U_OR_V_DIMENSION,
          D3D11_REQ_TEXTURE2D_U_OR_V_DIMENSION};
}

void OpenXrGraphicsBindingD3D11::ResizeSharedBuffer(
    OpenXrCompositionLayer&,
    OpenXrSwapchainInfo& swap_chain_info,
    gpu::SharedImageInterface* sii) {
  // TODO(crbug.com/529457611): Current texture size needs to be updated.
}

bool OpenXrGraphicsBindingD3D11::SupportsLayers() const {
  return false;
}

std::unique_ptr<OpenXrCompositionLayer::GraphicsBindingData>
OpenXrGraphicsBindingD3D11::CreateLayerGraphicsBindingData() const {
  return std::make_unique<D3DLayerData>();
}

}  // namespace device