#include "render.h"
#include "arpa/inet.h"
#include "pipeline.h"
#include "vk_mem_alloc.h"
#include "vulkan/vk_enum_string_helper.h"
#include "vulkan/vulkan_core.h"

#include "ft2build.h"
#include FT_FREETYPE_H

ColoredRect colored_rect(float width, float height, float r, float g, float b, float a, float x, float y, float z) {
  ColoredRect rect = {
    .size = {width, height},
    .pos = {x, y, z},
    .color = {r, g, b, a},
  };
  return rect;
}

typedef struct TextPointersMemoryStruct {
  VmaAllocation pointers_memory;
  VmaAllocation draw_memory;
  VmaAllocation strings_memory;
  VmaAllocation codes_memory;
  VmaAllocation characters_memory;
  VmaAllocation draw_clear_memory;

  VkBuffer pointers_buffer;
  VkBuffer draw_buffer;
  VkBuffer draw_clear_buffer;
  VkBuffer strings_buffer;
  VkBuffer codes_buffer;
  VkBuffer characters_buffer;
} TextPointersMemory;

VkResult create_text_pointers(
    uint32_t max_strings,
    uint32_t max_characters,
    VkDevice device,
    VmaAllocator allocator,
    VkCommandPool transfer_pool,
    Queue transfer_queue,
    TextPointersMemory* memory,
    VkDeviceAddress* address) {
  VkResult result;

  VmaAllocationCreateInfo memory_info = {
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
  };

  VkBufferCreateInfo pointers_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = sizeof(StringPointers),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &pointers_info, &memory_info, &memory->pointers_buffer, &memory->pointers_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo draw_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = sizeof(DrawCommand),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &draw_info, &memory_info, &memory->draw_buffer, &memory->draw_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo draw_clear_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = sizeof(DrawCommand),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &draw_clear_info, &memory_info, &memory->draw_clear_buffer, &memory->draw_clear_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo strings_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = max_strings*sizeof(String),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &strings_info, &memory_info, &memory->strings_buffer, &memory->strings_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo codes_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = max_characters*sizeof(DrawCommand),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &codes_info, &memory_info, &memory->codes_buffer, &memory->codes_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo characters_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .size = max_characters*sizeof(DrawCommand),
    .usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  result = vmaCreateBuffer(allocator, &characters_info, &memory_info, &memory->characters_buffer, &memory->characters_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferDeviceAddressInfo pointers_address = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = memory->pointers_buffer,
  };

  VkBufferDeviceAddressInfo strings_address = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = memory->strings_buffer,
  };

  VkBufferDeviceAddressInfo codes_address = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = memory->codes_buffer,
  };

  VkBufferDeviceAddressInfo characters_address = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = memory->characters_buffer,
  };

  VkBufferDeviceAddressInfo draw_address = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = memory->draw_buffer,
  };

  VkBuffer temp_buffer;
  VmaAllocation temp_memory;
  VkBufferCreateInfo temp_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = sizeof(StringPointers) +
            sizeof(DrawCommand),
  };
  VmaAllocationCreateInfo temp_memory_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };
  result = vmaCreateBuffer(allocator, &temp_buffer_info, &temp_memory_info, &temp_buffer, &temp_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

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

  DrawCommand draw = {
    .index_count = 6,
  };

  memcpy(mapped + 0, &draw, sizeof(DrawCommand));
  StringPointers pointers = {
    .strings = vkGetBufferDeviceAddress(device, &strings_address),
    .codes = vkGetBufferDeviceAddress(device, &codes_address),
    .characters = vkGetBufferDeviceAddress(device, &characters_address),
    .draw = vkGetBufferDeviceAddress(device, &draw_address),
  };
  memcpy(mapped + sizeof(DrawCommand), &pointers, sizeof(StringPointers));
  vmaUnmapMemory(allocator, temp_memory);

  VkCommandBuffer command_buffer = command_begin_single(device, transfer_pool);
  VkBufferCopy copy_draw = {
    .size = sizeof(DrawCommand),
    .srcOffset = 0,
  };
  vkCmdCopyBuffer(command_buffer, temp_buffer, memory->draw_buffer, 1, &copy_draw);
  vkCmdCopyBuffer(command_buffer, temp_buffer, memory->draw_clear_buffer, 1, &copy_draw);
  VkBufferCopy copy_pointers = {
    .size = sizeof(StringPointers),
    .srcOffset = sizeof(DrawCommand),
  };
  vkCmdCopyBuffer(command_buffer, temp_buffer, memory->pointers_buffer, 1, &copy_pointers);
  result = command_end_single(device, command_buffer, transfer_pool, transfer_queue);
  if(result != VK_SUCCESS) {
    return result;
  }
  vmaDestroyBuffer(allocator, temp_buffer, temp_memory);
  *address = vkGetBufferDeviceAddress(device, &pointers_address);

  return VK_SUCCESS;
}

VkResult create_transfer_buffer(VkDeviceSize size, VmaAllocator allocator, VkBuffer* buffer, VmaAllocation* memory, void** mapped) {
  VkBufferCreateInfo buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
    .size = size,
  };
  VmaAllocationCreateInfo memory_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };
  
  VkResult result = vmaCreateBuffer(allocator, &buffer_info, &memory_info, buffer, memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  result = vmaMapMemory(allocator, *memory, mapped);
  if(result != VK_SUCCESS) {
    vmaDestroyBuffer(allocator, *buffer, *memory);
    return result;
  }

  return VK_SUCCESS;
}

VkResult render_thread(GLFWwindow* window, RenderContext* render_context) {
  VkResult result;
  UIContext ui_context;
  VkBuffer colored_rect_buffer;
  VmaAllocation colored_rect_memory;

  result = init_pipelines(render_context->device, render_context->allocator, render_context->swapchain_extent, render_context->window_scale, render_context->render_pass, render_context->transfer_queue, render_context->transfer_pool, &ui_context);
  if(result != VK_SUCCESS) {
    return result;
  }

  FT_Library library;
  if(FT_Init_FreeType(&library) != FT_Err_Ok) {
    return VK_ERROR_UNKNOWN;
  }

  FontDescriptor test_font;
  uint32_t* charmap;
  result = load_font(render_context->device, render_context->allocator, ui_context.font_layout, ui_context.font_pool, render_context->transfer_pool, render_context->transfer_queue, library, "test.ttf", 16, VK_TRUE, &charmap, &test_font);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBufferCreateInfo colored_rect_buffer_info = {
    .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
    .usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
    .size = 3*sizeof(ColoredRect),
    .sharingMode = VK_SHARING_MODE_EXCLUSIVE,
  };

  VmaAllocationCreateInfo colored_rect_memory_info = {
    .usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
  };

  result = vmaCreateBuffer(render_context->allocator, &colored_rect_buffer_info, &colored_rect_memory_info, &colored_rect_buffer, &colored_rect_memory, NULL);
  if(result != VK_SUCCESS) {
    return result;
  }

  ColoredRect* colored_rects;

  result = vmaMapMemory(render_context->allocator, colored_rect_memory, (void**)&colored_rects);
  if(result != VK_SUCCESS) {
    return result;
  }

  colored_rects[0] = colored_rect(100.0, 100.0, 1.0, 0.0, 0.0, 1.0,   0.0,   0.0, 0.5);
  colored_rects[1] = colored_rect(100.0, 100.0, 0.0, 1.0, 0.0, 1.0,   0.0, 100.0, 0.5);
  colored_rects[2] = colored_rect(100.0, 100.0, 0.0, 0.0, 1.0, 1.0, 100.0,   0.0, 0.5);

  vmaUnmapMemory(render_context->allocator, colored_rect_memory);

  TextPointersMemory text_pointers;
  VkDeviceAddress text_pointers_address;
  result = create_text_pointers(10, 100, render_context->device, render_context->allocator, render_context->transfer_pool, render_context->transfer_queue, &text_pointers, &text_pointers_address);
  if(result != VK_SUCCESS) {
    return result;
  }

  VkBuffer temp_buffer;
  VmaAllocation temp_memory;
  void* mapped;
  result = create_transfer_buffer(sizeof(String) + 100*sizeof(uint32_t), render_context->allocator, &temp_buffer, &temp_memory, (void**)&mapped);
  if(result != VK_SUCCESS) {
    return result;
  }

  String* mapped_string = (String*)mapped;
  mapped_string->pos[0] = 200.0;
  mapped_string->pos[1] = 200.0;
  mapped_string->pos[2] = 0.5;
  mapped_string->color[0] = 1.0;
  mapped_string->color[1] = 1.0;
  mapped_string->color[2] = 1.0;
  mapped_string->color[3] = 1.0;
  mapped_string->size = 50.0;
  mapped_string->length = 100;
  mapped_string->offset = 0;

  uint32* mapped_codes = (uint32_t*)(mapped + sizeof(String));
  for(uint32_t i = 0; i < 100; i++) {
    mapped_codes[i]  = i;
  }

  VkCommandBuffer command_buffer = command_begin_single(render_context->device, render_context->transfer_pool);
  VkBufferCopy copy_string = {
    .size = sizeof(String),
    .srcOffset = 0,
  };
  vkCmdCopyBuffer(command_buffer, temp_buffer, text_pointers.strings_buffer, 1, &copy_string);
  VkBufferCopy copy_codes = {
    .size = 100*sizeof(uint32_t),
    .srcOffset = sizeof(String),
  };
  vkCmdCopyBuffer(command_buffer, temp_buffer, text_pointers.codes_buffer, 1, &copy_codes);
  result = command_end_single(render_context->device, command_buffer, render_context->transfer_pool, render_context->transfer_queue);
  if(result != VK_SUCCESS) {
    return result;
  }

  vkQueueWaitIdle(render_context->transfer_queue.handle);
  vmaUnmapMemory(render_context->allocator, temp_memory);
  vmaDestroyBuffer(render_context->allocator, temp_buffer, temp_memory);

  UILayer test_layer = {
    .colored_rects = colored_rect_buffer,
    .colored_rect_count = 3,

    .font = test_font,
    .chars_size = 10*sizeof(Character),
    .chars = text_pointers.characters_buffer,
    .string_count = 1,
    .string_pointers = text_pointers_address,
    .string_draw = text_pointers.draw_buffer,
    .string_draw_clear = text_pointers.draw_clear_buffer,
  };

  while(glfwWindowShouldClose(window) == 0) {
    glfwPollEvents();

    result = draw_frame(render_context, &ui_context, &test_layer, 1);
    if(result != VK_SUCCESS) {
      fprintf(stderr, "draw_frame error: %s\n", string_VkResult(result));
      return result;
    }
  }
  return 0;
}

int logic_thread() {
  return 0;
}

int network_thread() {
  return 0;
}

int main() {
  GLFWwindow* window = init_window();
  if(window == NULL) {
    return 1;
  }

  RenderContext render_context = {};
  if(init_vulkan(window, &render_context) != VK_SUCCESS) {
    return 2;
  }

  if(render_thread(window, &render_context) != VK_SUCCESS) {
    return 3;
  }

  return 0;
}