// 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/android/openxr_graphics_binding_open_gles.h"

#include <algorithm>
#include <vector>

#include "components/viz/common/resources/shared_image_format.h"
#include "device/vr/android/xr_image_transport_base.h"
#include "device/vr/openxr/openxr_api_wrapper.h"
#include "device/vr/openxr/openxr_platform.h"
#include "device/vr/openxr/openxr_util.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 "gpu/command_buffer/service/ahardwarebuffer_utils.h"
#include "gpu/ipc/common/android/android_hardware_buffer_utils.h"
#include "third_party/openxr/src/include/openxr/openxr.h"
#include "ui/gfx/buffer_types.h"
#include "ui/gfx/color_space.h"
#include "ui/gfx/gpu_fence.h"
#include "ui/gl/gl_context.h"
#include "ui/gl/gl_context_egl.h"
#include "ui/gl/gl_fence.h"
#include "ui/gl/gl_surface.h"
#include "ui/gl/gl_utils.h"
#include "ui/gl/init/gl_factory.h"
#include "ui/gl/scoped_egl_image.h"

namespace device {
namespace {
// Per the WebXR spec, the OS should treat the content generated by the page
// as sRGB, so we request/only support the sRGB format from the runtime.
constexpr int64_t kSwapchainFormat = GL_SRGB8_ALPHA8;

struct OpenGLESLayerData : public OpenXrCompositionLayer::GraphicsBindingData {
  OpenGLESLayerData() { type = kOpenGLES; }
  ~OpenGLESLayerData() override {
    if (back_buffer_fbo) {
      glDeleteFramebuffersEXT(1, &back_buffer_fbo);
    }
  }

  GLuint back_buffer_fbo = 0;
  bool using_shared_images = false;
};

OpenGLESLayerData& GetOpenGLESLayerData(OpenXrCompositionLayer& layer) {
  OpenGLESLayerData* data =
      static_cast<OpenGLESLayerData*>(layer.graphics_binding_data());
  // Layers are always created by OpenXrGraphicsBindingOpenGLES so this
  // should be safe.
  CHECK(data);
  CHECK(data->type == OpenXrCompositionLayer::GraphicsBindingData::kOpenGLES);
  return *data;
}

}  // namespace

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

OpenXrGraphicsBindingOpenGLES::OpenXrGraphicsBindingOpenGLES(
    const OpenXrExtensionEnumeration* extension_enum)
    : OpenXrGraphicsBinding(extension_enum) {}

OpenXrGraphicsBindingOpenGLES::~OpenXrGraphicsBindingOpenGLES() {
  if (overlay_texture_.id) {
    glDeleteTextures(1, &overlay_texture_.id);
  }
}

bool OpenXrGraphicsBindingOpenGLES::InitializeGl() {
  if (gl_initialized_) {
    return true;
  }

  // None of the other runtimes support ANGLE, so we disable it too for now.
  // TODO(crbug.com/529457533): Investigate if we can support ANGLE or if we'll
  // run into similar problems as cardboard.
  gl::DisableANGLE();

  // Everything below is a hacky first pass at making a session and likely needs
  // to be re-written with proper context/surfaces/types.
  gl::GLDisplay* display = nullptr;
  if (gl::GetGLImplementation() == gl::kGLImplementationNone) {
    display = gl::init::InitializeGLOneOff(
        /*gpu_preference=*/gl::GpuPreference::kDefault);
    if (!display) {
      DLOG(ERROR) << "gl::init::InitializeGLOneOff failed";
      return false;
    }
  }
  display = gl::GetDefaultDisplayEGL();

  DCHECK(gl::GetGLImplementation() != gl::kGLImplementationEGLANGLE);

  scoped_refptr<gl::GLSurface> surface;
  surface = gl::init::CreateOffscreenGLSurface(display, {0, 0});

  DVLOG(3) << "surface=" << surface.get();
  if (!surface.get()) {
    DLOG(ERROR) << "gl::init::CreateViewGLSurface failed";
    return false;
  }

  gl::GLContextAttribs context_attribs;
  // OpenXr's shared EGL context needs to be compatible with ours.
  // Any mismatches result in a EGL_BAD_MATCH error, including different reset
  // notification behavior according to
  // https://www.khronos.org/registry/EGL/specs/eglspec.1.5.pdf page 56.
  // Chromium defaults to lose context on reset when the robustness extension is
  // present, even if robustness features are not requested specifically.
  context_attribs.lose_context_on_reset = false;

  scoped_refptr<gl::GLContextEGL> egl_context = new gl::GLContextEGL(nullptr);
  scoped_refptr<gl::GLContext> context = gl::InitializeGLContext(
      egl_context.get(), surface.get(), context_attribs);
  if (!context.get()) {
    DLOG(ERROR) << "gl::init::CreateGLContext failed";
    return false;
  }
  if (!context->MakeCurrent(surface.get())) {
    DLOG(ERROR) << "gl::GLContext::MakeCurrent() failed";
    return false;
  }

  // Assign the surface and context members now that initialization has
  // succeeded.
  surface_ = std::move(surface);
  context_ = std::move(context);
  egl_context_ = std::move(egl_context);

  binding_.display = display->GetDisplay();
  binding_.config = (EGLConfig)0;
  binding_.context = egl_context_->GetHandle();

  renderer_ = std::make_unique<XrRenderer>();

  gl_initialized_ = true;
  return true;
}

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

  PFN_xrGetOpenGLESGraphicsRequirementsKHR get_graphics_requirements_fn{
      nullptr};
  if (XR_FAILED(xrGetInstanceProcAddr(
          instance, "xrGetOpenGLESGraphicsRequirementsKHR",
          (PFN_xrVoidFunction*)(&get_graphics_requirements_fn)))) {
    return false;
  }

  // TODO(crbug.com/529458212): Validate/set version based on the output here.
  XrGraphicsRequirementsOpenGLESKHR graphics_requirements = {
      XR_TYPE_GRAPHICS_REQUIREMENTS_OPENGL_ES_KHR};
  if (XR_FAILED(get_graphics_requirements_fn(instance, system,
                                             &graphics_requirements))) {
    return false;
  }

  initialized_ = true;
  return true;
}

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

int64_t OpenXrGraphicsBindingOpenGLES::GetSwapchainFormat(
    XrSession session) const {
  uint32_t format_length = 0;
  RETURN_IF_XR_FAILED(
      xrEnumerateSwapchainFormats(session, 0, &format_length, nullptr));
  std::vector<int64_t> swapchain_formats(format_length);
  RETURN_IF_XR_FAILED(
      xrEnumerateSwapchainFormats(session, (uint32_t)swapchain_formats.size(),
                                  &format_length, swapchain_formats.data()));
  DCHECK(!swapchain_formats.empty());

  // If the runtime doesn't support the swapchain format we need, log it. We'll
  // still return it anyway as that will cause an error with creating the
  // swapchain, which is better than arbitrarily returning a type that the
  // runtime supports, but we don't.
  if (!std::ranges::contains(swapchain_formats, kSwapchainFormat)) {
    LOG(ERROR) << "No matching supported swapchain formats with OpenXr Runtime";
  }

  return kSwapchainFormat;
}

XrResult OpenXrGraphicsBindingOpenGLES::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<XrSwapchainImageOpenGLESKHR> xr_color_swapchain_images(
      chain_length, {XR_TYPE_SWAPCHAIN_IMAGE_OPENGL_ES_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.image);
  }
  layer.SetSwapchainImages(std::move(color_swapchain_images));

  return XR_SUCCESS;
}

bool OpenXrGraphicsBindingOpenGLES::CanUseSharedImages() const {
  return true;
}

// This is more or less copied from XrImageTransportBase::ResizeSharedBuffer,
// with just the types changed as needed, and logic extracted out of the
// mailbox_to_surface_bridge.
void OpenXrGraphicsBindingOpenGLES::ResizeSharedBuffer(
    OpenXrCompositionLayer& layer,
    OpenXrSwapchainInfo& swap_chain_info,
    gpu::SharedImageInterface* sii) {
  CHECK(sii);
  const OpenGLESLayerData& open_gles_layer_data = GetOpenGLESLayerData(layer);
  auto transfer_size = layer.GetTransferSize();

  // TODO(crbug.com/459811463): We don't use GL_TEXTURE_CUBE_MAP because
  // AHARDWAREBUFFER_USAGE_GPU_CUBE_MAP is not always supported. So we put
  // 6 faces inside a 2D texture, 3 tiles per row, 2 rows in total.
  if (layer.type() == OpenXrCompositionLayer::Type::kCube) {
    transfer_size.set_width(transfer_size.width() * 3);
    transfer_size.set_height(transfer_size.height() * 2);
  }

  if (!open_gles_layer_data.using_shared_images ||
      swap_chain_info.shared_buffer_size == transfer_size) {
    return;
  }

  // Unbind previous image (if any).
  if (swap_chain_info.shared_image) {
    DVLOG(2) << ": DestroySharedImage, mailbox="
             << swap_chain_info.shared_image->mailbox().ToDebugString();
    // Note: The sync token in shared image may not be accurate. See comment
    // in XrImageTransportBase::TransferFrame.
    swap_chain_info.shared_image->UpdateDestructionSyncToken(
        swap_chain_info.sync_token);
    swap_chain_info.shared_image.reset();
  }

  // Remove reference to previous image (if any).
  swap_chain_info.local_eglimage.reset();

  static constexpr viz::SharedImageFormat format =
      viz::SinglePlaneFormat::kRGBA_8888;
  static constexpr gfx::BufferUsage usage = gfx::BufferUsage::SCANOUT;

  // The SharedImages created here will eventually be transferred to other
  // processes to have their contents written by WebGL and read via GL by
  // OpenXR.
  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 the XRSession is producing frames with WebGPU then the appropriate usage
  // also needs to be added.
  if (IsWebGPUSession()) {
    shared_image_usage |= gpu::SHARED_IMAGE_USAGE_WEBGPU_READ |
                          gpu::SHARED_IMAGE_USAGE_WEBGPU_WRITE;
  }

  swap_chain_info.scoped_ahb_handle =
      gpu::CreateScopedHardwareBufferHandle(transfer_size, format, usage);
  swap_chain_info.shared_buffer_size = transfer_size;

  // Create a GMB Handle from scoped_ahb_handle.
  gfx::GpuMemoryBufferHandle gmb_handle;
  gmb_handle.type = gfx::ANDROID_HARDWARE_BUFFER;
  gmb_handle.android_hardware_buffer =
      swap_chain_info.scoped_ahb_handle.Clone();

  // Though we've created the OpenXR swapchain to have an SRGB format (see
  // `kSwapchainFormat` above), all of the compositing that the page does is
  // actually done in the linear space. So set up our shared image to be linear.
  // Note that this matches the behavior in the D3D11 graphics binding as well.
  swap_chain_info.shared_image = sii->CreateSharedImage(
      {viz::SinglePlaneFormat::kRGBA_8888, transfer_size,
       gfx::ColorSpace(gfx::ColorSpace::PrimaryID::BT709,
                       gfx::ColorSpace::TransferID::LINEAR),
       shared_image_usage, "OpenXrGraphicsBinding"},
      std::move(gmb_handle));
  CHECK(swap_chain_info.shared_image);
  swap_chain_info.sync_token =
      swap_chain_info.shared_image->creation_sync_token();
  sii->VerifySyncToken(swap_chain_info.sync_token);
  DCHECK_EQ(swap_chain_info.shared_image->GetTextureTarget(),
            static_cast<uint32_t>(GL_TEXTURE_2D));

  DVLOG(2) << ": CreateSharedImage, mailbox="
           << swap_chain_info.shared_image->mailbox().ToDebugString()
           << ", SyncToken=" << swap_chain_info.sync_token.ToDebugString();

  // Create an EGLImage for the buffer.
  auto egl_image = gpu::CreateEGLImageFromAHardwareBuffer(
      swap_chain_info.scoped_ahb_handle.get());
  if (!egl_image.is_valid()) {
    DLOG(WARNING) << __func__ << ": ERROR: failed to initialize image!";
    return;
  }

  swap_chain_info.shared_buffer_texture.target = GL_TEXTURE_2D;
  glGenTextures(1, &swap_chain_info.shared_buffer_texture.id);
  glBindTexture(swap_chain_info.shared_buffer_texture.target,
                swap_chain_info.shared_buffer_texture.id);
  glEGLImageTargetTexture2DOES(swap_chain_info.shared_buffer_texture.target,
                               egl_image.get());
  swap_chain_info.local_eglimage = std::move(egl_image);
}

void OpenXrGraphicsBindingOpenGLES::CreateSharedImages(
    OpenXrCompositionLayer& layer,
    gpu::SharedImageInterface* sii) {
  CHECK(sii);
  OpenGLESLayerData& open_gles_layer_data = GetOpenGLESLayerData(layer);
  open_gles_layer_data.using_shared_images = true;
  for (auto& swap_chain_info : layer.GetSwapchainImages()) {
    ResizeSharedBuffer(layer, swap_chain_info, sii);
  }
}

bool OpenXrGraphicsBindingOpenGLES::RenderLayer(
    OpenXrCompositionLayer& layer,
    const scoped_refptr<viz::ContextProvider>& context_provider) {
  OpenXrSwapchainInfo* swap_chain_info = layer.GetActiveSwapchainImage();
  if (!swap_chain_info) {
    return false;
  }

  gfx::Size swapchain_image_size = layer.GetSwapchainImageSize();

  OpenGLESLayerData& open_gles_layer_data = GetOpenGLESLayerData(layer);
  if (!open_gles_layer_data.back_buffer_fbo) {
    glGenFramebuffersEXT(1, &open_gles_layer_data.back_buffer_fbo);
  }
  glBindFramebufferEXT(GL_FRAMEBUFFER, open_gles_layer_data.back_buffer_fbo);

  glDisable(GL_CULL_FACE);
  glDisable(GL_SCISSOR_TEST);
  glDisable(GL_POLYGON_OFFSET_FILL);
  glDisable(GL_BLEND);

  // TODO(https://crbug.com/324596270): This shouldn't be necessary, but we
  // can't seem to get the image set up to be treated as linear any other way.
  glDisable(GL_FRAMEBUFFER_SRGB_EXT);
  glViewport(0, 0, swapchain_image_size.width(), swapchain_image_size.height());

  gfx::Transform transform;
  float transform_floats[16];
  transform.GetColMajorF(transform_floats);

  if (layer.type() == OpenXrCompositionLayer::Type::kCube) {
    if (webxr_visible_) {
      // We call glFramebufferTexture2DEXT inside DrawCubemap for 6 faces.
      renderer_->DrawCubemap(swap_chain_info->shared_buffer_texture,
                             swap_chain_info->openxr_texture, transform_floats,
                             layer.mutable_data().opacity);
    }
  } else {
    glFramebufferTexture2DEXT(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
                              GL_TEXTURE_2D, swap_chain_info->openxr_texture,
                              0);
    if (webxr_visible_) {
      renderer_->Draw(swap_chain_info->shared_buffer_texture, transform_floats,
                      layer.mutable_data().opacity);
    }
  }

  // The overlay is rendred to the base layer, which has an invalid layer id.
  if (overlay_visible_ && layer.GetLayerId() == kInvalidLayerId) {
    glEnable(GL_FRAMEBUFFER_SRGB_EXT);
    if (webxr_visible_) {
      glEnable(GL_BLEND);
      glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
    }

    auto egl_image = gpu::CreateEGLImageFromAHardwareBuffer(
        overlay_handle_.android_hardware_buffer.get());
    if (!egl_image.is_valid()) {
      return false;
    }

    if (!overlay_texture_.id) {
      overlay_texture_.target = GL_TEXTURE_2D;
      glGenTextures(1, &overlay_texture_.id);
    }

    // If the image is going to be flipped during compositing then the overlay
    // needs to be rendered flipped as well.
    if (ShouldFlipSubmittedImage(layer)) {
      transform.Translate(0, 1);
      transform.Scale(1, -1);
      transform.GetColMajorF(transform_floats);
    }

    glBindTexture(overlay_texture_.target, overlay_texture_.id);
    glEGLImageTargetTexture2DOES(overlay_texture_.target, egl_image.get());
    renderer_->Draw(overlay_texture_, transform_floats);
  }

  return true;
}

void OpenXrGraphicsBindingOpenGLES::CleanupWithoutSubmit() {
  // Nothing to do, we store all of our necessary data on the active swapchain
  // image, which gets recycled when the frame ends.
}

bool OpenXrGraphicsBindingOpenGLES::WaitOnFence(OpenXrCompositionLayer& layer,
                                                gfx::GpuFence& gpu_fence) {
  std::unique_ptr<gl::GLFence> local_fence =
      gl::GLFence::CreateFromGpuFence(gpu_fence);
  local_fence->ServerWait();

  return true;
}

bool OpenXrGraphicsBindingOpenGLES::ShouldFlipSubmittedImage(
    OpenXrCompositionLayer& layer) const {
  bool should_flip = layer.flip_y();
  // WebGPU produces textures that are y-flipped relative to WebGL, which needs
  // to be accounted for during frame submission.
  if (IsWebGPUSession()) {
    should_flip = !should_flip;
  }
  return should_flip;
}

void OpenXrGraphicsBindingOpenGLES::OnSwapchainImageReady(
    OpenXrCompositionLayer& layer,
    gpu::SharedImageInterface* sii) {
  OpenXrSwapchainInfo* swap_chain_info = layer.GetActiveSwapchainImage();
  CHECK(swap_chain_info);
  ResizeSharedBuffer(layer, *swap_chain_info, sii);
}

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

  CHECK(texture.type == gfx::ANDROID_HARDWARE_BUFFER);
  overlay_handle_ = std::move(texture);
  return true;
}

gfx::Size OpenXrGraphicsBindingOpenGLES::GetMaxTextureSize() {
  CHECK(gl_initialized_);
  int max_texture_size;
  glGetIntegerv(GL_MAX_TEXTURE_SIZE, &max_texture_size);
  DVLOG(1) << __func__ << " Max size=" << max_texture_size;
  return {max_texture_size, max_texture_size};
}

bool OpenXrGraphicsBindingOpenGLES::SupportsLayers() const {
  return true;
}

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

}  // namespace device