/*
 * Copyright 2023 Rive
 */

#include "rive/renderer/vulkan/vulkan_context.hpp"

#include "rive/rive_types.hpp"
#include <vk_mem_alloc.h>

namespace rive::gpu
{
// Returns null if the driver can't give us an allocator, leaving it to the
// caller to decide whether that is recoverable.
static VmaAllocator make_vma_allocator(
    VkInstance instance,
    VkPhysicalDevice physicalDevice,
    VkDevice device,
    uint32_t apiVersion,
    PFN_vkGetInstanceProcAddr pfnvkGetInstanceProcAddr)
{
  VmaVulkanFunctions vmaVulkanFunctions = {
      .vkGetInstanceProcAddr = pfnvkGetInstanceProcAddr,
      .vkGetDeviceProcAddr = reinterpret_cast<PFN_vkGetDeviceProcAddr>(
          pfnvkGetInstanceProcAddr(instance, "vkGetDeviceProcAddr")),
      .vkGetPhysicalDeviceProperties = reinterpret_cast<
          PFN_vkGetPhysicalDeviceProperties>(
          pfnvkGetInstanceProcAddr(instance, "vkGetPhysicalDeviceProperties")),
  };

  VmaAllocatorCreateInfo vmaCreateInfo = {
      // We are single-threaded.
      .flags = VMA_ALLOCATOR_CREATE_EXTERNALLY_SYNCHRONIZED_BIT,
      .physicalDevice = physicalDevice,
      .device = device,
      .pVulkanFunctions = &vmaVulkanFunctions,
      .instance = instance,
      .vulkanApiVersion = apiVersion,
  };

  VmaAllocator vmaAllocator;
  if (!VK_SUCCEEDED(
          vmaCreateAllocator(&vmaCreateInfo, &vmaAllocator),
          "the VMA allocator"))
  {
    return VK_NULL_HANDLE;
  }
  return vmaAllocator;
}

rcp<VulkanContext> VulkanContext::make(
    VkInstance instance,
    VkPhysicalDevice physicalDevice,
    VkDevice device,
    const VulkanFeatures& features,
    PFN_vkGetInstanceProcAddr pfnvkGetInstanceProcAddr,
    bool enableDebugNames)
{
  // Building the allocator first means a failure never constructs a context
  // at all.
  VmaAllocator vmaAllocator = make_vma_allocator(
      instance,
      physicalDevice,
      device,
      features.apiVersion,
      pfnvkGetInstanceProcAddr);
  if (vmaAllocator == VK_NULL_HANDLE)
  {
    return nullptr;
  }
  return rcp<VulkanContext>(new VulkanContext(
      instance,
      physicalDevice,
      device,
      features,
      pfnvkGetInstanceProcAddr,
      vmaAllocator,
      enableDebugNames));
}

VulkanContext::VulkanContext(
    VkInstance instance,
    VkPhysicalDevice physicalDevice_,
    VkDevice device_,
    const VulkanFeatures& features_,
    PFN_vkGetInstanceProcAddr pfnvkGetInstanceProcAddr,
    VmaAllocator vmaAllocator,
    bool enableDebugNames)
    : instance(instance)
    , physicalDevice(physicalDevice_)
    , device(device_)
    ,
#define LOAD_VULKAN_INSTANCE_COMMAND(CMD)                                      \
  CMD(reinterpret_cast<PFN_vk##CMD>(                                           \
      pfnvkGetInstanceProcAddr(instance, "vk" #CMD))),
    RIVE_VULKAN_INSTANCE_COMMANDS(LOAD_VULKAN_INSTANCE_COMMAND)
#undef LOAD_VULKAN_INSTANCE_COMMAND
        SetDebugUtilsObjectNameEXT(
            enableDebugNames
                ? reinterpret_cast<PFN_vkSetDebugUtilsObjectNameEXT>(
                      pfnvkGetInstanceProcAddr(
                          instance,
                          "vkSetDebugUtilsObjectNameEXT"))
                : nullptr)
    ,
#define LOAD_VULKAN_DEVICE_COMMAND(CMD)                                        \
  CMD(reinterpret_cast<PFN_vk##CMD>(GetDeviceProcAddr(device, "vk" #CMD))),
    RIVE_VULKAN_DEVICE_COMMANDS(LOAD_VULKAN_DEVICE_COMMAND)
#undef LOAD_VULKAN_DEVICE_COMMAND
        physicalDeviceProperties([this]() {
          VkPhysicalDeviceProperties props;
          GetPhysicalDeviceProperties(physicalDevice, &props);
          return props;
        }())
    , features([this, mutableFeatures = features_]() mutable {
      if (physicalDeviceProperties.vendorID == vkutil::vendors::Qualcomm ||
          physicalDeviceProperties.vendorID == vkutil::vendors::Imagination)
      {
        // We have observed rendering issues on both Adreno and PowerVR when
        // using VK_EXT_color_write_enable.
        mutableFeatures.colorWriteEnable = false;
      }
      return mutableFeatures;
    }())
    , m_vmaAllocator(
          vmaAllocator != VK_NULL_HANDLE ? vmaAllocator
                                         : make_vma_allocator(
                                               instance,
                                               physicalDevice_,
                                               device_,
                                               features_.apiVersion,
                                               pfnvkGetInstanceProcAddr))
{
  // A constructor has nowhere to report this; make() is the path that can
  // recover. make_vma_allocator() already printed why.
  if (m_vmaAllocator == VK_NULL_HANDLE)
  {
    abort();
  }

  // Check that we weren't told the device was more capable than it is
  assert(
      physicalDeviceProperties.apiVersion >= features.apiVersion &&
      "Supplied API version should not be newer than the physical device");

  // VK spec says between D24_S8 and D32_S8, one of them must be supported
  m_supportsD24S8 = isFormatSupportedWithFeatureFlags(
      VK_FORMAT_D24_UNORM_S8_UINT,
      VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT);

  // This assert should never fire unless some hardware is breaking the VK
  // spec by reporting that it does not support one of these formats.
  assert(
      (m_supportsD24S8 ||
       isFormatSupportedWithFeatureFlags(
           VK_FORMAT_D32_SFLOAT_S8_UINT,
           VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT)) &&
      "No suitable depth format supported!");
}

VulkanContext::~VulkanContext() { vmaDestroyAllocator(m_vmaAllocator); }

bool VulkanContext::isFormatSupportedWithFeatureFlags(
    VkFormat format,
    VkFormatFeatureFlagBits featureFlags)
{
  // Can flesch this out, but currently just checks if format's optimal tiling
  // features include the provided bits.
  VkFormatProperties properties;
  GetPhysicalDeviceFormatProperties(physicalDevice, format, &properties);
  return properties.optimalTilingFeatures & featureFlags;
}

rcp<vkutil::Buffer> VulkanContext::makeBuffer(
    const VkBufferCreateInfo& info,
    vkutil::Mappability mappability)
{
  return rcp(new vkutil::Buffer(ref_rcp(this), info, mappability));
}

rcp<vkutil::Image> VulkanContext::makeImage(
    const VkImageCreateInfo& info,
    const char* name)
{
  return rcp(new vkutil::Image(ref_rcp(this), info, name));
}

rcp<vkutil::Image> VulkanContext::makeExternalImage(
    VkImage existingImage,
    const VkImageCreateInfo& info,
    const char* name)
{
  return rcp(new vkutil::Image(ref_rcp(this), existingImage, info, name));
}

rcp<vkutil::Framebuffer> VulkanContext::makeFramebuffer(
    const VkFramebufferCreateInfo& info)
{
  return rcp(new vkutil::Framebuffer(ref_rcp(this), info));
}

static VkImageViewType image_view_type_for_image_type(VkImageType type)
{
  switch (type)
  {
    case VK_IMAGE_TYPE_2D:
      return VK_IMAGE_VIEW_TYPE_2D;
    default:
      fprintf(
          stderr,
          "vkutil::image_view_type_for_image_type: VkImageType %u is not "
          "supported\n",
          type);
  }
  RIVE_UNREACHABLE();
}

static VkImageAspectFlags image_aspect_flags_for_format(VkFormat format)
{
  switch (format)
  {
    case VK_FORMAT_D24_UNORM_S8_UINT:
    case VK_FORMAT_D32_SFLOAT_S8_UINT:
      return VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT;
    default:
      return VK_IMAGE_ASPECT_COLOR_BIT;
  }
  RIVE_UNREACHABLE();
}

rcp<vkutil::ImageView> VulkanContext::makeImageView(
    rcp<vkutil::Image> image,
    const char* name)
{
  const VkImageCreateInfo& texInfo = image->info();

  return makeImageView(
      image,
      {
          .viewType = image_view_type_for_image_type(texInfo.imageType),
          .format = texInfo.format,
          .subresourceRange =
              {
                  .aspectMask = image_aspect_flags_for_format(texInfo.format),
                  .levelCount = texInfo.mipLevels,
                  .layerCount = 1,
              },
      },
      name);
}

rcp<vkutil::ImageView> VulkanContext::makeImageView(
    rcp<vkutil::Image> image,
    const VkImageViewCreateInfo& info,
    const char* name)
{
  assert(image);
  return rcp(
      new vkutil::ImageView(ref_rcp(this), std::move(image), info, name));
}

rcp<vkutil::ImageView> VulkanContext::makeExternalImageView(
    const VkImageViewCreateInfo& info,
    const char* name)
{
  return rcp<vkutil::ImageView>(
      new vkutil::ImageView(ref_rcp(this), nullptr, info, name));
}

rcp<vkutil::Texture2D> VulkanContext::makeTexture2D(
    const VkImageCreateInfo& info,
    const char* name)
{
  return rcp<vkutil::Texture2D>(
      new vkutil::Texture2D(ref_rcp(this), info, name));
}

rcp<vkutil::Texture2D> VulkanContext::makeTexture2D(
    rcp<vkutil::Image> existingImage,
    const char* name)
{
  return rcp<vkutil::Texture2D>(
      new vkutil::Texture2D(ref_rcp(this), std::move(existingImage), name));
}

void VulkanContext::updateImageDescriptorSets(
    VkDescriptorSet vkDescriptorSet,
    VkWriteDescriptorSet writeSet,
    std::initializer_list<VkDescriptorImageInfo> imageInfos)
{
  writeSet.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
  writeSet.dstSet = vkDescriptorSet;
  writeSet.descriptorCount = math::lossless_numeric_cast<uint32_t>(
      imageInfos.size());
  writeSet.pImageInfo = imageInfos.begin();
  UpdateDescriptorSets(device, 1, &writeSet, 0, nullptr);
}

void VulkanContext::updateBufferDescriptorSets(
    VkDescriptorSet vkDescriptorSet,
    VkWriteDescriptorSet writeSet,
    std::initializer_list<VkDescriptorBufferInfo> bufferInfos)
{
  writeSet.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
  writeSet.dstSet = vkDescriptorSet;
  writeSet.descriptorCount = math::lossless_numeric_cast<uint32_t>(
      bufferInfos.size());
  writeSet.pBufferInfo = bufferInfos.begin();
  UpdateDescriptorSets(device, 1, &writeSet, 0, nullptr);
}

void VulkanContext::memoryBarrier(
    VkCommandBuffer commandBuffer,
    VkPipelineStageFlags srcStageMask,
    VkPipelineStageFlags dstStageMask,
    VkDependencyFlags dependencyFlags,
    VkMemoryBarrier memoryBarrier)
{
  memoryBarrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
  CmdPipelineBarrier(
      commandBuffer,
      srcStageMask,
      dstStageMask,
      dependencyFlags,
      1,
      &memoryBarrier,
      0,
      nullptr,
      0,
      nullptr);
}

void VulkanContext::imageMemoryBarriers(
    VkCommandBuffer commandBuffer,
    VkPipelineStageFlags srcStageMask,
    VkPipelineStageFlags dstStageMask,
    VkDependencyFlags dependencyFlags,
    uint32_t count,
    VkImageMemoryBarrier* imageMemoryBarriers)
{
  for (uint32_t i = 0; i < count; ++i)
  {
    auto& imageMemoryBarrier = imageMemoryBarriers[i];
    imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    if (imageMemoryBarrier.subresourceRange.aspectMask == 0)
    {
      imageMemoryBarrier.subresourceRange
          .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    }
    if (imageMemoryBarrier.subresourceRange.levelCount == 0)
    {
      imageMemoryBarrier.subresourceRange.levelCount = VK_REMAINING_MIP_LEVELS;
    }
    if (imageMemoryBarrier.subresourceRange.layerCount == 0)
    {
      imageMemoryBarrier.subresourceRange
          .layerCount = VK_REMAINING_ARRAY_LAYERS;
    }
  }
  CmdPipelineBarrier(
      commandBuffer,
      srcStageMask,
      dstStageMask,
      dependencyFlags,
      0,
      nullptr,
      0,
      nullptr,
      count,
      imageMemoryBarriers);
}

const vkutil::ImageAccess& VulkanContext::simpleImageMemoryBarrier(
    VkCommandBuffer commandBuffer,
    const vkutil::ImageAccess& srcAccess,
    const vkutil::ImageAccess& dstAccess,
    VkImage image,
    vkutil::ImageAccessAction imageAccessAction,
    VkDependencyFlags dependencyFlags)
{
  assert(image != VK_NULL_HANDLE);
  if (srcAccess != dstAccess)
  {
    imageMemoryBarrier(
        commandBuffer,
        srcAccess.pipelineStages,
        dstAccess.pipelineStages,
        dependencyFlags,
        {
            .srcAccessMask = srcAccess.accessMask,
            .dstAccessMask = dstAccess.accessMask,
            .oldLayout = imageAccessAction ==
                                 vkutil::ImageAccessAction::preserveContents
                             ? srcAccess.layout
                             : VK_IMAGE_LAYOUT_UNDEFINED,
            .newLayout = dstAccess.layout,
            .image = image,
        });
  }
  return dstAccess;
}

void VulkanContext::bufferMemoryBarrier(
    VkCommandBuffer commandBuffer,
    VkPipelineStageFlags srcStageMask,
    VkPipelineStageFlags dstStageMask,
    VkDependencyFlags dependencyFlags,
    VkBufferMemoryBarrier bufferMemoryBarrier)
{
  bufferMemoryBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
  if (bufferMemoryBarrier.size == 0)
  {
    bufferMemoryBarrier.size = VK_WHOLE_SIZE;
  }
  CmdPipelineBarrier(
      commandBuffer,
      srcStageMask,
      dstStageMask,
      dependencyFlags,
      0,
      nullptr,
      1,
      &bufferMemoryBarrier,
      0,
      nullptr);
}

void VulkanContext::clearColorImage(
    VkCommandBuffer commandBuffer,
    ColorInt clearColor,
    VkImage image,
    VkImageLayout imageLayout)
{
  const VkClearColorValue colorClearValue = {
      vkutil::color_clear_rgba32f(clearColor)};

  const VkImageSubresourceRange colorClearRange = {
      .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
      .levelCount = 1,
      .layerCount = 1,
  };

  CmdClearColorImage(
      commandBuffer,
      image,
      imageLayout,
      &colorClearValue,
      1,
      &colorClearRange);
}

void VulkanContext::blitSubRect(
    VkCommandBuffer commandBuffer,
    VkImage srcImage,
    VkImageLayout srcImageLayout,
    VkImage dstImage,
    VkImageLayout dstImageLayout,
    const IAABB& blitBounds)
{
  if (blitBounds.empty())
  {
    return;
  }

  VkImageBlit imageBlit = {
      .srcSubresource =
          {
              .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
              .layerCount = 1,
          },
      .dstSubresource =
          {
              .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
              .layerCount = 1,
          },
  };

  imageBlit.srcOffsets[0] = {blitBounds.left, blitBounds.top, 0};
  imageBlit.srcOffsets[1] = {blitBounds.right, blitBounds.bottom, 1};

  imageBlit.dstOffsets[0] = {blitBounds.left, blitBounds.top, 0};
  imageBlit.dstOffsets[1] = {blitBounds.right, blitBounds.bottom, 1};

  CmdBlitImage(
      commandBuffer,
      srcImage,
      srcImageLayout,
      dstImage,
      dstImageLayout,
      1,
      &imageBlit,
      VK_FILTER_NEAREST);
}

void VulkanContext::setDebugNameIfEnabled(
    uint64_t handle,
    VkObjectType objectType,
    const char* name)
{
  // A null handle means creation failed; naming it would be invalid usage.
  if (SetDebugUtilsObjectNameEXT != nullptr && name != nullptr && handle != 0)
  {
    VkDebugUtilsObjectNameInfoEXT nameInfo = {
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = objectType,
        .objectHandle = handle,
        .pObjectName = name,
    };

    SetDebugUtilsObjectNameEXT(device, &nameInfo);
  }
}

} // namespace rive::gpu