/*
 * gulkan
 * Copyright 2018 Collabora Ltd.
 * Author: Lubosz Sarnecki <lubosz.sarnecki@collabora.com>
 * SPDX-License-Identifier: MIT
 */

#include "gulkan-frame-buffer.h"

struct _GulkanFrameBuffer
{
  GObject parent;

  GulkanDevice *device;

  /*
  GulkanTexture *color_texture;
  GulkanTexture *depth_stencil_texture;
  */
  VkImage        color_image;
  VkDeviceMemory color_memory;
  VkImageView    color_image_view;
  VkImage        depth_stencil_image;
  VkDeviceMemory depth_stencil_memory;
  VkImageView    depth_stencil_image_view;

  VkRenderPass   render_pass;
  VkFramebuffer  framebuffer;

  uint32_t width;
  uint32_t height;
};

G_DEFINE_TYPE (GulkanFrameBuffer, gulkan_frame_buffer, G_TYPE_OBJECT)

static void
gulkan_frame_buffer_finalize (GObject *gobject);

static void
gulkan_frame_buffer_class_init (GulkanFrameBufferClass *klass)
{
  GObjectClass *object_class = G_OBJECT_CLASS (klass);

  object_class->finalize = gulkan_frame_buffer_finalize;
}

static void
gulkan_frame_buffer_init (GulkanFrameBuffer *self)
{
  self->color_image = VK_NULL_HANDLE;
  self->color_memory = VK_NULL_HANDLE;
}

GulkanFrameBuffer *
gulkan_frame_buffer_new (void)
{
  return (GulkanFrameBuffer*) g_object_new (GULKAN_TYPE_FRAME_BUFFER, 0);
}

static void
gulkan_frame_buffer_finalize (GObject *gobject)
{
  GulkanFrameBuffer *self = GULKAN_FRAME_BUFFER (gobject);
  if (self->color_image_view == VK_NULL_HANDLE)
    return;

  VkDevice device = gulkan_device_get_handle (self->device);

  vkDestroyImageView (device, self->color_image_view, NULL);
  if (self->color_image != VK_NULL_HANDLE)
    vkDestroyImage (device, self->color_image, NULL);
  if (self->color_memory != VK_NULL_HANDLE)
    vkFreeMemory (device, self->color_memory, NULL);
  vkDestroyImageView (device, self->depth_stencil_image_view, NULL);
  vkDestroyImage (device, self->depth_stencil_image, NULL);
  vkFreeMemory (device, self->depth_stencil_memory, NULL);
  vkDestroyRenderPass (device, self->render_pass, NULL);
  vkDestroyFramebuffer (device, self->framebuffer, NULL);
}

static gboolean
_create_image_view (GulkanFrameBuffer     *self,
                    VkImage                image,
                    VkFormat               format,
                    VkImageAspectFlagBits  aspect_flag,
                    VkImageView           *out_image_view)
{
  VkImageViewCreateInfo image_view_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
    .flags = 0,
    .image = image,
    .viewType = VK_IMAGE_VIEW_TYPE_2D,
    .format = format,
    .components = {
      .r = VK_COMPONENT_SWIZZLE_IDENTITY,
      .g = VK_COMPONENT_SWIZZLE_IDENTITY,
      .b = VK_COMPONENT_SWIZZLE_IDENTITY,
      .a = VK_COMPONENT_SWIZZLE_IDENTITY
    },
    .subresourceRange = {
      .aspectMask = aspect_flag,
      .baseMipLevel = 0,
      .levelCount = 1,
      .baseArrayLayer = 0,
      .layerCount = 1,
    }
  };

  VkResult res = vkCreateImageView (gulkan_device_get_handle (self->device),
                                   &image_view_info, NULL, out_image_view);
  if (res != VK_SUCCESS)
    {
      g_print ("vkCreateImageView failed: %d\n", res);
      return FALSE;
    }

  return TRUE;
}

static gboolean
_init_target (GulkanDevice          *device,
              uint32_t               width,
              uint32_t               height,
              VkSampleCountFlagBits  sample_count,
              VkFormat               format,
              VkImageUsageFlags      usage,
              VkImage               *out_image,
              VkDeviceMemory        *out_memory)
{
  VkDevice vk_device = gulkan_device_get_handle (device);

  /* Depth/stencil target */
  VkImageCreateInfo image_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
    .imageType = VK_IMAGE_TYPE_2D,
    .extent = {
      .width = width,
      .height = height,
      .depth = 1,
    },
    .mipLevels = 1,
    .arrayLayers = 1,
    .format = format,
    .tiling = VK_IMAGE_TILING_OPTIMAL,
    .samples = sample_count,
    .usage = usage,
    .flags = 0
  };

  VkResult res;
  res = vkCreateImage (vk_device, &image_info, NULL, out_image);
  if (res != VK_SUCCESS)
    {
      g_print ("vkCreateImage failed for depth buffer: %d\n", res);
      return FALSE;
    }

  VkMemoryRequirements memory_requirements;
  vkGetImageMemoryRequirements (vk_device, *out_image, &memory_requirements);

  VkMemoryAllocateInfo memory_info = {
    .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
    .allocationSize = memory_requirements.size
  };
  if (!gulkan_device_memory_type_from_properties (
          device, memory_requirements.memoryTypeBits,
          VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
          &memory_info.memoryTypeIndex))
    {
      g_print ("Failed to find memory type matching requirements.\n");
      return FALSE;
    }

  res = vkAllocateMemory (vk_device, &memory_info, NULL, out_memory);
  if (res != VK_SUCCESS)
    {
      g_print ("Failed to find memory for image.\n");
      return FALSE;
    }

  res = vkBindImageMemory (vk_device, *out_image, *out_memory, 0);
  if (res != VK_SUCCESS)
    {
      g_print ("Failed to bind memory for image.\n");
      return FALSE;
    }

  return TRUE;
}

static gboolean
_init_render_pass (GulkanFrameBuffer    *self,
                   VkSampleCountFlagBits samples,
                   VkFormat              color_format)
{
  VkDevice vk_device = gulkan_device_get_handle (self->device);

  VkRenderPassCreateInfo renderpass_info = {
    .sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO,
    .flags = 0,
    .attachmentCount = 2,
    .pAttachments = (VkAttachmentDescription[]) {
      {
        .format = color_format,
        .samples = samples,
        .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
        .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
        .finalLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        .flags = 0
      },
      {
        .format = VK_FORMAT_D32_SFLOAT,
        .samples = samples,
        .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
        .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
        .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
        .finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
        .flags = 0
      },
    },
    .subpassCount = 1,
    .pSubpasses = &(VkSubpassDescription) {
      .pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
      .colorAttachmentCount = 1,
      .pColorAttachments = &(VkAttachmentReference) {
        .attachment = 0,
        .layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
      },
      .pDepthStencilAttachment = &(VkAttachmentReference) {
        .attachment = 1,
        .layout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
      },
      .pResolveAttachments = NULL
     },
    .dependencyCount = 0,
    .pDependencies = NULL
  };

  VkResult res = vkCreateRenderPass (vk_device, &renderpass_info,
                                     NULL, &self->render_pass);
  if (res != VK_SUCCESS)
    {
      g_print ("vkCreateRenderPass failed with error %d.\n", res);
      return false;
    }
  return TRUE;
}

static gboolean
_init_frame_buffer (GulkanFrameBuffer *self, uint32_t width, uint32_t height)
{
  VkDevice vk_device = gulkan_device_get_handle (self->device);

  VkFramebufferCreateInfo framebuffer_info = {
    .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
    .renderPass = self->render_pass,
    .attachmentCount = 2,
    .pAttachments = (VkImageView[]) {
      self->color_image_view,
      self->depth_stencil_image_view
    },
    .width = width,
    .height = height,
    .layers = 1
  };

  VkResult res = vkCreateFramebuffer (vk_device,
                                     &framebuffer_info, NULL,
                                     &self->framebuffer);
  if (res != VK_SUCCESS)
    {
      g_print ("vkCreateFramebuffer failed with error %d.\n", res);
      return FALSE;
    }
  return TRUE;
}

bool
gulkan_frame_buffer_initialize (GulkanFrameBuffer    *self,
                                GulkanDevice         *device,
                                uint32_t              width,
                                uint32_t              height,
                                VkSampleCountFlagBits sample_count,
                                VkFormat              color_format)
{
  VkImageUsageFlags usage =
        VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT |
        VK_IMAGE_USAGE_TRANSFER_SRC_BIT;

  if (!_init_target (device, width, height, sample_count, color_format,
                     usage, &self->color_image, &self->color_memory))
    return FALSE;

  if (!gulkan_frame_buffer_initialize_from_image (self, device,
                                                  self->color_image,
                                                  width, height,
                                                  sample_count, color_format))
    return FALSE;

  return TRUE;
}

bool
gulkan_frame_buffer_initialize_from_image (GulkanFrameBuffer    *self,
                                           GulkanDevice         *device,
                                           VkImage               color_image,
                                           uint32_t              width,
                                           uint32_t              height,
                                           VkSampleCountFlagBits sample_count,
                                           VkFormat              color_format)
{
  self->device = device;
  self->width = width;
  self->height = height;

  if (!_create_image_view (self, color_image, color_format,
                           VK_IMAGE_ASPECT_COLOR_BIT,
                          &self->color_image_view))
    return FALSE;

  VkFormat depth_format = VK_FORMAT_D32_SFLOAT;

  if (!_init_target (device, width, height, sample_count,
                     depth_format,
                     VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                     &self->depth_stencil_image, &self->depth_stencil_memory))
    return FALSE;

  if (!_create_image_view (self, self->depth_stencil_image,
                           depth_format, VK_IMAGE_ASPECT_DEPTH_BIT,
                           &self->depth_stencil_image_view))
    return FALSE;

  if (!_init_render_pass (self, sample_count, color_format))
    return FALSE;

  if (!_init_frame_buffer (self, width, height))
    return FALSE;

  return TRUE;
}

void
gulkan_frame_buffer_begin_pass (GulkanFrameBuffer *self,
                                VkCommandBuffer    cmd_buffer)
{
  VkRenderPassBeginInfo render_pass_info = {
    .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
    .renderPass = self->render_pass,
    .framebuffer = self->framebuffer,
    .renderArea = {
      .offset = {
        .x = 0,
        .y = 0,
      },
      .extent = {
        .width = self->width,
        .height = self->height,
      }
    },
    .clearValueCount = 2,
    .pClearValues = (VkClearValue[]) {
      {
        .color = {
          .float32 = {0.0f, 0.0f, 0.0f, 1.0f}
        }
      },
      {
        .depthStencil = {
          .depth = 1.0f,
          .stencil = 0
        }
      },
    }
  };

  vkCmdBeginRenderPass (cmd_buffer, &render_pass_info,
                        VK_SUBPASS_CONTENTS_INLINE);
}

VkImage
gulkan_frame_buffer_get_color_image (GulkanFrameBuffer *self)
{
  return self->color_image;
}

VkRenderPass
gulkan_frame_buffer_get_render_pass (GulkanFrameBuffer *self)
{
  return self->render_pass;
}
