#include "draw.h"
#include "gpu.h"

/*
 * Must be updated if any of the following changes
 * - ui_context->max_containers <- uint32
 * - ui_context->containers[*].layer_count <- uint32
 * - ui_context->containers[*].layers[*].address <- device address
 * Basically, needs to be re-run whenever the number of layers/containers changes
 */
void record_ui_compute(VkCommandBuffer command_buffer, UIContext* ui_context) {
  VkDeviceAddress push[2] = {ui_context->address, 0};

  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, ui_context->string_pipeline.pipeline);
  for(uint32_t i = 0; i < ui_context->max_containers; i++) {
    if(ui_context->containers[i].id != 0x00000000) {
      for(uint32_t j = 0; j < ui_context->containers[i].layer_count; j++) {
        push[1] = ui_context->containers[i].layers[j].address;
        VkBufferMemoryBarrier draw_command_barrier_1 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_INDIRECT_COMMAND_READ_BIT,
          .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, NULL, 1, &draw_command_barrier_1, 0, NULL);
        command_copy_buffer(command_buffer, ui_context->containers[i].layers[j].layer, ui_context->containers[i].layers[j].layer, offsetof(GPULayer, num_drawables), offsetof(GPULayer, draw) + offsetof(DrawCommand, instance_count), sizeof(uint32_t));
        VkBufferMemoryBarrier draw_command_barrier_2 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
          .dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, NULL, 1, &draw_command_barrier_2, 0, NULL);
        vkCmdPushConstants(command_buffer, ui_context->string_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, 16, push);
        vkCmdDispatchIndirect(command_buffer, ui_context->containers[i].layers[j].layer, offsetof(GPULayer, dispatch_strings));
        VkBufferMemoryBarrier draw_command_barrier_3 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
          .dstAccessMask = VK_ACCESS_INDIRECT_COMMAND_READ_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT, 0, 0, NULL, 1, &draw_command_barrier_3, 0, NULL);

        VkBufferMemoryBarrier drawables_barrier = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].drawables,
          .offset = 0,
          .size = sizeof(GPUDrawable)*ui_context->containers[i].layers[j].data.max_drawables,
          .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
          .dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_VERTEX_SHADER_BIT, 0, 0, NULL, 1, &drawables_barrier, 0, NULL);
      }
    }
  }
}

/*
 * Must be updated if any of the following changes
 * - ui_context->max_containers <- uint32
 * - ui_context->containers[*].layer_count <- uint32
 * - ui_context->containers[*].layers[*].address <- device address
 * Basically, needs to be re-run whenever the number of layers/containers changes
 */
void record_ui_draw(VkCommandBuffer command_buffer, UIContext* ui_context) {
  VkDeviceAddress push[2] = {ui_context->address, 0};

  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, ui_context->pipeline.pipeline);
  vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, ui_context->pipeline.layout, 0, 1, &ui_context->font_samplers, 0, NULL);
  vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, ui_context->pipeline.layout, 1, 1, &ui_context->font_textures, 0, NULL);
  vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, ui_context->pipeline.layout, 2, 1, &ui_context->samplers, 0, NULL);
  vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, ui_context->pipeline.layout, 3, 1, &ui_context->textures, 0, NULL);
  for(uint32_t i = 0; i < ui_context->max_containers; i++) {
    if(ui_context->containers[i].id != 0x00000000) {
      for(uint32_t j = 0; j < ui_context->containers[i].layer_count; j++) {
        push[1] = ui_context->containers[i].layers[j].address;
        vkCmdPushConstants(command_buffer, ui_context->pipeline.layout, VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, 16, push);
        vkCmdDrawIndirect(command_buffer, ui_context->containers[i].layers[j].layer, offsetof(GPULayer, draw), 1, 0);
      }
    }
  }
}

VkResult record_draw_commands(
    RenderContext* context,
    UIContext* ui_context) {
  VkResult result;
  for(uint32_t image_index = 0; image_index < context->swapchain_image_count; image_index++) {
    VkCommandBuffer command_buffer = context->swapchain_command_buffers[image_index];
    VK_RESULT(vkResetCommandBuffer(command_buffer, 0));

    VkCommandBufferBeginInfo begin_info = {
      .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
    }; 
    VK_RESULT(vkBeginCommandBuffer(command_buffer, &begin_info));

    VkViewport viewport = {
      .width = context->swapchain_extent.width,
      .height = context->swapchain_extent.height,
      .maxDepth = 1.0f,
      .minDepth = 0.0f,
    };
    vkCmdSetViewport(command_buffer, 0, 1, &viewport);

    VkRect2D scissor = {
      .extent = context->swapchain_extent,
    };
    vkCmdSetScissor(command_buffer, 0, 1, &scissor);

    VkClearValue clear_values[2] = {{.color={{0.0f, 0.0f, 0.0f, 0.0f}}}, {.depthStencil={1.0f, 0.0f}}};

    VkRenderPassBeginInfo render_pass_begin = {
      .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
      .renderPass = context->render_pass,
      .framebuffer = context->swapchain_framebuffers[image_index],
      .renderArea.offset = {0, 0},
      .renderArea.extent = context->swapchain_extent,
      .clearValueCount = 2,
      .pClearValues = clear_values,
    }; 

    record_ui_compute(command_buffer, ui_context);
    vkCmdBeginRenderPass(command_buffer, &render_pass_begin, VK_SUBPASS_CONTENTS_INLINE);
    // Render World
    vkCmdNextSubpass(command_buffer, VK_SUBPASS_CONTENTS_INLINE);
    // Render UI
    record_ui_draw(command_buffer, ui_context);
    vkCmdEndRenderPass(command_buffer);

    VK_RESULT(vkEndCommandBuffer(command_buffer));
  }


  return VK_SUCCESS;
}

VkResult draw_frame(
    RenderContext* context,
    double time) {
  (void)time;
  VkResult result;

  result = vkWaitForFences(context->device, 1, &context->in_flight_fences[context->current_frame], VK_TRUE, UINT64_MAX);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = vkResetFences(context->device, 1, &context->in_flight_fences[context->current_frame]);
  if(result != VK_SUCCESS) {
    return result;
  }

  uint32_t image_index;
  result = vkAcquireNextImageKHR(context->device, context->swapchain, UINT64_MAX, context->image_available_semaphores[context->current_frame], VK_NULL_HANDLE, &image_index);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkPipelineStageFlags wait_stages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
  VkSubmitInfo submit_info = {
    .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
    .waitSemaphoreCount = 1,
    .pWaitSemaphores = &context->image_available_semaphores[context->current_frame],
    .pWaitDstStageMask = wait_stages,
    .commandBufferCount = 1,
    .pCommandBuffers = &context->swapchain_command_buffers[image_index],
    .signalSemaphoreCount = 1,
    .pSignalSemaphores = &context->render_finished_semaphores[context->current_frame],
  };

  result = vkQueueSubmit(context->graphics_queue.handle, 1, &submit_info, context->in_flight_fences[context->current_frame]);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkPresentInfoKHR present_info = {
    .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
    .waitSemaphoreCount = 1,
    .pWaitSemaphores = &context->render_finished_semaphores[context->current_frame],
    .swapchainCount = 1,
    .pSwapchains = &context->swapchain,
    .pImageIndices = &image_index,
    .pResults = 0,
  };

  result = vkQueuePresentKHR(context->present_queue.handle, &present_info);
  if(result != VK_SUCCESS) {
    return result;
  }

  context->current_frame = (context->current_frame + 1) % MAX_FRAMES_IN_FLIGHT;

  return VK_SUCCESS;
}