#include "GLFW/glfw3.h"
#include "ui.h"
#include "gpu.h"
#include "draw.h"

#include "arpa/inet.h"
#include "vk_mem_alloc.h"
#include "vulkan/vk_enum_string_helper.h"
#include "vulkan/vulkan_core.h"
#include "pthread.h"
#include "stdatomic.h"

typedef struct ClientContextStruct {
  GLFWwindow* window;
  RenderContext render;
  UIContext ui;
} ClientContext;

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

  vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, ui_context->string_pipeline.pipeline);
  for(uint32_t i = 0; i < ui_context->max_containers; i++) {
    if(ui_context->containers[i].id != 0x00000000) {
      for(uint32_t j = 0; j < ui_context->containers[i].layer_count; j++) {
        push.layer = ui_context->containers[i].layers[j].address;
        VkBufferMemoryBarrier draw_command_barrier_1 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_INDIRECT_COMMAND_READ_BIT,
          .dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, NULL, 1, &draw_command_barrier_1, 0, NULL);
        command_copy_buffer(command_buffer, ui_context->containers[i].layers[j].layer, ui_context->containers[i].layers[j].layer, offsetof(GPULayer, num_drawables), offsetof(GPULayer, draw) + offsetof(DrawCommand, instance_count), sizeof(uint32_t));
        VkBufferMemoryBarrier draw_command_barrier_2 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT,
          .dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, NULL, 1, &draw_command_barrier_2, 0, NULL);
        vkCmdPushConstants(command_buffer, ui_context->string_pipeline.layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, 16, &push);
        vkCmdDispatchIndirect(command_buffer, ui_context->containers[i].layers[j].layer, offsetof(GPULayer, dispatch_strings));
        VkBufferMemoryBarrier draw_command_barrier_3 = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].layer,
          .offset = offsetof(GPULayer, draw),
          .size = sizeof(DrawCommand),
          .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
          .dstAccessMask = VK_ACCESS_INDIRECT_COMMAND_READ_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT, 0, 0, NULL, 1, &draw_command_barrier_3, 0, NULL);

        VkBufferMemoryBarrier drawables_barrier = {
          .sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER,
          .buffer = ui_context->containers[i].layers[j].drawables,
          .offset = 0,
          .size = sizeof(GPUDrawable)*ui_context->containers[i].layers[j].data.max_drawables,
          .srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT,
          .dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
        };
        vkCmdPipelineBarrier(command_buffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_VERTEX_SHADER_BIT, 0, 0, NULL, 1, &drawables_barrier, 0, NULL);
      }
    }
  }
}

VkResult test_ui(RenderContext* gpu, UIContext* ui) {
  VkResult result;

  GPUString context_strings[] = {
    {
      .pos = {0, 32},
      .size = 32,
      .color = {1.0, 1.0, 1.0, 1.0},
      .offset = 0,
      .length = strlen("Example Text"),
      .font = 0,
    },
    {
      .pos = {0, 64},
      .size = 32,
      .color = {1.0, 1.0, 1.0, 1.0},
      .offset = strlen("Example Text"),
      .length = strlen("Example Text"),
      .font = 1,
    },
    {
      .pos = {0, 96},
      .size = 32,
      .color = {1.0, 1.0, 1.0, 1.0},
      .offset = 2*strlen("Example Text"),
      .length = strlen("Example Text"),
      .font = 2,
    },
  };

  uint32_t context_codes[256];
  map_string("Example Text", context_codes, 0, 0, ui);
  map_string("Example Text", context_codes, strlen("Example Text"), 1, ui);
  map_string("Example Text", context_codes, 2*strlen("Example Text"), 2, ui);

  LayerInput context_layers[] = {
    {
      .num_strings = sizeof(context_strings)/sizeof(GPUString),
      .strings = context_strings,
      .num_codes = sizeof(context_codes)/sizeof(uint32_t),
      .codes = context_codes,
    },
  };

  ContainerInput context_info = {
    .id = 0x01,
    .anchor = ANCHOR_TOP_LEFT,
    .size = {WINDOW_MIN_WIDTH, WINDOW_MIN_HEIGHT},
    .layer_count = sizeof(context_layers)/sizeof(LayerInput),
    .layers = context_layers,
  };

  vec2 map_size = {250, 200};

  GPUDrawable map_rects[] = {
    {
      .pos = {0.0, 0.0},
      .size = {map_size[0], map_size[1]},
      .color = {0.0, 0.5,0.8, 1.0},
      .type = 0,
    },
  };

  LayerInput map_layers[] = {
    {
      .num_drawables = sizeof(map_rects)/sizeof(GPUDrawable),
      .drawables = map_rects,
    },
  };

  ContainerInput map_info = {
    .id = 0x02,
    .anchor = ANCHOR_TOP_RIGHT,
    .size = {map_size[0], map_size[1]},
    .layer_count = sizeof(map_layers)/sizeof(LayerInput),
    .layers = map_layers,
  };

  vec2 inventory_size = {map_size[0], WINDOW_MIN_HEIGHT-map_size[1]};

  GPUDrawable inventory_rects[] = {
    {
      .pos = {0, 0},
      .size = {inventory_size[0], inventory_size[1]},
      .color = {0.5, 0.8, 0.0, 1.0},
      .type = 0,
    },
  };

  LayerInput inventory_layers[] = {
    {
      .num_drawables = sizeof(inventory_rects)/sizeof(GPUDrawable),
      .drawables = inventory_rects,
    },
  };

  ContainerInput inventory_info = {
    .id = 0x03,
    .anchor = ANCHOR_BOTTOM_RIGHT,
    .size = {inventory_size[0], inventory_size[1]},
    .layer_count = sizeof(inventory_layers)/sizeof(LayerInput),
    .layers = inventory_layers,
  };

  vec2 chat_size = {WINDOW_MIN_WIDTH-inventory_size[0], 200};

  GPUDrawable chat_rects[] = {
    {
      .pos = {0, 0},
      .size = {chat_size[0], chat_size[1]},
      .color = {0.8, 0.0, 0.5, 1.0},
      .type = 0,
    },
  };

  LayerInput chat_layers[] = {
    {
      .num_drawables = sizeof(chat_rects)/sizeof(GPUDrawable),
      .drawables = chat_rects,
    },
  };

  ContainerInput chat_info = {
    .id = 0x04,
    .anchor = ANCHOR_BOTTOM_LEFT,
    .size = {chat_size[0], chat_size[1]},
    .layer_count = sizeof(chat_layers)/sizeof(LayerInput),
    .layers = chat_layers,
  };

  VK_RESULT(create_container(&context_info, gpu, ui));
  VK_RESULT(create_container(&map_info, gpu, ui));
  VK_RESULT(create_container(&inventory_info, gpu, ui));
  VK_RESULT(create_container(&chat_info, gpu, ui));

  for(uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    VkCommandBuffer command_buffer = command_begin_single(gpu->device, gpu->transfer_pool);
    record_ui_compute(command_buffer, ui);
    VK_RESULT(command_end_single(gpu->device, command_buffer, gpu->transfer_pool, gpu->transfer_queue)); 
  }


  return VK_SUCCESS;
}

// Threads:
// 1. render
//  - Submits the draw buffer to the GPU as soon as it can
// 2. network
//  - Handles packets to/from the network to/from the main thread
// 3. main
//  - updates the data in the GPU that's being drawn from
//  - updates the data in the GPU from network requests
//
//  Data:
//  Render thread reads Render and UI context 
//  Main thread reads and writes UI context

void* render_thread(void* data) {
  ClientContext* context = (ClientContext*)data;

  double last_frame_time = glfwGetTime();
  while(glfwWindowShouldClose(context->window) == 0) {
    double frame_time = glfwGetTime();
    double delta_time = frame_time - last_frame_time;
    (void)delta_time;

    VkResult result = draw_frame(&context->render, &context->ui, frame_time);
    if(result != VK_SUCCESS) {
      fprintf(stderr, "draw_frame error: %s\n", string_VkResult(result));
      glfwDestroyWindow(context->window);
    }
    last_frame_time = frame_time;
  }

  return NULL;
}

void* network_thread(void* data) {
  ClientContext* context = (ClientContext*)data;
  (void)context;

  return NULL;
}

int main_thread(void* data) {
  ClientContext* context = (ClientContext*)data;

  int x = 0;
  while(glfwWindowShouldClose(context->window) == 0) {
    glfwPollEvents();
    if(x == 0 && glfwGetTime() > 0.0) {
      x = 1;
      test_ui(&context->render, &context->ui);
    }
  }

  return 0;
}

void key_callback(GLFWwindow* window, int key, int scancode, int action, int mods) {
  (void)window;
  (void)key;
  (void)scancode;
  (void)action;
  (void)mods;
}

void mouse_button_callback(GLFWwindow* window, int button, int action, int mods) {
  ClientContext* context = (ClientContext*)glfwGetWindowUserPointer(window);

  (void)mods;
  (void)context;
  switch(button) {
  // Handle camera hover
  case GLFW_MOUSE_BUTTON_MIDDLE:
    if(action == GLFW_PRESS) {
      glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_DISABLED);
    } else {
      glfwSetInputMode(window, GLFW_CURSOR, GLFW_CURSOR_NORMAL);
    }
    break;
  case GLFW_MOUSE_BUTTON_LEFT:
    if(action == GLFW_PRESS) {
      // NOTE: this logic is the same as it would be for right-click(just with different outcomes), so dont repeat yourself
      // 1. Search through the UI context for the first element that the contains the mouse point
      // 2. If no UI element intersection, cast a ray through the world scene to find the "clicked entity" 
      // 3. Based on what was clicked, start the mouse click animation at the target area
    }
    break;
  case GLFW_MOUSE_BUTTON_RIGHT:
    if(action == GLFW_PRESS) {
    }
    break;
  }
}

void scroll_callback(GLFWwindow* window, double xoffset, double yoffset) {
  (void)window;
  (void)xoffset;
  (void)yoffset;
}

void cursor_pos_callback(GLFWwindow* window, double xpos, double ypos) {
  (void)window;
  (void)xpos;
  (void)ypos;
}

int main() {
  ClientContext context = {};
  context.window = init_window();
  if(context.window == NULL) {
    return 1;
  }

  glfwSetWindowUserPointer(context.window, &context);
  glfwSetKeyCallback(context.window, key_callback);
  glfwSetMouseButtonCallback(context.window, mouse_button_callback);
  glfwSetScrollCallback(context.window, scroll_callback);
  glfwSetCursorPosCallback(context.window, cursor_pos_callback);

  int error;
  VkResult result;
  VK_RESULT(init_vulkan(context.window, &context.render));
  
  // TODO: make # of fonts/textures/containers scaling, recreate GPU buffers as necessary
  VK_RESULT(create_ui_context(10, 10, 10, &context.render, &context.ui));

  // Start threads
  pthread_t render_thread_handle;
  pthread_t network_thread_handle;

  error = pthread_create(&render_thread_handle, NULL, &render_thread, &context);
  if(error != 0) {
    return error;
  }

  error = pthread_create(&network_thread_handle, NULL, &network_thread, &context);
  if(error != 0) {
    return error;
  }

  error = main_thread(&context);
  if(error != 0) {
    return error;
  }

  error = pthread_join(render_thread_handle, NULL);
  if(error != 0) {
    return error;
  }

  error = pthread_join(network_thread_handle, NULL);
  if(error != 0) {
    return error;
  }

  return 0;
}