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

VkResult create_hex_context(
  RenderContext* gpu,
  HexContext* context) {
  VkResult result;
  
  // Compute Pipeline

  VkShaderModule compute_shader = load_shader_file("shader/hex.comp.spv", gpu->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 compute_push = {
    .size = sizeof(HexPushConstant),
    .offset = 0,
    .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
  };

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

  VK_RESULT(vkCreatePipelineLayout(gpu->device, &compute_layout_info, NULL, &context->compute.layout));

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

  VK_RESULT(vkCreateComputePipelines(gpu->device, VK_NULL_HANDLE, 1, &compute_pipeline_info, NULL, &context->compute.pipeline));

  // Graphics Pipeline
  
  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 graphics_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 graphics_push = {
    .size = sizeof(HexPushConstant),
    .offset = 0,
    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
  };

  VkPipelineLayoutCreateInfo graphics_layout_info = {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .pPushConstantRanges = &graphics_push,
    .pushConstantRangeCount = 1,
  };

  VK_RESULT(vkCreatePipelineLayout(gpu->device, &graphics_layout_info, NULL, &context->graphics.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 = context->graphics.layout,
    .stageCount = sizeof(graphics_stages)/sizeof(VkPipelineShaderStageCreateInfo),
    .pStages = graphics_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, &context->graphics.pipeline));

  glm_perspective(-M_PI*1/8, (float)gpu->swapchain_extent.width/(float)gpu->swapchain_extent.height, 0.01, 1000, context->data.proj);
  glm_mat4_identity(context->data.view);
  VK_RESULT(create_storage_buffer(
        gpu->allocator, 
        0,
        sizeof(GPUHexContext),
        &context->context,
        &context->context_memory));
  fprintf(stderr, "Created hex context %p\n", context->context);
  context->address = buffer_address(gpu->device, context->context);
  VK_RESULT(add_transfer(&context->data, context->context, 0, sizeof(GPUHexContext) - (sizeof(GPUHexContext) - offsetof(GPUHexContext, regions)), 0, gpu));

  return VK_SUCCESS;
}

VkResult create_hex_region(int32_t q, int32_t r, HexRegion** region, HexContext* hex, RenderContext* gpu) {
  VkResult result;

  uint32_t i = 0;
  for(; i < MAX_LOADED_REGIONS; i++) {
    if(hex->regions[i].address == 0) {
      *region = &hex->regions[i];
      break;
    }
  }

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

  (*region)->q = q;
  (*region)->r = r;
  VK_RESULT(create_storage_buffer(
        gpu->allocator,
        0,
        sizeof(GPUHexRegion),
        &(*region)->region,
        &(*region)->region_memory));
  VK_RESULT(add_transfer(&(*region)->q, (*region)->region, offsetof(GPUHexRegion, q), sizeof(uint32_t)*2, 0, 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), 0, gpu)); 

  return VK_SUCCESS;
}