#include "ui.h"
#include "cglm/affine.h"
#include "cglm/mat4.h"
#include "command.h"
#include "render.h"
#include "stdio.h"
#include "stdlib.h"
#include "string.h"
#include "vk_mem_alloc.h"
#include "vulkan/vulkan_core.h"
#include "spng.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,
    VkDescriptorSetLayout samplers_layout,
    VkDescriptorSetLayout textures_layout,
    GraphicsPipeline* pipeline,
    ComputePipeline* compute) {
  VkResult result;
  VkShaderModule compute_shader = load_shader_file("shader_src/string.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,
    .size = 16,
  };

  VkPipelineLayoutCreateInfo compute_layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pushConstantRangeCount = 1,
    .pPushConstantRanges = &push_constant,
  };

  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.vert.spv", device);
  if(vert_shader == VK_NULL_HANDLE) {
    return VK_ERROR_UNKNOWN;
  }
  VkShaderModule frag_shader = load_shader_file("shader_src/ui.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[] = {};
  VkVertexInputAttributeDescription attributes[] = {};

  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),
  };
 
  VkPipelineInputAssemblyStateCreateInfo input_assembly_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
    .topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
    .primitiveRestartEnable = VK_FALSE,
  };

  VkPushConstantRange push_constants[] = {
    {
      .stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT,
      .size = 16,
      .offset = 0,
    },
  };

  VkDescriptorSetLayout set_layouts[] = {samplers_layout, textures_layout, samplers_layout, textures_layout};

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pushConstantRangeCount = sizeof(push_constants)/sizeof(VkPushConstantRange),
    .pPushConstantRanges = push_constants,
    .setLayoutCount = sizeof(set_layouts)/sizeof(VkDescriptorSetLayout),
    .pSetLayouts = set_layouts,
  };

  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,
  };

  // additive blending
  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,
    .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,
  };

  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 = sizeof(shader_stages)/sizeof(VkPipelineShaderStageCreateInfo),
    .pStages = shader_stages,
    .pVertexInputState = &input_info,
    .pInputAssemblyState = &input_assembly_info,
    .pViewportState = &viewport_state,
    .pRasterizationState = &raster_info,
    .pColorBlendState = &color_blend_info,
    .pDynamicState = &dynamic_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_container(
    float x,
    float y,
    float width,
    float height,
    VkDevice device,
    VmaAllocator allocator,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    UIContainerStorage* memory) {
  VkResult result;

  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(UIContainer), &memory->container, &memory->container_memory));

  VkBuffer transfer;
  VmaAllocation transfer_memory;
  void* mapped;
  VK_RESULT(create_transfer_buffer(allocator, sizeof(UIContainer), &transfer, &transfer_memory, &mapped));

  memory->data.pos[0] = x;
  memory->data.pos[1] = y;
  memory->data.size[0] = width;
  memory->data.size[1] = height;
  memcpy(mapped, &memory->data, sizeof(UIContainer));
  
  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);
  command_copy_buffer(command_buffer, transfer, memory->container, 0, 0, sizeof(UIContainer));
  VK_RESULT(command_end_single(device, command_buffer, transfer_pool, transfer_queue));
  destroy_transfer_buffer(allocator, transfer, transfer_memory);

  memory->address = buffer_address(device, memory->container);

  return VK_SUCCESS;
}

VkResult create_layer(
    uint32_t max_strings,
    uint32_t max_codes,
    uint32_t max_drawables,
    uint32_t font_index,
    VkDevice device,
    VmaAllocator allocator,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    UIContainerStorage* container,
    UILayerStorage* memory) {
  VkResult result;

  VK_RESULT(create_storage_buffer(allocator, VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT, sizeof(UILayer), &memory->layer, &memory->layer_memory));
  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(UIString)*max_strings, &memory->strings, &memory->strings_memory))
  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(UIDrawable)*(max_drawables + max_codes), &memory->drawables, &memory->drawables_memory));
  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(uint32_t)*max_codes, &memory->codes, &memory->codes_memory));

  VkBuffer transfer;
  VmaAllocation transfer_memory;
  void* mapped;
  VK_RESULT(create_transfer_buffer(allocator, sizeof(UILayer), &transfer, &transfer_memory, &mapped));

  memory->data.strings = buffer_address(device, memory->strings);
  memory->data.codes = buffer_address(device, memory->codes);

  memory->data.drawables = buffer_address(device, memory->drawables);

  memory->data.draw.first_vertex = 0;
  memory->data.draw.vertex_count = 6;
  memory->data.draw.first_instance = 0;
  memory->data.draw.instance_count = 0;

  memory->data.dispatch_strings.x = 0;
  memory->data.dispatch_strings.y = 1;
  memory->data.dispatch_strings.z = 1;

  memory->data.font_index = font_index;
  memory->data.max_drawables = max_drawables;
  memory->data.max_strings = max_strings;
  memory->data.num_drawables = 0;
  memory->data.container = container->address;
  memcpy(mapped, &memory->data, sizeof(UILayer));

  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);
  command_copy_buffer(command_buffer, transfer, memory->layer, 0, 0, sizeof(UILayer));
  vkCmdFillBuffer(command_buffer, memory->strings, 0, sizeof(UIString)*max_strings, 0x00000000);
  vkCmdFillBuffer(command_buffer, memory->drawables, 0, sizeof(UIDrawable)*max_drawables, 0x00000000);
  vkCmdFillBuffer(command_buffer, memory->codes, 0, sizeof(uint32_t)*max_codes, 0x00000000);
  VK_RESULT(command_end_single(device, command_buffer, transfer_pool, transfer_queue));
  vkQueueWaitIdle(transfer_queue.handle);
  destroy_transfer_buffer(allocator, transfer, transfer_memory);

  memory->address = buffer_address(device, memory->layer);

  return VK_SUCCESS;
}

VkResult load_texture(
    VkDevice device,
    VmaAllocator allocator,
    UIContextStorage* context,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    const char* png_path,
    uint32_t* index,
    TextureStorage* memory) {
  *index = 0xFFFFFFFF;
  for(uint32_t i = 0; i < context->max_textures; i++) {
    if(context->texture_slots[i] == 0) {
      context->texture_slots[i] = 1;
      *index = i;
      break;
    }
  }
  if(*index == 0xFFFFFFFF) {
    return VK_ERROR_OUT_OF_DEVICE_MEMORY;
  }

  spng_ctx* spng = spng_ctx_new(0);
  if(spng == NULL) {
    return VK_ERROR_UNKNOWN;
  }

  FILE* png_file = fopen(png_path, "rb");
  if(png_file == NULL) {
    return VK_ERROR_UNKNOWN;
  }

  fseek(png_file, 0, SEEK_END);
  int error = fseek(png_file, 0, SEEK_END);
  if(error != 0) {
    return VK_ERROR_UNKNOWN;
  }

  size_t png_size = ftell(png_file);

  error = fseek(png_file, 0, SEEK_SET);
  if(error != 0) {
    return VK_ERROR_UNKNOWN;
  }


  void* png = malloc(png_size);
  if(png == NULL) {
    return VK_ERROR_UNKNOWN;
  }

  size_t read = fread(png, 1, png_size, png_file);
  if(read != png_size) {
    return VK_ERROR_UNKNOWN;
  }

  fclose(png_file);

  spng_set_png_buffer(spng, png, png_size);
  size_t out_size;
  spng_decoded_image_size(spng, SPNG_FMT_RGBA8, &out_size);
  struct spng_ihdr ihdr;
  spng_get_ihdr(spng, &ihdr);
  void* image_buffer = malloc(out_size);
  if(image_buffer == NULL) {
    return VK_ERROR_UNKNOWN;
  }

  spng_decode_image(spng, image_buffer, out_size, SPNG_FMT_RGBA8, 0);

  VkResult result;

  VkImageCreateInfo image_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
    .usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
    .extent = {
      .width = ihdr.width,
      .height = ihdr.height,
      .depth = 1,
    },
    .mipLevels = 1,
    .arrayLayers = 1,
    .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 memory_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };

  VkBuffer transfer;
  VmaAllocation transfer_memory;
  void* mapped;

  VK_RESULT(vmaCreateImage(allocator, &image_info, &memory_info, &memory->image, &memory->image_memory, NULL));
  VK_RESULT(create_transfer_buffer(allocator, sizeof(uint32_t)*ihdr.width*ihdr.height, &transfer, &transfer_memory, &mapped));
  memcpy(mapped, image_buffer, sizeof(uint32_t)*ihdr.height*ihdr.width);
  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);

  VkImageMemoryBarrier first_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = memory->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,
      .levelCount = 1,
      .layerCount = 1,
    },
    .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 = {
    .imageExtent = {
      .width = ihdr.width,
      .height = ihdr.height,
      .depth = 1,
    },
    .imageSubresource = {
      .layerCount = 1,
      .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
    },
  };
  vkCmdCopyBufferToImage(command_buffer, transfer, memory->image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_copy);

  VkImageMemoryBarrier second_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = memory->image,
    .oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
    .newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    .subresourceRange = {
      .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
      .layerCount = 1,
      .levelCount = 1,
    },
    .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);


  VK_RESULT(command_end_single(device, command_buffer, transfer_pool, transfer_queue));
  vkQueueWaitIdle(transfer_queue.handle);
  destroy_transfer_buffer(allocator, transfer, transfer_memory);
  free(image_buffer);
  free(png);
  spng_ctx_free(spng);

  VkImageViewCreateInfo view_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
    .image = memory->image,
    .viewType = VK_IMAGE_VIEW_TYPE_2D,
    .format = VK_FORMAT_R8G8B8A8_SRGB,
    .subresourceRange = {
      .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
      .layerCount = 1,
      .levelCount = 1,
    },
  };

  VK_RESULT(vkCreateImageView(device, &view_info, NULL, &memory->view));

  VkSamplerCreateInfo sampler_info = {
    .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
    .magFilter = VK_FILTER_LINEAR,
    .minFilter = VK_FILTER_LINEAR,
    .addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT,
    .addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT,
    .addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT,
  };

  VK_RESULT(vkCreateSampler(device, &sampler_info, NULL, &memory->sampler));

  VkDescriptorImageInfo desc_sampler_info = {
    .sampler = memory->sampler,
  };

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

  VkWriteDescriptorSet desc_writes[] = {
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = context->textures,
      .dstBinding = 0,
      .dstArrayElement = *index,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
      .descriptorCount = 1,
      .pImageInfo = &desc_image_info,
    },
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = context->samplers,
      .dstBinding = 0,
      .dstArrayElement = *index,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER,
      .descriptorCount = 1,
      .pImageInfo = &desc_sampler_info,
    },
  };

  vkUpdateDescriptorSets(device, sizeof(desc_writes)/sizeof(VkWriteDescriptorSet), desc_writes, 0, NULL);

  return VK_SUCCESS;
}

VkResult load_font(
    VkDevice device,
    VmaAllocator allocator,
    UIContextStorage* context,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    FT_Library library,
    const char* ttf_file,
    uint32_t size,
    VkBool32 antialias,
    uint32_t* index,
    FontStorage* memory) {  
  *index = 0xFFFFFFFF;
  for(uint32_t i = 0; i < context->max_fonts; i++) {
    if(context->font_slots[i] == 0) {
      context->font_slots[i] = 1;
      *index = i;
      break;
    }
  }
  if(*index == 0xFFFFFFFF) {
    return VK_ERROR_OUT_OF_DEVICE_MEMORY;
  }

  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);
  Font info;

  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;
    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);
  }

  info.width = max_width;
  info.height = max_height;
  info.num_symbols = symbol_count;
  memory->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);
  memory->charmap = malloc(sizeof(uint32_t)*symbol_count);
  memcpy(memory->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, memory->charmap[i]);
    FT_Load_Glyph(face, glyph_index, load_flags);
    symbols[i].width = (float)face->glyph->bitmap.width/(float)max_width;
    symbols[i].height = (float)face->glyph->bitmap.rows/(float)max_height;
    symbols[i].advance = (float)face->glyph->advance.x*16.0/(float)face->units_per_EM/(float)max_width;
    symbols[i].left = (float)face->glyph->bitmap_left/(float)max_width;
    symbols[i].top = (float)face->glyph->bitmap_top/(float)max_height;
    for(uint32_t y = 0; y < face->glyph->bitmap.rows; y++) {
      for(uint32_t x = 0; x < face->glyph->bitmap.width; x++) {
        uint64_t level = face->glyph->bitmap.buffer[y*face->glyph->bitmap.width+x];
        level = ((level * 0xFFFFFFFF) / 256);
        images[max_width*max_height*i + max_width*y + x] = level;
      }
    }
  }


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

  VkResult result;
  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(SymbolInfo)*info.num_symbols, &memory->symbols, &memory->symbol_memory));

  VkImageCreateInfo image_info = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
    .usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
    .extent.depth  = 1,
    .extent.width  = info.width,
    .extent.height = info.height,
    .mipLevels = 1,
    .arrayLayers = info.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,
  };

  VK_RESULT(vmaCreateImage(allocator, &image_info, &image_memory_info, &memory->image, &memory->image_memory, NULL));

  VkBuffer transfer;
  VmaAllocation transfer_memory;
  void* mapped;
  VK_RESULT(create_transfer_buffer(allocator, sizeof(Font) + image_size*info.num_symbols + sizeof(SymbolInfo)*info.num_symbols, &transfer, &transfer_memory, &mapped));
  
	info.symbol_list = buffer_address(device, memory->symbols);

  memcpy(mapped, images, image_size*info.num_symbols);
  memcpy(mapped + image_size*info.num_symbols, &info, sizeof(Font));
  memcpy(mapped + image_size*info.num_symbols + sizeof(Font), symbols, sizeof(SymbolInfo)*info.num_symbols);

  free(images);
  free(symbols);

  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);
  command_copy_buffer(command_buffer, transfer, context->font_infos, image_size*info.num_symbols, *index*sizeof(Font), sizeof(Font));
  command_copy_buffer(command_buffer, transfer, memory->symbols, image_size*info.num_symbols + sizeof(Font), 0, sizeof(SymbolInfo)*info.num_symbols);

  VkImageMemoryBarrier first_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = memory->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 = info.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 = info.num_symbols,
    .imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
    .imageExtent = image_info.extent,
  };
  vkCmdCopyBufferToImage(command_buffer, transfer, memory->image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &image_copy);

  VkImageMemoryBarrier second_barrier = {
    .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
    .image = memory->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 = info.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;
  }
  vkQueueWaitIdle(transfer_queue.handle);
  destroy_transfer_buffer(allocator, transfer, transfer_memory);

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

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

  VkDescriptorImageInfo desc_sampler_info = {
    .sampler = memory->sampler,
  };

  VkDescriptorImageInfo desc_texture_info = {
    .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
    .imageView = memory->view,
  };

  VkWriteDescriptorSet descriptor_writes[] = {
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = context->font_textures,
      .dstBinding = 0,
      .dstArrayElement = *index,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
      .descriptorCount = 1,
      .pImageInfo = &desc_texture_info,
    },
    {
      .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
      .dstSet = context->font_samplers,
      .dstBinding = 0,
      .dstArrayElement = *index,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER,
      .descriptorCount = 1,
      .pImageInfo = &desc_sampler_info,
    },
  };

  vkUpdateDescriptorSets(device, sizeof(descriptor_writes)/sizeof(VkWriteDescriptorSet), descriptor_writes, 0, NULL);
  return VK_SUCCESS;
}

VkResult create_ui_descriptor(
    VkDevice device,
    uint32_t max_fonts,
    uint32_t max_textures,
    VkDescriptorSetLayout* sampler_layout,
    VkDescriptorSetLayout* texture_layout,
    VkDescriptorSet* samplers,
    VkDescriptorSet* textures,
    VkDescriptorSet* font_samplers,
    VkDescriptorSet* font_textures,
    VkDescriptorPool* font_pool,
    VkDescriptorPool* texture_pool) { 
  VkResult result;
  VkDescriptorSetLayoutBinding sampler_bindings[] = {
    {
      .binding = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER,
      .descriptorCount = max_fonts,
      .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
    },
  };

  VkDescriptorBindingFlags bindless_flags[] = {
    VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT | VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT,
  };

  VkDescriptorSetLayoutBindingFlagsCreateInfo sampler_binding_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
    .bindingCount = sizeof(sampler_bindings)/sizeof(VkDescriptorSetLayoutBinding),
    .pBindingFlags = bindless_flags,
  };

  VkDescriptorSetLayoutCreateInfo sampler_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .pBindings = sampler_bindings,
    .bindingCount = sizeof(sampler_bindings)/sizeof(VkDescriptorSetLayoutBinding),
    .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
    .pNext = &sampler_binding_info,
  }; 

  VK_RESULT(vkCreateDescriptorSetLayout(device, &sampler_info, NULL, sampler_layout));

  VkDescriptorSetLayoutBinding texture_bindings[] = {
    {
      .binding = 0,
      .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
      .descriptorCount = max_fonts,
      .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
    },
  };

  VkDescriptorSetLayoutBindingFlagsCreateInfo texture_bindings_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
    .bindingCount = sizeof(texture_bindings)/sizeof(VkDescriptorSetLayoutBinding),
    .pBindingFlags = bindless_flags,
  };

  VkDescriptorSetLayoutCreateInfo texture_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .pBindings = texture_bindings,
    .bindingCount = sizeof(texture_bindings)/sizeof(VkDescriptorSetLayoutBinding),
    .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
    .pNext = &texture_bindings_info,
  }; 

  VK_RESULT(vkCreateDescriptorSetLayout(device, &texture_info, NULL, texture_layout));

  VkDescriptorPoolSize font_pool_sizes[] = {
    {
      .type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
      .descriptorCount = max_fonts,
    },
    {
      .type = VK_DESCRIPTOR_TYPE_SAMPLER,
      .descriptorCount = max_fonts,
    },
  };
  VkDescriptorPoolCreateInfo font_pool_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .pPoolSizes = font_pool_sizes,
    .poolSizeCount = 2,
    .maxSets = 2,
    .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
  };

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

  uint32_t max_font_binding = max_fonts - 1;
  VkDescriptorSetVariableDescriptorCountAllocateInfo count_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO,
    .descriptorSetCount = 1,
    .pDescriptorCounts = &max_font_binding,
  };

  VkDescriptorSetAllocateInfo allocate_font_samplers = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = sampler_layout,
    .descriptorSetCount = 1,
    .descriptorPool = *font_pool,
    .pNext = &count_info,
  };

  VkDescriptorSetAllocateInfo allocate_font_textures = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = texture_layout,
    .descriptorSetCount = 1,
    .descriptorPool = *font_pool,
    .pNext = &count_info,
  };

  result = vkAllocateDescriptorSets(device, &allocate_font_samplers, font_samplers);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = vkAllocateDescriptorSets(device, &allocate_font_textures, font_textures);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkDescriptorPoolSize texture_pool_sizes[] = {
    {
      .type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
      .descriptorCount = max_textures,
    },
    {
      .type = VK_DESCRIPTOR_TYPE_SAMPLER,
      .descriptorCount = max_textures,
    },
  };
  VkDescriptorPoolCreateInfo texture_pool_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .pPoolSizes = texture_pool_sizes,
    .poolSizeCount = 2,
    .maxSets = 2,
    .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
  };

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

  uint32_t max_texture_binding = max_textures - 1;
  VkDescriptorSetVariableDescriptorCountAllocateInfo texture_count_info = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO,
    .descriptorSetCount = 1,
    .pDescriptorCounts = &max_texture_binding,
  };

  VkDescriptorSetAllocateInfo allocate_samplers = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = sampler_layout,
    .descriptorSetCount = 1,
    .descriptorPool = *texture_pool,
    .pNext = &texture_count_info,
  };

  VkDescriptorSetAllocateInfo allocate_textures = {
    .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .pSetLayouts = texture_layout,
    .descriptorSetCount = 1,
    .descriptorPool = *texture_pool,
    .pNext = &texture_count_info,
  };

  result = vkAllocateDescriptorSets(device, &allocate_samplers, samplers);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = vkAllocateDescriptorSets(device, &allocate_textures, textures);
  if(result != VK_SUCCESS) {
    return result;
  }

  return VK_SUCCESS;
}

VkResult create_ui_context(
    VkDevice device,
    VmaAllocator allocator,
    VkRenderPass render_pass,
    uint32_t max_fonts,
    uint32_t max_textures,
    VkExtent2D swapchain_extent,
    vec2 window_scale,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    UIContextStorage* memory) {

  VkResult result;

  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(UIContext), &memory->context, &memory->context_memory)); 
  VK_RESULT(create_storage_buffer(allocator, 0, sizeof(Font)*max_fonts, &memory->font_infos, &memory->font_infos_memory)); 

  VkBuffer transfer;
  VmaAllocation transfer_memory;
  UIContext* mapped;

  VK_RESULT(create_transfer_buffer(allocator, sizeof(UIContext), &transfer, &transfer_memory, (void**)&mapped));

  memory->data.font_infos = buffer_address(device, memory->font_infos);
  memory->data.screen[0] = window_scale[0] / swapchain_extent.width;
  memory->data.screen[1] = window_scale[1] / swapchain_extent.height;
  memory->data.scale[0] = window_scale[0];
  memory->data.scale[1] = window_scale[1];
  memcpy(mapped, &memory->data, sizeof(UIContext));

  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);
  command_copy_buffer(command_buffer, transfer, memory->context, 0, 0, sizeof(UIContext));
  VK_RESULT(command_end_single(device, command_buffer, transfer_pool, transfer_queue));
  destroy_transfer_buffer(allocator, transfer, transfer_memory);
  memory->address = buffer_address(device, memory->context);

  VK_RESULT(create_ui_descriptor(device, max_fonts, max_textures, &memory->samplers_layout, &memory->textures_layout, &memory->samplers, &memory->textures, &memory->font_samplers, &memory->font_textures, &memory->fonts_pool, &memory->textures_pool));

  VK_RESULT(create_ui_pipeline(device, render_pass, memory->samplers_layout, memory->textures_layout, &memory->pipeline, &memory->string_pipeline));

  memory->max_textures = max_textures;
  memory->max_fonts = max_fonts;
  memory->texture_slots = malloc(max_textures);
  memory->font_slots = malloc(max_fonts);

  return VK_SUCCESS;
}

void set_ui_rect(
    float x,
    float y,
    float width,
    float height,
    float r,
    float g,
    float b,
    float a,
    UIDrawable* drawable) {
  drawable->pos[0] = x;
  drawable->pos[1] = y;
  drawable->size[0] = width;
  drawable->size[1] = height;
  drawable->color[0] = r;
  drawable->color[1] = g;
  drawable->color[2] = b;
  drawable->color[3] = a;
  drawable->type = 0;
  drawable->code = 0;
}

void set_ui_image(
    float x,
    float y,
    float width,
    float height,
    float r,
    float g,
    float b,
    float a,
    uint32_t index,
    UIDrawable* drawable) {
  drawable->pos[0] = x;
  drawable->pos[1] = y;
  drawable->size[0] = width;
  drawable->size[1] = height;
  drawable->color[0] = r;
  drawable->color[1] = g;
  drawable->color[2] = b;
  drawable->color[3] = a;
  drawable->type = 2;
  drawable->code = index;
}

void set_ui_string(
    float x,
    float y,
    float size,
    float r,
    float g,
    float b,
    float a,
    uint32_t length,
    uint32_t offset,
    UIString* string) {
  string->pos[0] = x;
  string->pos[1] = y;
  string->size = size;
  string->color[0] = r;
  string->color[1] = g;
  string->color[2] = b;
  string->color[3] = a;
  string->length = length;
  string->offset = offset;
}

VkResult set_ui_codes(
    const char * text,
    uint32_t* buffer,
    uint32_t offset,
    uint32_t* charmap,
    uint32_t charmap_size) {
  size_t i = 0;
  while(text[i] != '\0') {
    uint32_t mapped = 0xFFFFFFFF;
    for(uint32_t j = 0; j < charmap_size; j++) {
      if(charmap[j] == (uint32_t)text[i]) {
        mapped = j;
        break;
      }
    }
    if(mapped == 0xFFFFFFFF) {
      return VK_ERROR_UNKNOWN;
    }
    buffer[i + offset] = mapped;
    i += 1;
  }
  
  return VK_SUCCESS;
}