#include "draw.h"
#include "hex.h"
#include "vulkan/vulkan_core.h"

void record_ui_draw(VkCommandBuffer command_buffer, UIContext* ui_context, double time, uint32_t frame) {
  UIPushConstant push = {
    .time = (float)time,
    .layer = 0,
    .pad = frame,
  };

  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++) {
    for(uint32_t j = 0; j < ui_context->containers[i].layer_count; j++) {
      push.layer = ui_context->containers[i].layers[j].address[frame];
      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[frame], offsetof(GPULayer, draw), 1, 0);
    }
  }
}

void record_hex_draw(VkCommandBuffer command_buffer, HexContext* hex, double time, uint32_t frame) {
  HexPushConstant push = {
    .context = hex->address,
    .time = (float)time,
    .region = 0,
  };
  vkCmdPushConstants(command_buffer, hex->graphics.layout, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(HexPushConstant), &push);
  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, hex->graphics.pipeline);
  vkCmdDraw(command_buffer, 18, REGION_HEX_COUNT*MAX_LOADED_REGIONS, 0, 0);
}

void record_ui_compute(VkCommandBuffer command_buffer, UIContext* ui, uint32_t frame) {
  UIPushConstant push = {
    .time = 0.0,
    .layer = 0,
    .pad = frame,
  };

  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, ui->string_pipeline.pipeline);
  for(uint32_t i = 0; i < ui->max_containers; i++) {
    if(ui->containers[i].id != 0x00000000) {
      for(uint32_t j = 0; j < ui->containers[i].layer_count; j++) {
        command_copy_buffer(command_buffer, ui->containers[i].layers[j].layer[frame], ui->containers[i].layers[j].layer[frame], offsetof(GPULayer, num_drawables), offsetof(GPULayer, draw) + offsetof(DrawCommand, instance_count), sizeof(uint32_t));
      }
    }
  }

  for(uint32_t i = 0; i < ui->max_containers; i++) {
    if(ui->containers[i].id != 0x00000000) {
      for(uint32_t j = 0; j < ui->containers[i].layer_count; j++) {
        push.layer = ui->containers[i].layers[j].address[frame];
        vkCmdPushConstants(command_buffer, ui->string_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, 16, &push);
        vkCmdDispatchIndirect(command_buffer, ui->containers[i].layers[j].layer[frame], offsetof(GPULayer, dispatch_strings));
      }
    }
  }

}

VkResult draw_frame(
    RenderContext* context,
    UIContext* ui,
    HexContext* hex,
    double time) {
  VkResult result;

  FrameContext* frame = &context->frame[context->current_frame];

  VkFence fences[] = {frame->ready};
  VK_RESULT(vkWaitForFences(context->device, 1, fences, VK_TRUE, UINT64_MAX));
  VK_RESULT(vkResetFences(context->device, 1, fences));

  VkCommandBufferBeginInfo begin_info = {
    .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
  }; 

  TransferBuffer* transfer = &frame->transfers[frame->transfer_index % 2];
  if(transfer->count > 0) {
    VkCommandBuffer transfer_commands = frame->transfer_commands;
    VK_RESULT(vkResetCommandBuffer(transfer_commands, 0));
    VK_RESULT(vkBeginCommandBuffer(transfer_commands, &begin_info));

    VkDeviceSize src_offset = 0;
    for(uint32_t tid = 0; tid < transfer->count; tid++) {
      command_copy_buffer(
          transfer_commands,
          transfer->buffer,
          transfer->infos[tid].buffer,
          src_offset,
          transfer->infos[tid].offset,
          transfer->infos[tid].size);
      src_offset += transfer->infos[tid].size;
    }
    VK_RESULT(vkEndCommandBuffer(transfer_commands));

    frame->transfer_index += 1;
    transfer->count = 0;
    transfer->written = 0;

    VkTimelineSemaphoreSubmitInfo transfer_timeline = {
      .sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO,
      .signalSemaphoreValueCount = 1,
      .pSignalSemaphoreValues = &frame->transfer_index,
    };
    VkSubmitInfo transfer_submit = {
      .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
      .commandBufferCount = 1,
      .pCommandBuffers = &transfer_commands,
      .pSignalSemaphores = &frame->transfer,
      .signalSemaphoreCount = 1,
      .pNext = &transfer_timeline,
    };
    VK_RESULT(vkQueueSubmit(context->transfer_queue.handle, 1, &transfer_submit, VK_NULL_HANDLE));

    VkCommandBuffer compute_commands = frame->compute_commands;
    VK_RESULT(vkResetCommandBuffer(compute_commands, 0));
    VK_RESULT(vkBeginCommandBuffer(compute_commands, &begin_info));
    VkMemoryBarrier compute_barrier = {
      .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
      .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
      .dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT,
    };
    vkCmdPipelineBarrier(compute_commands, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, &compute_barrier, 0, NULL, 0, NULL);
    record_ui_compute(compute_commands, ui, context->current_frame);
    VK_RESULT(vkEndCommandBuffer(compute_commands));

    frame->compute_index += 1;
    VkPipelineStageFlags compute_wait_stages[] = {VK_PIPELINE_STAGE_ALL_COMMANDS_BIT};
    VkTimelineSemaphoreSubmitInfo compute_timeline = {
      .sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO,
      .waitSemaphoreValueCount = 1,
      .pWaitSemaphoreValues = &frame->transfer_index,
      .signalSemaphoreValueCount = 1,
      .pSignalSemaphoreValues = &frame->compute_index,
    };
    VkSubmitInfo compute_submit = {
      .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
      .commandBufferCount = 1,
      .pCommandBuffers = &compute_commands,
      .pSignalSemaphores = &frame->compute,
      .signalSemaphoreCount = 1,
      .pWaitSemaphores = &frame->transfer,
      .pWaitDstStageMask = compute_wait_stages,
      .waitSemaphoreCount = 1,
      .pNext = &compute_timeline,
    };
    VK_RESULT(vkQueueSubmit(context->transfer_queue.handle, 1, &compute_submit, VK_NULL_HANDLE));
  }

  uint32_t image_index;
  result = vkAcquireNextImageKHR(context->device, context->swapchain, UINT64_MAX, frame->image, VK_NULL_HANDLE, &image_index);
  if(result == VK_ERROR_OUT_OF_DATE_KHR) {
    VK_RESULT(recreate_framebuffer(context));
  } else if(result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
    return result;
  }

  VkCommandBuffer command_buffer = context->swapchain_command_buffers[image_index];
  VK_RESULT(vkResetCommandBuffer(command_buffer, 0));
  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,
  }; 

  vkCmdBeginRenderPass(command_buffer, &render_pass_begin, VK_SUBPASS_CONTENTS_INLINE);
  record_hex_draw(command_buffer, hex, time, context->current_frame);
  vkCmdNextSubpass(command_buffer, VK_SUBPASS_CONTENTS_INLINE);
  record_ui_draw(command_buffer, ui, time, context->current_frame);
  vkCmdEndRenderPass(command_buffer);
  VK_RESULT(vkEndCommandBuffer(command_buffer));

  VkPipelineStageFlags wait_stages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT};
  VkSemaphore wait_semaphores[] = {frame->image, frame->compute};
  VkSemaphore signal_semaphores[] = {frame->render};
  uint64_t wait_values[] = {0, frame->compute_index};
  uint64_t signal_values[] = {0};
  VkTimelineSemaphoreSubmitInfo timeline_info = {
    .sType = VK_STRUCTURE_TYPE_TIMELINE_SEMAPHORE_SUBMIT_INFO,
    .signalSemaphoreValueCount = sizeof(signal_values)/sizeof(uint64_t),
    .pSignalSemaphoreValues = signal_values,
    .waitSemaphoreValueCount = sizeof(wait_values)/sizeof(uint64_t),
    .pWaitSemaphoreValues = wait_values,
  };
  VkSubmitInfo submit_info = {
    .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
    .waitSemaphoreCount = sizeof(wait_semaphores)/sizeof(VkSemaphore),
    .pWaitSemaphores = wait_semaphores,
    .pWaitDstStageMask = wait_stages,
    .commandBufferCount = 1,
    .pCommandBuffers = &context->swapchain_command_buffers[image_index],
    .signalSemaphoreCount = sizeof(signal_semaphores)/sizeof(VkSemaphore),
    .pSignalSemaphores = signal_semaphores,
    .pNext = &timeline_info,
  };

  VK_RESULT(vkQueueSubmit(context->graphics_queue.handle, 1, &submit_info, frame->ready));

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

  result = vkQueuePresentKHR(context->present_queue.handle, &present_info);
  if(result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR) {
    VK_RESULT(recreate_framebuffer(context));
  } else if(result != VK_SUCCESS) {
    return result;
  }

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

  return VK_SUCCESS;
}