#include <evk/evk.h>
#include <ev_numeric.h>
#include <ev_helpers.h>

#define GLFW_INCLUDE_NONE
#include <GLFW/glfw3.h>

evstring PROJECT_NAME = evstr("evk");

int main(void) 
{
  u32 width = 1024;
  u32 height = 1024;

  evkInstance instance = evkCreateInstance((evkInstanceCreateInfo){
    .applicationInfo = EV_DEFAULT(evkApplicationInfo),
    .layers = svec_init(evstring, {
      evstr("VK_LAYER_KHRONOS_validation"),
    }),
    .extensions = svec_init(evstring, {
      evstr("VK_KHR_surface"),
// TODO Get these from GLFW
#if EV_OS_WINDOWS
      evstr("VK_KHR_win32_surface"),
#elif EV_OS_LINUX
      evstr("VK_KHR_xcb_surface"),
#endif
    }),
  });

  if(instance.vk == EV_INVALID(VkInstance))
  {
    puts("Instance creation failed.");
    goto InstanceCreationFailed;
  }
  puts("Instance was created successfully.");

  evkDevice device = evkCreateDevice((evkDeviceCreateInfo) {
      .instance = instance,
      // TODO Add a fallback physical device option.
      // .physicalDeviceType = VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU,
      .physicalDeviceType = VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU,
      .queueRequirements = svec_init(evkDeviceQueueRequirement, {
          { VK_QUEUE_GRAPHICS_BIT, 1 },
          { VK_QUEUE_COMPUTE_BIT , 1 },
      }),
      .deviceExtensions = svec_init(evstring, {
        evstr("VK_KHR_swapchain"),
        evstr("VK_KHR_dynamic_rendering"),
        evstr("VK_KHR_depth_stencil_resolve"),
        evstr("VK_KHR_create_renderpass2"),
        evstr("VK_KHR_synchronization2"),
        evstr("VK_KHR_buffer_device_address"),
        evstr("VK_EXT_descriptor_indexing"),
        evstr("VK_EXT_descriptor_buffer"),
      }),
  });

  if(device.vk == EV_INVALID(VkDevice))
  {
    puts("Couldn't create a VkDevice");
    goto DeviceCreationFailed;
  }

  puts("Logical Device created successfully.");
  VkQueue graphicsQueue;
  vkGetDeviceQueue(device.vk, device.queueFamilies[VK_QUEUE_GRAPHICS_BIT].familyIndex, 0, &graphicsQueue);

  if (!glfwInit())
  {
    puts("GLFW Initialization failed.");
    goto GLFWInitFailed;
  }

  glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
  GLFWwindow* window = glfwCreateWindow(1024,1024, "evk", NULL, NULL);
  if(!window)
  {
    puts("Window Creation Failed.");
    goto WindowCreationFailed;
  }

  VkSurfaceKHR surface;
  VkResult err = glfwCreateWindowSurface(instance.vk, window, NULL, &surface);
  if (err)
  {
    goto VKSurfaceCreationFailed;
    puts("Surface creation failed.");
  }

  evkGPUAllocator allocator = evkGPUCreateAllocator(device);

  evkSwapChain swapChain = evkCreateSwapChain((evkSwapChainCreateInfo){
    .device = device,
    .surface = surface,
    .width = width,
    .height = height,
    .imageCount = 3,
  });

  evkCommandPool commandPool = evkCreateCommandPool((evkCommandPoolCreateInfo) { 
    .device = device, 
    .queueFlags = VK_QUEUE_GRAPHICS_BIT,
    .poolFlags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT
  });

  vec(evkCommandBuffer) commandBuffers = evkAllocateCommandBuffers(device, commandPool, vec_len(&swapChain.images), true);

  evkShaderCompiler compiler = evkCreateShaderCompiler();

  evkShader vertShader = evkInitShaderFromFile(device, compiler, "shaders/tri.vert");
  evkShader fragShader = evkInitShaderFromFile(device, compiler, "shaders/tri.frag");
  
  evkDestroyShaderCompiler(compiler);

  evkColorAttachment colorAttachment0 = {
    swapChain.surfaceFormat.format,
    EV_DEFAULT(VkPipelineColorBlendAttachmentState)
  };

  // TODO Get this from shader reflection data
  // evkDescriptorSetLayout setLayout_0 = evkCreateDescriptorSetLayout(
  //   &device, svec_init(evkDescriptorBinding, {
  //     { 
  //       .name = evstr("positions"), 
  //       .binding = 0, 
  //       .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 
  //       .descriptorCount = 1, 
  //       .stageFlags = VK_SHADER_STAGE_ALL_GRAPHICS, 
  //     },
  //   })
  // );
  // evkDescriptorSetLayout setLayout_0 = evkCreateDescriptorSetLayoutFromBindings(&device, vertShader.reflect.bindings);
  evkDescriptorSetLayout setLayout_0 = evkCreateDescriptorSetLayoutFromShaders(&device, svec_init(evkShader, {vertShader, fragShader}));

  evkPipelineCreateInfo pipelineCreateInfo = EV_DEFAULT(evkPipelineCreateInfo,
    dynamicStates = svec_init(VkDynamicState, { 
      VK_DYNAMIC_STATE_VIEWPORT, 
      VK_DYNAMIC_STATE_SCISSOR,
    }),
    shaderStages = svec_init(evkShader, {
      vertShader, 
      fragShader,
    }),
    colorAttachments = svec_init(evkColorAttachment, {
      colorAttachment0,
    }),
    viewportCountOverride = 1,

    vertexBufferLayouts = svec_init(evkVertexBufferLayout, {
      {{ // VB #0
        { EVK_VERTEX_ATTRIBUTE_POSITION, EVK_FMT_FLOAT32, 2 },
      }},
    }),

    setLayouts = svec_init(evkDescriptorSetLayout, {setLayout_0}),
  );

  evkDescriptorSet set_0 = evkCreateDescriptorSet(&(evkDescriptorSetCreateInfo){
    .device = &device, 
    .allocator = &allocator, 
    .layout = &setLayout_0
  });


  // // if stageflags is 0, it's all graphics
  // // if descriptor count is 0, it's 1
  // // if
  // evkDescriptorSetLayout set_0 = evkCreateDescriptorSetLayout(&device, svec_init(evkDescriptorBinding, {
  //     { "positions", VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER },
  //   })
  // );

  evkBuffer vertBuf = evkCreateBuffer(&device, (evkBufferCreateInfo) {
    .queueFamilyIndices = svec_init(u32, {device.queueFamilies[VK_QUEUE_GRAPHICS_BIT].familyIndex}),
    .sizeInBytes = 6 * 4,
    .usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
    .allocationCreateInfo = {
      .allocationFlags = EVK_GPU_ALLOCATION_CREATE_MAPPED_BIT | EVK_GPU_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
      .allocator = allocator,
    },
    .exclusive = true,
  });

  evkBuffer uniBuf = evkCreateBuffer(&device, (evkBufferCreateInfo) {
    .sizeInBytes = 12 * 4,
    .usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
    .allocationCreateInfo = {
      .allocationFlags = EVK_GPU_ALLOCATION_CREATE_MAPPED_BIT | EVK_GPU_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
      .allocator = allocator,
    },
    .exclusive = true,
  });

  float* vert_arr = (float*)vertBuf.allocData.allocationInfo.vma.pMappedData;
  vert_arr[0] =  0.0f; vert_arr[1] = -0.5f;
  vert_arr[2] =  0.5f; vert_arr[3] =  0.5f;
  vert_arr[4] = -0.5f; vert_arr[5] =  0.5f;

  float* uni_arr = (float*)uniBuf.allocData.allocationInfo.vma.pMappedData;
  uni_arr[0] =  0.0f; uni_arr[1] = -0.5f; uni_arr[2]  = 0.0f; uni_arr[3]  = 1.0f;
  uni_arr[4] =  0.5f; uni_arr[5] =  0.5f; uni_arr[6]  = 0.0f; uni_arr[7]  = 1.0f;
  uni_arr[8] = -0.5f; uni_arr[9] =  0.5f; uni_arr[10] = 0.0f; uni_arr[11] = 1.0f;

  evkSetDescriptor(&set_0, evstr("vertexData"), &uniBuf);

  VkViewport viewport = EV_DEFAULT(VkViewport, width=width, height=height);
  VkRect2D scissor = EV_DEFAULT(VkRect2D,extent.width=width, extent.height=height);

  evkPipeline graphicsPipeline = evkCreatePipeline(device, pipelineCreateInfo);

  u32 imageCount = vec_len(&swapChain.images);

  VkFence drawFence = evkCreateFence(device, false);
  vec(VkSemaphore) imageAcquiredSemaphores = vec_init(VkSemaphore);
  vec(VkSemaphore) drawFinishedSemaphores = vec_init(VkSemaphore);
  for(u32 i = 0; i < imageCount; i++)
  {
    imageAcquiredSemaphores[i] = evkCreateSemaphore(device);
    drawFinishedSemaphores[i] = evkCreateSemaphore(device);
  }
  // VkSemaphore imageAcquiredSemaphore = evkCreateSemaphore(device);
  // VkSemaphore drawFinishedSemaphore = evkCreateSemaphore(device);

  i32 imageIdx = -1;
  while(!glfwWindowShouldClose(window))
  {
    imageIdx = (imageIdx + 1) % imageCount;
    u32 swapChainImageIdx;
    EVK_ASSERT(vkAcquireNextImageKHR(device.vk, swapChain.vk, UInt64.MAX, imageAcquiredSemaphores[imageIdx], VK_NULL_HANDLE, &swapChainImageIdx));

    evkCommandBuffer cmdbuf = commandBuffers[imageIdx];

    VkRenderingAttachmentInfoKHR colorAttachment = EV_DEFAULT(VkRenderingAttachmentInfoKHR,
      imageView = swapChain.imageViews[swapChainImageIdx].vk,
      loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
      storeOp= VK_ATTACHMENT_STORE_OP_STORE,
    );

    VkRenderingInfo renderingInfo = EV_DEFAULT(VkRenderingInfo,
      renderArea.extent = ((VkExtent2D){width, height}),
      pColorAttachments = &colorAttachment,
      colorAttachmentCount = 1,
    );

    evkBeginPrimaryCommandBuffer(&cmdbuf);
    {
      VkDeviceSize offset = 0;
      vkCmdBindVertexBuffers(cmdbuf.vk, 0, 1, &vertBuf.vk, &offset);

      evkCmdBindDescriptorSets(&cmdbuf, &graphicsPipeline, svec_init(evkDescriptorSet, { set_0 }), svec_init(u32, { 0 }));

      vkCmdSetScissor(cmdbuf.vk, 0, 1, &scissor);
      vkCmdSetViewport(cmdbuf.vk, 0, 1, &viewport);

      vkCmdBeginRenderingKHR(cmdbuf.vk, &renderingInfo);

      evkCmdBindPipeline(&cmdbuf, &graphicsPipeline);
      vkCmdDraw(cmdbuf.vk, 3, 1, 0, 0);

      vkCmdEndRenderingKHR(cmdbuf.vk);

      VkImageMemoryBarrier imageMemoryBarrier = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        .image = swapChain.images[swapChainImageIdx].vk,
        .newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
        .subresourceRange = {
          .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
          .layerCount = 1,
          .levelCount = 1,
        },
      };

      vkCmdPipelineBarrier(cmdbuf.vk, 0, 0, 0, 0, NULL, 0, NULL, 1, &imageMemoryBarrier);
    }
    evkEndCommandBuffer(&cmdbuf);

    VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};

    VkSubmitInfo submitInfo = {
      .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
      .pCommandBuffers = &cmdbuf.vk, // TODO This won't work with more than a single cmdbuf
      .commandBufferCount = 1,
      .waitSemaphoreCount = 1,
      .pWaitSemaphores = &imageAcquiredSemaphores[imageIdx],
      .pWaitDstStageMask = waitStages,
      .signalSemaphoreCount = 1,
      .pSignalSemaphores = &drawFinishedSemaphores[swapChainImageIdx],
    };

    vkQueueSubmit(graphicsQueue, 1, &submitInfo, drawFence);
    vkWaitForFences(device.vk, 1, &drawFence, VK_TRUE, UInt64.MAX);
    vkResetFences(device.vk, 1, &drawFence);

    VkPresentInfoKHR presentInfo = {
      .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
      .waitSemaphoreCount = 1,
      .pWaitSemaphores = &drawFinishedSemaphores[swapChainImageIdx],
      .swapchainCount = 1,
      .pSwapchains = &swapChain.vk,
      .pImageIndices = &swapChainImageIdx,
    };
    vkQueuePresentKHR(graphicsQueue, &presentInfo);

    vkResetCommandBuffer(cmdbuf.vk,VK_COMMAND_BUFFER_RESET_RELEASE_RESOURCES_BIT);

    // Main Loop
    glfwPollEvents();
  }

  vkQueueWaitIdle(graphicsQueue);

  vkDestroyFence(device.vk, drawFence, NULL);
  for(u32 i = 0; i < imageCount; i++)
  {
    vkDestroySemaphore(device.vk, imageAcquiredSemaphores[i], NULL);
    vkDestroySemaphore(device.vk, drawFinishedSemaphores[i], NULL);
  }
  // vkDestroySemaphore(device.vk, imageAcquiredSemaphore, NULL);
  // vkDestroySemaphore(device.vk, drawFinishedSemaphore, NULL);

  evkDestroyPipeline(graphicsPipeline);

  evkDestroyDescriptorSet(&device, &set_0);
  evkDestroyDescriptorSetLayout(&device, &setLayout_0);

  evkDestroyShader(device, vertShader);
  evkDestroyShader(device, fragShader);

  evkFreeCommandBuffers(device, commandPool, commandBuffers);
  evkDestroyCommandPool(device, commandPool);

  evkDestroySwapChain(device, swapChain);

// SwapchainCreationFailed:
  vkDestroySurfaceKHR(instance.vk, swapChain.surface, NULL);

VKSurfaceCreationFailed:
WindowCreationFailed:
  glfwTerminate();

GLFWInitFailed:
  evkDestroyDevice(device);

DeviceCreationFailed:
  evkDestroyInstance(instance);
InstanceCreationFailed:
  return 0;
}