#include "pipeline.h"
#include "cglm/affine.h"
#include "cglm/mat4.h"
#include "command.h"
#include "stdio.h"
#include "stdlib.h"
#include "string.h"
#include "vk_mem_alloc.h"
#include "vulkan/vulkan_core.h"

VkShaderModule load_shader_file(const char* path, VkDevice device) {
  FILE* file;
  file = fopen(path, "rb");
  if(file == 0) {
    return VK_NULL_HANDLE;
  }

  int result = fseek(file, 0, SEEK_END);
  if(result != 0) {
    return VK_NULL_HANDLE;
  }

  long buffer_size = ftell(file);

  result = fseek(file, 0, SEEK_SET);
  if(result != 0) {
    return VK_NULL_HANDLE;
  }

  char * buffer = malloc(buffer_size);
  if(buffer == 0) {
    return VK_NULL_HANDLE;
  }

  size_t read = fread(buffer, 1, buffer_size, file);

  VkShaderModuleCreateInfo shader_info = {
    .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
    .codeSize = read,
    .pCode = (uint32_t*)buffer,
  };

  VkShaderModule shader;
  result = vkCreateShaderModule(device, &shader_info, 0, &shader);
  free(buffer);
  if(result != VK_SUCCESS) {
    return VK_NULL_HANDLE;
  }

  return shader;
}

VkResult create_ui_pipeline(
    VkDevice device,
    VkRenderPass render_pass,
    VkPipelineShaderStageCreateInfo* shader_stages,
    uint32_t shader_stage_count,
    VkPipelineVertexInputStateCreateInfo input_info,
    VkPipelineLayoutCreateInfo layout_info,
    VkPipelineInputAssemblyStateCreateInfo input_assembly_info,
    GraphicsPipeline* pipeline) {
  VkResult result;

  result = vkCreatePipelineLayout(device, &layout_info, 0, &pipeline->layout);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkViewport viewport = {
    .x = 0.0f,
    .y = 0.0f,
    .width = (float)(100),
    .height = (float)(100),
    .minDepth = 0.0f,
    .maxDepth = 1.0f,
  };

  VkRect2D scissor = {
    .offset = {
      .x = 0,
      .y = 0,
    },
    .extent = {
      .width = 100,
      .height = 100,
    },
  };

  VkPipelineViewportStateCreateInfo viewport_state = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
    .viewportCount = 1,
    .pViewports = &viewport,
    .scissorCount = 1,
    .pScissors = &scissor,
  };

  VkPipelineRasterizationStateCreateInfo raster_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
    .depthClampEnable = VK_FALSE,
    .rasterizerDiscardEnable = VK_FALSE,
    .polygonMode = VK_POLYGON_MODE_FILL,
    .lineWidth = 1.0f,
    .cullMode = VK_CULL_MODE_BACK_BIT,
    .frontFace = VK_FRONT_FACE_CLOCKWISE,
    .depthBiasEnable = VK_FALSE,
    .depthBiasConstantFactor = 0.0f,
    .depthBiasClamp = 0.0f,
    .depthBiasSlopeFactor = 0.0f,
  };

  VkPipelineMultisampleStateCreateInfo multisample_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
    .sampleShadingEnable = VK_FALSE,
    .rasterizationSamples = VK_SAMPLE_COUNT_1_BIT,
    .minSampleShading = 1.0f,
    .pSampleMask = 0,
    .alphaToCoverageEnable = VK_FALSE,
    .alphaToOneEnable = VK_FALSE,
  };

  VkPipelineDepthStencilStateCreateInfo depth_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
    .depthTestEnable = VK_TRUE,
    .depthWriteEnable = VK_TRUE,
    .depthCompareOp = VK_COMPARE_OP_LESS,
    .depthBoundsTestEnable = VK_FALSE,
    .maxDepthBounds = 1.0f,
    .minDepthBounds = 0.0f,
    .stencilTestEnable = VK_FALSE,
    .front = {},
    .back = {},
  };

  VkPipelineColorBlendAttachmentState color_blend_attachment = {
    .colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT,
    .blendEnable = VK_TRUE,
    .srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA,
    .dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
    .colorBlendOp = VK_BLEND_OP_ADD,
    .srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE,
    .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
    .alphaBlendOp = VK_BLEND_OP_ADD,
  };

  VkPipelineColorBlendStateCreateInfo color_blend_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
    .logicOpEnable = VK_FALSE,
    .logicOp = VK_LOGIC_OP_COPY,
    .attachmentCount = 1,
    .pAttachments = &color_blend_attachment,
    .blendConstants[0] = 0.0f,
    .blendConstants[1] = 0.0f,
    .blendConstants[2] = 0.0f,
    .blendConstants[3] = 0.0f,
  };

  VkDynamicState dynamic_states[] = {
    VK_DYNAMIC_STATE_VIEWPORT,
    VK_DYNAMIC_STATE_SCISSOR,
  };

  uint32_t dynamic_state_count = sizeof(dynamic_states)/sizeof(VkDynamicState);

  VkPipelineDynamicStateCreateInfo dynamic_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    .dynamicStateCount = dynamic_state_count,
    .pDynamicStates = dynamic_states,
  };

  VkGraphicsPipelineCreateInfo draw_pipeline_info = {
    .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
    .stageCount = shader_stage_count,
    .pStages = shader_stages,
    .pVertexInputState = &input_info,
    .pInputAssemblyState = &input_assembly_info,
    .pViewportState = &viewport_state,
    .pRasterizationState = &raster_info,
    .pColorBlendState = &color_blend_info,
    .pDynamicState = &dynamic_info,
    .pDepthStencilState = &depth_info,
    .pMultisampleState = &multisample_info,
    .layout = pipeline->layout,
    .renderPass = render_pass,
    .subpass = 1,
    .basePipelineHandle = VK_NULL_HANDLE,
    .basePipelineIndex = -1,
  };

  result = vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &draw_pipeline_info, 0, &pipeline->pipeline);
  if(result != VK_SUCCESS) {
    return result;
  }
  return VK_SUCCESS;
}

VkResult create_ui_colored_rect_pipeline(VkDevice device, VkRenderPass render_pass, VkDescriptorSetLayout ui_descriptor_layout, GraphicsPipeline* pipeline) {
  VkShaderModule vert_shader = load_shader_file("shader_src/ui_colored_rect.vert.spv", device);
  if(vert_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }
  VkShaderModule frag_shader = load_shader_file("shader_src/ui_colored_rect.frag.spv", device);
  if(frag_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }
  VkPipelineShaderStageCreateInfo shader_stages[] = {
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_VERTEX_BIT,
      .module = vert_shader,
      .pName = "main",
    },
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
      .module = frag_shader,
      .pName = "main",
    },
  };

  VkVertexInputBindingDescription bindings[] = {
    {
      .binding = 0,
      .stride = sizeof(vec2),
      .inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
    },
    {
      .binding = 1,
      .stride = sizeof(ColoredRect),
      .inputRate = VK_VERTEX_INPUT_RATE_INSTANCE,
    },
  };

  VkVertexInputAttributeDescription attributes[] = {
    {
      .binding = 0,
      .location = 0,
      .format = VK_FORMAT_R32G32_SFLOAT,
      .offset = 0,
    },
    {
      .binding = 1,
.location = 1,
      .format = VK_FORMAT_R32G32B32_SFLOAT,
      .offset = offsetof(ColoredRect, pos),
    },
    {
      .binding = 1,
      .location = 2,
      .format = VK_FORMAT_R32G32_SFLOAT,
      .offset = offsetof(ColoredRect, size),
    },
    {
      .binding = 1,
      .location = 3,
      .format = VK_FORMAT_R32G32B32A32_SFLOAT,
      .offset = offsetof(ColoredRect, color),
    },
  };

VkPipelineVertexInputStateCreateInfo input_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
    .pVertexBindingDescriptions = bindings,
    .vertexBindingDescriptionCount = sizeof(bindings)/sizeof(VkVertexInputBindingDescription),
    .pVertexAttributeDescriptions = attributes,
    .vertexAttributeDescriptionCount = sizeof(attributes)/sizeof(VkVertexInputAttributeDescription),
  };

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .setLayoutCount = 1,
    .pSetLayouts = &ui_descriptor_layout,
  };

  VkPipelineInputAssemblyStateCreateInfo input_assembly_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
    .topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
    .primitiveRestartEnable = VK_FALSE,
  };

  VkResult result = create_ui_pipeline(device, render_pass, shader_stages, sizeof(shader_stages)/sizeof(VkPipelineShaderStageCreateInfo), input_info, layout_info, input_assembly_info, pipeline);
  if(result != VK_SUCCESS) {
    return result;
  }

  return VK_SUCCESS;
}

VkResult create_ui_text_pipeline(VkDevice device, VkRenderPass render_pass, VkDescriptorSetLayout ui_descriptor_layout, VkDescriptorSetLayout font_layout, GraphicsPipeline* pipeline, ComputePipeline* compute) {
  VkResult result;
  VkShaderModule compute_shader = load_shader_file("shader_src/ui_text.comp.spv", device);
  if(compute_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }

  VkPipelineShaderStageCreateInfo compute_shader_stage = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
    .stage = VK_SHADER_STAGE_COMPUTE_BIT,
    .pName = "main",
    .module = compute_shader,
  };

  VkPushConstantRange push_constant = {
    .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT | VK_SHADER_STAGE_VERTEX_BIT,
    .size = 8,
  };

  VkDescriptorSetLayout compute_descriptors[] = {font_layout};

  VkPipelineLayoutCreateInfo compute_layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pushConstantRangeCount = 1,
    .pPushConstantRanges = &push_constant,
    .setLayoutCount = sizeof(compute_descriptors)/sizeof(VkDescriptorSetLayout),
    .pSetLayouts = compute_descriptors,
  };

  result = vkCreatePipelineLayout(device, &compute_layout_info, NULL, &compute->layout);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkComputePipelineCreateInfo compute_pipeline_info = {
    .sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
    .layout = compute->layout,
    .stage = compute_shader_stage,
  };

  result = vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &compute_pipeline_info, NULL, &compute->pipeline);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkShaderModule vert_shader = load_shader_file("shader_src/ui_text.vert.spv", device);
  if(vert_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }
  VkShaderModule frag_shader = load_shader_file("shader_src/ui_text.frag.spv", device);
  if(frag_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }
  
  // TODO: add a compute stage before the shader stages that sets up the draw buffer and commands for the characters based off of strings + positions
  // 1. Reserve a buffer for array of draw command params + Char on the GPU
  // 2. Reserve a buffer for Text on the GPU
  VkPipelineShaderStageCreateInfo shader_stages[] = {
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_VERTEX_BIT,
      .module = vert_shader,
      .pName = "main",
    },
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
      .module = frag_shader,
      .pName = "main",
    },
  };

  VkVertexInputBindingDescription bindings[] = {
    {
      .binding = 0,
      .stride = sizeof(vec2),
      .inputRate = VK_VERTEX_INPUT_RATE_VERTEX,
    },
  };
  VkVertexInputAttributeDescription attributes[] = {
    {
      .binding = 0,
      .location = 0,
      .format = VK_FORMAT_R32G32_SFLOAT,
      .offset = 0,
    },
  };

  VkPipelineVertexInputStateCreateInfo input_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
    .pVertexBindingDescriptions = bindings,
    .vertexBindingDescriptionCount = sizeof(bindings)/sizeof(VkVertexInputBindingDescription),
    .pVertexAttributeDescriptions = attributes,
    .vertexAttributeDescriptionCount = sizeof(attributes)/sizeof(VkVertexInputAttributeDescription),
  };
 
  VkDescriptorSetLayout all_layouts[] = {ui_descriptor_layout, font_layout};

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .setLayoutCount = sizeof(all_layouts)/sizeof(VkDescriptorSetLayout),
    .pSetLayouts = all_layouts,
    .pPushConstantRanges = &push_constant,
    .pushConstantRangeCount = 1,
  };

  VkPipelineInputAssemblyStateCreateInfo input_assembly_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
    .topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
    .primitiveRestartEnable = VK_FALSE,
  };

  result = create_ui_pipeline(device, render_pass, shader_stages, sizeof(shader_stages)/sizeof(VkPipelineShaderStageCreateInfo), input_info, layout_info, input_assembly_info, pipeline);
  if(result != VK_SUCCESS) {
    return result;
  }

  return VK_SUCCESS;
}

VkResult create_font_descriptor_pools(VkDevice device, uint32_t max_sets, VkDescriptorPool* font_pool, VkDescriptorSetLayout* font_layout) {
  VkResult result;
  VkDescriptorPoolSize font_pool_sizes[] = {
    {
      .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = 1,
    },
    {
      .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
      .descriptorCount = 1,
    },
  };

  VkDescriptorPoolCreateInfo font_pool_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .pPoolSizes = font_pool_sizes,
    .poolSizeCount = sizeof(font_pool_sizes)/sizeof(VkDescriptorPoolSize),
    .maxSets = max_sets,
  };

  result = vkCreateDescriptorPool(device, &font_pool_info, NULL, font_pool);
  if(result != VK_SUCCESS) {
    return result;
  }
  
  VkDescriptorSetLayoutBinding font_descriptor_bindings[] = {
    {
      .binding = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = 1,
      .stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_COMPUTE_BIT,
    },
    {
      .binding = 1,
      .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
      .descriptorCount = 1,
      .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
    },
  };

  VkDescriptorSetLayoutCreateInfo font_descriptor_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .pBindings = font_descriptor_bindings,
    .bindingCount = sizeof(font_descriptor_bindings)/sizeof(VkDescriptorSetLayoutBinding),
  };

  result = vkCreateDescriptorSetLayout(device, &font_descriptor_info, NULL, font_layout);
  if(result != VK_SUCCESS) {
    return result;
  }

  return VK_SUCCESS;
}

VkResult load_font(VkDevice device, VmaAllocator allocator, VkDescriptorSetLayout layout, VkDescriptorPool pool,VkCommandPool transfer_pool, Queue transfer_queue, FT_Library library, const char* ttf_file, uint32_t size, VkBool32 antialias, uint32_t** charmap, FontDescriptor* descriptor) {  
  FT_Face face;

  int error;
  error = FT_New_Face(library, ttf_file, 0, &face);
  if(error != FT_Err_Ok) {
    return VK_ERROR_UNKNOWN;
  }

  error = FT_Set_Pixel_Sizes(face, 0, size);
  if(error != FT_Err_Ok) {
    return VK_ERROR_UNKNOWN;
  }

  uint32_t* tmp_charmap = malloc(sizeof(uint32_t)*face->num_glyphs);
  SymbolInfo* symbols = malloc(sizeof(SymbolInfo)*face->num_glyphs);
  FontUniform uniform;

  uint32_t glyph_index;
  uint32_t max_height = 0;
  uint32_t max_width = 0;
  uint32_t symbol_count = 0;
  uint32_t c;
  FT_Int32 load_flags = FT_LOAD_RENDER;
  if(antialias == VK_TRUE) {
    load_flags += FT_RENDER_MODE_NORMAL;
  } else {
    load_flags += FT_RENDER_MODE_MONO;
  }

  // TODO: worry about variants(that's why num_glyphs doesn't match symbol_count)
  c = FT_Get_First_Char(face, &glyph_index);
  for(uint32_t i = 0; i < face->num_glyphs; i++) {
    if(glyph_index == 0) {
      break;
    }
    FT_Load_Glyph(face, glyph_index, load_flags);
    uint32_t width = face->glyph->bitmap.width;
    uint32_t height = face->glyph->bitmap.rows;
    tmp_charmap[i] = c;
    symbols[i].width = width;
    symbols[i].height = height;
    symbols[i].left = face->glyph->bitmap_left;
    symbols[i].top = face->glyph->bitmap_top;
    symbols[i].advance = face->glyph->advance.x*16/face->units_per_EM;
    max_width = width > max_width ? width : max_width; 
    max_height = height > max_height ? height : max_height;
    symbol_count += 1;
    c = FT_Get_Next_Char(face, c, &glyph_index);
  }

  uniform.width = max_width;
  uniform.height = max_height;
  uniform.num_symbols = symbol_count;
  uint32_t image_size = max_width*max_height*sizeof(uint32_t);
  uint32_t* images = malloc(image_size*symbol_count);
  memset(images, 0x00, image_size*symbol_count);
  *charmap = malloc(sizeof(uint32_t)*symbol_count);
  memcpy(*charmap, tmp_charmap, sizeof(uint32_t)*symbol_count);
  free(tmp_charmap);

  for(uint32_t i = 0; i < symbol_count; i++) {
    glyph_index = FT_Get_Char_Index(face, (*charmap)[i]);
    FT_Load_Glyph(face, glyph_index, load_flags);
    for(uint32_t y = 0; y < face->glyph->bitmap.rows; y++) {
      for(uint32_t x = 0; x < face->glyph->bitmap.width; x++) {
        uint32_t level = face->glyph->bitmap.buffer[y*face->glyph->bitmap.width+x];
        level = ((level * 0xFFFFFF) / 256) << 8;
        images[max_width*max_height*i + max_width*y + x] = 0x000000FF + level;
      }
    }
  }


  error = FT_Done_Face(face);
  if(error != FT_Err_Ok) {
    return VK_ERROR_UNKNOWN;
  }

  VkResult result;
  VkDescriptorSetAllocateInfo set_allocate_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = &layout,
    .descriptorSetCount = 1,
    .descriptorPool = pool,
  };

  result = vkAllocateDescriptorSets(device, &set_allocate_info, &descriptor->set);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo symbol_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = sizeof(SymbolInfo)*uniform.num_symbols,
  };
  VmaAllocationCreateInfo symbol_memory_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };
  result = vmaCreateBuffer(allocator, &symbol_buffer_info, &symbol_memory_info, &descriptor->symbols, &descriptor->symbol_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo uniform_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
    .size = sizeof(FontUniform),
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  VmaAllocationCreateInfo uniform_memory_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };

  result = vmaCreateBuffer(allocator, &uniform_buffer_info, &uniform_memory_info, &descriptor->uniform, &descriptor->uniform_memory, NULL);
  VkImageCreateInfo image_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
    .extent.depth  = 1,
    .extent.width  = uniform.width,
    .extent.height = uniform.height,
    .mipLevels = 1,
    .arrayLayers = uniform.num_symbols,
    .format = VK_FORMAT_R8G8B8A8_SRGB,
    .tiling = VK_IMAGE_TILING_OPTIMAL,
    .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
    .samples = VK_SAMPLE_COUNT_1_BIT,
    .imageType = VK_IMAGE_TYPE_2D,
  };

  VmaAllocationCreateInfo image_memory_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };

  result = vmaCreateImage(allocator, &image_info, &image_memory_info, &descriptor->image, &descriptor->image_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo staging_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = sizeof(FontUniform) + image_size*uniform.num_symbols + sizeof(SymbolInfo)*uniform.num_symbols,
  };

  VmaAllocationCreateInfo staging_memory_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };

  VkBuffer staging_buffer;
  VmaAllocation staging_memory;
  result = vmaCreateBuffer(allocator, &staging_info, &staging_memory_info, &staging_buffer, &staging_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  void* mapped_staging;
  result = vmaMapMemory(allocator, staging_memory, &mapped_staging);
  if(result != VK_SUCCESS) {
    return result;
  }
  memcpy(mapped_staging + image_size*uniform.num_symbols + sizeof(FontUniform), symbols, sizeof(SymbolInfo)*uniform.num_symbols);
  VkBufferDeviceAddressInfo address_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = descriptor->symbols,
  };
	uniform.symbol_list = vkGetBufferDeviceAddress(device, &address_info);
  memcpy(mapped_staging + image_size*uniform.num_symbols, &uniform, sizeof(FontUniform));
  memcpy(mapped_staging, images, image_size*uniform.num_symbols);
  vmaUnmapMemory(allocator, staging_memory);
  free(images);
  free(symbols);

  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);

  VkBufferCopy uniform_copy_info = {
    .size = sizeof(FontUniform),
    .srcOffset = image_size*uniform.num_symbols,
    .dstOffset = 0,
  };
  vkCmdCopyBuffer(command_buffer, staging_buffer, descriptor->uniform, 1, &uniform_copy_info);

  VkBufferCopy symbol_copy_info = {
    .size = sizeof(SymbolInfo)*uniform.num_symbols,
    .srcOffset = image_size*uniform.num_symbols + sizeof(FontUniform),
    .dstOffset = 0,
  };
  vkCmdCopyBuffer(command_buffer, staging_buffer, descriptor->symbols, 1, &symbol_copy_info);

  VkImageMemoryBarrier first_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = descriptor->image,
    .oldLayout = VK_IMAGE_LAYOUT_UNDEFINED,
    .newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
    .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
    .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
    .subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
    .subresourceRange.levelCount = 1,
    .subresourceRange.layerCount = uniform.num_symbols,
    .srcAccessMask = 0,
    .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
  };
  vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, NULL, 0, NULL, 1, &first_barrier);

  VkBufferImageCopy image_copy = {
    .imageSubresource.layerCount = uniform.num_symbols,
    .imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
    .imageExtent = image_info.extent,
  };
  vkCmdCopyBufferToImage(command_buffer, staging_buffer, descriptor->image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_copy);

  VkImageMemoryBarrier second_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = descriptor->image,
    .oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
    .newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
    .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
    .subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
    .subresourceRange.levelCount = 1,
    .subresourceRange.layerCount = uniform.num_symbols,
    .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT,
    .dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
  };
  vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, NULL, 0, NULL, 1, &second_barrier);

  result = command_end_single(device, command_buffer, transfer_pool, transfer_queue);
  if(result != VK_SUCCESS) {
    return result;
  }
  vmaDestroyBuffer(allocator, staging_buffer, staging_memory);

  VkImageViewCreateInfo view_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
    .image = descriptor->image,
    .viewType = VK_IMAGE_VIEW_TYPE_2D_ARRAY,
    .format = VK_FORMAT_R8G8B8A8_SRGB,
    .subresourceRange = {
      .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
      .layerCount = uniform.num_symbols,
      .baseMipLevel = 0,
      .levelCount = 1,
      .baseArrayLayer = 0,
    },
  };
  result = vkCreateImageView(device, &view_info, NULL, &descriptor->view);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkSamplerCreateInfo sampler_info = {
    .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
    .magFilter = VK_FILTER_NEAREST,
    .minFilter = VK_FILTER_LINEAR,
    .addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT,
    .addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT,
    .addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT,
  };
  result = vkCreateSampler(device, &sampler_info, NULL, &descriptor->sampler);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkDescriptorBufferInfo desc_uniform_info = {
    .offset = 0,
    .range = sizeof(FontUniform),
    .buffer = descriptor->uniform,
  };

  VkDescriptorImageInfo desc_image_info = {
    .sampler = descriptor->sampler,
    .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    .imageView = descriptor->view,
  };

  VkWriteDescriptorSet descriptor_writes[] = {
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = descriptor->set,
      .dstBinding = 0,
      .dstArrayElement = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = 1,
      .pBufferInfo = &desc_uniform_info,
    },
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = descriptor->set,
      .dstBinding = 1,
      .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
      .descriptorCount = 1,
      .pImageInfo = &desc_image_info,
    }
  };

  vkUpdateDescriptorSets(device, 2, descriptor_writes, 0, NULL);
  return VK_SUCCESS;
}

VkResult create_ui_descriptor_set(VkDevice device, VmaAllocator allocator, VkExtent2D swapchain_extent, vec2 window_scale, VkDescriptorSetLayout* ui_descriptor_layout, VkDescriptorPool* ui_descriptor_pool, VkDescriptorSet* ui_descriptor_set, VmaAllocation* ui_descriptor_memory, VkBuffer* ui_descriptor_buffer) { 
  VkDescriptorSetLayoutBinding ui_descriptor_bindings[] = {
    {
      .binding = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
      .descriptorCount = 1,
      .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
    }
  };

  VkDescriptorSetLayoutCreateInfo ui_descriptor_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .pBindings = ui_descriptor_bindings,
    .bindingCount = sizeof(ui_descriptor_bindings)/sizeof(VkDescriptorSetLayoutBinding),
  }; 

  VkResult result;
  result = vkCreateDescriptorSetLayout(device, &ui_descriptor_info, NULL, ui_descriptor_layout);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkDescriptorPoolSize pool_size = {
    .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
    .descriptorCount = 1,
  };

  VkDescriptorPoolCreateInfo ui_pool_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .pPoolSizes = &pool_size,
    .poolSizeCount = 1,
    .maxSets = 1,
  };

  result = vkCreateDescriptorPool(device, &ui_pool_info, NULL, ui_descriptor_pool);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkDescriptorSetAllocateInfo ui_descriptor_allocate_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = ui_descriptor_layout,
    .descriptorSetCount = 1,
    .descriptorPool = *ui_descriptor_pool,
  };

  result = vkAllocateDescriptorSets(device, &ui_descriptor_allocate_info, ui_descriptor_set);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo ui_uniform_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
    .size = sizeof(UIUniform),
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  VmaAllocationCreateInfo ui_uniform_memory_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };

  result = vmaCreateBuffer(allocator, &ui_uniform_buffer_info, &ui_uniform_memory_info, ui_descriptor_buffer, ui_descriptor_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  void* mapped;
  result = vmaMapMemory(allocator, *ui_descriptor_memory, &mapped);
  if(result != VK_SUCCESS) {
    return result;
  }
  UIUniform ui_uniform;
  vec3 screen_offset = {-1.0, -1.0, 0.0};
  vec3 screen_scale = {1.0/(float)swapchain_extent.width*window_scale[0], 1.0/(float)swapchain_extent.height*window_scale[1], 1.0};
  glm_mat4_identity(ui_uniform.screen);
  glm_translate(ui_uniform.screen, screen_offset);
  glm_scale(ui_uniform.screen, screen_scale);
  memcpy(mapped, &ui_uniform, sizeof(ui_uniform));
  vmaUnmapMemory(allocator, *ui_descriptor_memory);

  VkDescriptorBufferInfo ui_uniform_info = {
    .offset = 0,
    .range = sizeof(ui_uniform),
    .buffer = *ui_descriptor_buffer,
  };

  VkWriteDescriptorSet ui_uniform_write = {
    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
    .dstSet = *ui_descriptor_set,
    .dstBinding = 0,
    .dstArrayElement = 0,
    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
    .descriptorCount = 1,
    .pBufferInfo = &ui_uniform_info,
  };

  vkUpdateDescriptorSets(device, 1, &ui_uniform_write, 0, NULL);
  return VK_SUCCESS;
}

VkResult create_ui_rect_buffer(VkDevice device, Queue transfer_queue, VkCommandPool transfer_pool, VmaAllocator allocator, struct RectBuffer* rect) {

  uint32_t vertex_buffer_size = 4 * sizeof(vec2);
  uint32_t index_buffer_size = 6 * sizeof(uint32_t);

  // Create temp buffer
  VkBufferCreateInfo temp_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = vertex_buffer_size + index_buffer_size,
  };

  VmaAllocationCreateInfo temp_allocation_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };

  VkBuffer temp_buffer;
  VmaAllocation temp_buffer_memory;
  VkResult result;

  result = vmaCreateBuffer(allocator, &temp_buffer_info, &temp_allocation_info, &temp_buffer, &temp_buffer_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  // Create buffers
  VmaAllocationCreateInfo allocation_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };

  VkBufferCreateInfo vertex_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = vertex_buffer_size,
  };

  result = vmaCreateBuffer(allocator, &vertex_buffer_info, &allocation_info, &rect->vertex, &rect->vertex_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo index_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = index_buffer_size,
  };

  result = vmaCreateBuffer(allocator, &index_buffer_info, &allocation_info, &rect->index, &rect->index_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  void* mapped;
  result = vmaMapMemory(allocator, temp_buffer_memory, &mapped);
  if(result != VK_SUCCESS) {
    return result;
  }

  vec2* mapped_vertex = (vec2*)mapped;
  mapped_vertex[0][0] = 0.0f;
  mapped_vertex[0][1] = 0.0f;
  mapped_vertex[1][0] = 1.0f;
  mapped_vertex[1][1] = 0.0f;
  mapped_vertex[2][0] = 0.0f;
  mapped_vertex[2][1] = 1.0f;
  mapped_vertex[3][0] = 1.0f;
  mapped_vertex[3][1] = 1.0f;

  uint32_t* mapped_index = (uint32_t*)(mapped + vertex_buffer_size);
  mapped_index[0] = 0;
  mapped_index[1] = 1;
  mapped_index[2] = 2;
  mapped_index[3] = 1;
  mapped_index[4] = 3;
  mapped_index[5] = 2;

  vmaUnmapMemory(allocator, temp_buffer_memory);

  VkCommandBuffer copy_buffer = command_begin_single(device, transfer_pool);
  VkBufferCopy vertex_copy_region = {
    .size = vertex_buffer_size,
    .dstOffset = 0,
    .srcOffset = 0,
  };
  vkCmdCopyBuffer(copy_buffer, temp_buffer, rect->vertex, 1, &vertex_copy_region);

  VkBufferCopy index_copy_region = {
    .size = index_buffer_size,
    .dstOffset = 0,
    .srcOffset = vertex_buffer_size,
  };
  vkCmdCopyBuffer(copy_buffer, temp_buffer, rect->index, 1, &index_copy_region);

  result = command_end_single(device, copy_buffer, transfer_pool, transfer_queue);
  if(result != VK_SUCCESS) {
    return result;
  }

  vmaDestroyBuffer(allocator, temp_buffer, temp_buffer_memory);

  return VK_SUCCESS;
}

VkResult init_pipelines(VkDevice device, VmaAllocator allocator, VkExtent2D swapchain_extent, vec2 window_scale, VkRenderPass render_pass, Queue transfer_queue, VkCommandPool transfer_pool, UIContext* context) {

  VkResult result;
  result = create_ui_descriptor_set(device, allocator, swapchain_extent, window_scale, &context->ui_descriptor_layout, &context->ui_descriptor_pool, &context->ui_descriptor_set, &context->ui_descriptor_memory, &context->ui_descriptor_buffer);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = create_ui_rect_buffer(device, transfer_queue, transfer_pool, allocator, &context->ui_rect);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = create_ui_colored_rect_pipeline(device, render_pass, context->ui_descriptor_layout, &context->ui_pipeline_rect);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = create_font_descriptor_pools(device, 10, &context->font_pool, &context->font_layout);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = create_ui_text_pipeline(device, render_pass, context->ui_descriptor_layout, context->font_layout, &context->ui_pipeline_text, &context->ui_compute_text);
  if(result != VK_SUCCESS) {
    return result;
  }

  return VK_SUCCESS;
}