#include "hex.h"
#include "gpu.h"
#include "vulkan/vulkan_core.h"
#include <stdio.h>

VkResult create_ray_pipeline(
    RenderContext* gpu,
    GraphicsPipeline* pipeline) {
  VkResult result;

  VkShaderModule vert_shader = load_shader_file("shader/ray.vert.spv", gpu->device);
  VkShaderModule frag_shader = load_shader_file("shader/ray.frag.spv", gpu->device);

  VkPipelineShaderStageCreateInfo stages[] = {
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_VERTEX_BIT,
      .pName = "main",
      .module = vert_shader,
    },
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
      .pName = "main",
      .module = frag_shader,
    },
  };

  VkPushConstantRange push = {
    .size = sizeof(HexPushConstant),
    .offset = 0,
    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
  };

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pPushConstantRanges = &push,
    .pushConstantRangeCount = 1,
  };

  VK_RESULT(vkCreatePipelineLayout(gpu->device, &layout_info, NULL, &pipeline->layout));

  VkPipelineVertexInputStateCreateInfo vertex_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
  };

  VkPipelineInputAssemblyStateCreateInfo input_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
    .topology = VK_PRIMITIVE_TOPOLOGY_LINE_LIST,
    .primitiveRestartEnable = VK_FALSE,
  };

  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_info = {
    .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,
  };

  VkPipelineColorBlendAttachmentState blend_attachments = {
    .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_ZERO,
    .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
    .alphaBlendOp = VK_BLEND_OP_ADD,
  };

  VkPipelineColorBlendStateCreateInfo blend_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
    .logicOpEnable = VK_FALSE,
    .logicOp = VK_LOGIC_OP_COPY,
    .attachmentCount = 1,
    .pAttachments = &blend_attachments,
  };

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

  VkPipelineDynamicStateCreateInfo dynamic_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    .dynamicStateCount = sizeof(dynamic_states)/sizeof(VkDynamicState),
    .pDynamicStates = dynamic_states,
  };

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

  VkGraphicsPipelineCreateInfo graphics_pipeline_info = {
    .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
    .layout = pipeline->layout,
    .stageCount = sizeof(stages)/sizeof(VkPipelineShaderStageCreateInfo),
    .pStages = stages,
    .pVertexInputState = &vertex_info,
    .pInputAssemblyState = &input_info,
    .pViewportState = &viewport_info,
    .pRasterizationState = &raster_info,
    .pColorBlendState = &blend_info,
    .pDynamicState = &dynamic_info,
    .pMultisampleState = &multisample_info,
    .pDepthStencilState = &depth_info,
    .renderPass = gpu->render_pass,
    .subpass = 0,
    .basePipelineHandle = VK_NULL_HANDLE,
    .basePipelineIndex = -1,
  };

  VK_RESULT(vkCreateGraphicsPipelines(gpu->device, VK_NULL_HANDLE, 1, &graphics_pipeline_info, NULL, &pipeline->pipeline));
  return VK_SUCCESS;
}

VkResult create_hex_highlight_pipeline(
    RenderContext* gpu,
    GraphicsPipeline* pipeline) {
  VkResult result;

  VkShaderModule vert_shader = load_shader_file("shader/hex_highlight.vert.spv", gpu->device);
  VkShaderModule frag_shader = load_shader_file("shader/hex_highlight.frag.spv", gpu->device);

  VkPipelineShaderStageCreateInfo stages[] = {
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_VERTEX_BIT,
      .pName = "main",
      .module = vert_shader,
    },
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
      .pName = "main",
      .module = frag_shader,
    },
  };

  VkPushConstantRange push = {
    .size = sizeof(HexPushConstant),
    .offset = 0,
    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
  };

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pPushConstantRanges = &push,
    .pushConstantRangeCount = 1,
  };

  VK_RESULT(vkCreatePipelineLayout(gpu->device, &layout_info, NULL, &pipeline->layout));

  VkPipelineVertexInputStateCreateInfo vertex_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
  };

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

  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_info = {
    .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,
  };

  VkPipelineColorBlendAttachmentState blend_attachments = {
    .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_ZERO,
    .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
    .alphaBlendOp = VK_BLEND_OP_ADD,
  };

  VkPipelineColorBlendStateCreateInfo blend_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
    .logicOpEnable = VK_FALSE,
    .logicOp = VK_LOGIC_OP_COPY,
    .attachmentCount = 1,
    .pAttachments = &blend_attachments,
  };

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

  VkPipelineDynamicStateCreateInfo dynamic_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    .dynamicStateCount = sizeof(dynamic_states)/sizeof(VkDynamicState),
    .pDynamicStates = dynamic_states,
  };

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

  VkGraphicsPipelineCreateInfo graphics_pipeline_info = {
    .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
    .layout = pipeline->layout,
    .stageCount = sizeof(stages)/sizeof(VkPipelineShaderStageCreateInfo),
    .pStages = stages,
    .pVertexInputState = &vertex_info,
    .pInputAssemblyState = &input_info,
    .pViewportState = &viewport_info,
    .pRasterizationState = &raster_info,
    .pColorBlendState = &blend_info,
    .pDynamicState = &dynamic_info,
    .pMultisampleState = &multisample_info,
    .pDepthStencilState = &depth_info,
    .renderPass = gpu->render_pass,
    .subpass = 0,
    .basePipelineHandle = VK_NULL_HANDLE,
    .basePipelineIndex = -1,
  };

  VK_RESULT(vkCreateGraphicsPipelines(gpu->device, VK_NULL_HANDLE, 1, &graphics_pipeline_info, NULL, &pipeline->pipeline));
  return VK_SUCCESS;
}

VkResult create_hex_pipeline(
    RenderContext* gpu,
    GraphicsPipeline* pipeline) {
  VkResult result;

  VkShaderModule vert_shader = load_shader_file("shader/hex.vert.spv", gpu->device);
  VkShaderModule frag_shader = load_shader_file("shader/hex.frag.spv", gpu->device);

  VkPipelineShaderStageCreateInfo stages[] = {
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_VERTEX_BIT,
      .pName = "main",
      .module = vert_shader,
    },
    {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
      .pName = "main",
      .module = frag_shader,
    },
  };

  VkPushConstantRange push = {
    .size = sizeof(HexPushConstant),
    .offset = 0,
    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
  };

  VkPipelineLayoutCreateInfo layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pPushConstantRanges = &push,
    .pushConstantRangeCount = 1,
  };

  VK_RESULT(vkCreatePipelineLayout(gpu->device, &layout_info, NULL, &pipeline->layout));

  VkPipelineVertexInputStateCreateInfo vertex_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
  };

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

  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_info = {
    .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,
  };

  VkPipelineColorBlendAttachmentState blend_attachments = {
    .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_ZERO,
    .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
    .alphaBlendOp = VK_BLEND_OP_ADD,
  };

  VkPipelineColorBlendStateCreateInfo blend_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
    .logicOpEnable = VK_FALSE,
    .logicOp = VK_LOGIC_OP_COPY,
    .attachmentCount = 1,
    .pAttachments = &blend_attachments,
  };

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

  VkPipelineDynamicStateCreateInfo dynamic_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    .dynamicStateCount = sizeof(dynamic_states)/sizeof(VkDynamicState),
    .pDynamicStates = dynamic_states,
  };

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

  VkGraphicsPipelineCreateInfo graphics_pipeline_info = {
    .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
    .layout = pipeline->layout,
    .stageCount = sizeof(stages)/sizeof(VkPipelineShaderStageCreateInfo),
    .pStages = stages,
    .pVertexInputState = &vertex_info,
    .pInputAssemblyState = &input_info,
    .pViewportState = &viewport_info,
    .pRasterizationState = &raster_info,
    .pColorBlendState = &blend_info,
    .pDynamicState = &dynamic_info,
    .pMultisampleState = &multisample_info,
    .pDepthStencilState = &depth_info,
    .renderPass = gpu->render_pass,
    .subpass = 0,
    .basePipelineHandle = VK_NULL_HANDLE,
    .basePipelineIndex = -1,
  };

  VK_RESULT(vkCreateGraphicsPipelines(gpu->device, VK_NULL_HANDLE, 1, &graphics_pipeline_info, NULL, &pipeline->pipeline));
  return VK_SUCCESS;
}

VkResult create_hex_context(
  RenderContext* gpu,
  HexContext* context) {
  VkResult result;
   
  VK_RESULT(create_hex_pipeline(gpu, &context->graphics));
  VK_RESULT(create_ray_pipeline(gpu, &context->ray_pipeline));
  VK_RESULT(create_hex_highlight_pipeline(gpu, &context->highlight_pipeline));

  memset(&context->data, 0, sizeof(GPUHexContext));
  glm_perspective(PERSPECTIVE_FOVY, (float)gpu->swapchain_extent.width/(float)gpu->swapchain_extent.height, PERSPECTIVE_NEARZ, PERSPECTIVE_FARZ, context->data.proj);
  glm_mat4_identity(context->data.view);
  VK_RESULT(create_storage_buffer(
        gpu->allocator, 
        0,
        sizeof(GPUHexContext),
        &context->context,
        &context->context_memory));
  context->address = buffer_address(gpu->device, context->context);
  VK_RESULT(add_transfer(&context->data, context->context, 0, sizeof(GPUHexContext), gpu->current_frame, gpu));

  return VK_SUCCESS;
}

VkResult set_hex_region(int32_t q, int32_t r, int32_t y, uint32_t map, HexRegion** region, HexContext* hex, RenderContext* gpu) {
  VkResult result;

  uint32_t i = 0;
  for(; i < MAX_LOADED_REGIONS; i++) {
    if(hex->regions[i] == NULL) {
      hex->regions[i] = malloc(sizeof(GPUHexRegion));
      *region = hex->regions[i];
      break;
    }
  }

  if(*region == NULL) {
    return VK_ERROR_OUT_OF_HOST_MEMORY;
  }

  (*region)->data.q = q;
  (*region)->data.r = r;
  (*region)->data.y = y;
  (*region)->data.map = map;
  VK_RESULT(create_storage_buffer(
        gpu->allocator,
        0,
        sizeof(GPUHexRegion),
        &(*region)->region,
        &(*region)->region_memory));
  VK_RESULT(add_transfer(&(*region)->data.q, (*region)->region, offsetof(GPUHexRegion, q), sizeof(int32_t)*3 + sizeof(uint32_t), gpu->current_frame, gpu));

  (*region)->address = buffer_address(gpu->device, (*region)->region);

  VK_RESULT(add_transfer(&(*region)->address, hex->context, offsetof(GPUHexContext, regions) + sizeof(VkDeviceAddress)*i, sizeof(VkDeviceAddress), gpu->current_frame, gpu)); 

  return VK_SUCCESS;
}