#include "vk_pipelines.h"
#include <fstream>
#include "vk_initializers.h"
#include <print>

bool vkutil::load_shader_module(const char* filePath,
    VkDevice device,
    VkShaderModule* outShaderModule)
{
  std::ifstream file(filePath, std::ios::ate | std::ios::binary);

  if(!file.is_open()) {
    std::println("failed to open file: {}", filePath);
    return false;
  }

  size_t fileSize = (size_t)file.tellg();

  //spirv expects the buffer to uint32
  std::vector<uint32_t> buffer(fileSize / sizeof(uint32_t));
  file.seekg(0);
  file.read((char*)buffer.data(), fileSize);

  file.close();

  VkShaderModuleCreateInfo createInfo = {
    .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
    .pNext = nullptr,
    // codeSize has to be in bytes
    .codeSize = buffer.size() * sizeof(uint32_t),
    .pCode = buffer.data(),
  };

  VkShaderModule shaderModule;
  if(vkCreateShaderModule(device, &createInfo, nullptr, &shaderModule) != VK_SUCCESS) {
    std::println("Error after vkCreateShaderModule");
    return false;
  }

  *outShaderModule = shaderModule;
  return true;
}


void PipelineBuilder::clear() {
   // clear all of the structs we need back to 0 with their correct stype

  _inputAssembly = { .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO };

  _rasterizer = { .sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO };

  _colorBlendAttachment = {};

  _multisampling = { .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO };

  _pipelineLayout = {};

  _depthStencil = { .sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO };

  _renderInfo = { .sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO };

  _shaderStages.clear();
}

VkPipeline PipelineBuilder::build_pipeline(VkDevice device)
{
  VkPipelineViewportStateCreateInfo viewportState{
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
    .viewportCount = 1,
    .scissorCount = 1,
  };

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

  VkPipelineVertexInputStateCreateInfo _vertexInputInfo{
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO
  };

  VkDynamicState state[] = {
    VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR
  };

  VkPipelineDynamicStateCreateInfo dynamicInfo{
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
      .dynamicStateCount = 2,
      // ZED: we don't need to do this? state is also &state[0]
      .pDynamicStates = &state[0],
  };

  // build the actual pipeline
  VkGraphicsPipelineCreateInfo pipelineInfo = {
    .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
    .pNext = &_renderInfo,
    .stageCount = (uint32_t)_shaderStages.size(),
    .pStages = _shaderStages.data(),
    .pVertexInputState = &_vertexInputInfo,
    .pInputAssemblyState = &_inputAssembly,
    .pViewportState = &viewportState,
    .pRasterizationState = &_rasterizer,
    .pMultisampleState = &_multisampling,
    .pDepthStencilState = &_depthStencil,
    .pColorBlendState = &colorBlending,
    .pDynamicState = &dynamicInfo,
    .layout = _pipelineLayout,
  };

  // its easy to error out on create graphics pipeline, so we handle it a bit
  // better than the common VK_CHECK case
  
  VkPipeline newPipeline;
  auto works = vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &newPipeline);

  if(works != VK_SUCCESS) {
    std::println("failed to create pipeline");
    return VK_NULL_HANDLE;
  } else {
    return newPipeline;
  }
}

void PipelineBuilder::set_shaders(VkShaderModule vertexShader, VkShaderModule fragmentShader)
{
  _shaderStages.clear();

  _shaderStages.push_back(
      vkinit::pipeline_shader_stage_create_info(
        VK_SHADER_STAGE_VERTEX_BIT,
        vertexShader));

  _shaderStages.push_back(
      vkinit::pipeline_shader_stage_create_info(
        VK_SHADER_STAGE_FRAGMENT_BIT,
        fragmentShader));
}

void PipelineBuilder::set_input_topology(VkPrimitiveTopology topology)
{
  _inputAssembly.topology = topology;
  _inputAssembly.primitiveRestartEnable = VK_FALSE;
}

void PipelineBuilder::set_polygon_mode(VkPolygonMode mode)
{
  _rasterizer.polygonMode = mode;
  _rasterizer.lineWidth = 1.0f;
}

void PipelineBuilder::set_cull_mode(VkCullModeFlags cullMode, VkFrontFace frontFace)
{
  _rasterizer.cullMode = cullMode;
  _rasterizer.frontFace = frontFace;
}

void PipelineBuilder::set_multisampling_none()
{
  _multisampling.sampleShadingEnable = VK_FALSE;
  _multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
  _multisampling.minSampleShading = 1.0f;
  _multisampling.pSampleMask = nullptr;
  _multisampling.alphaToCoverageEnable = VK_FALSE;
  _multisampling.alphaToOneEnable = VK_FALSE;
}

void PipelineBuilder::disable_blending() 
{
  _colorBlendAttachment.colorWriteMask =
    VK_COLOR_COMPONENT_R_BIT
    | VK_COLOR_COMPONENT_G_BIT
    | VK_COLOR_COMPONENT_B_BIT
    | VK_COLOR_COMPONENT_A_BIT;

  _colorBlendAttachment.blendEnable = VK_FALSE;
}

void PipelineBuilder::set_color_attachment_format(VkFormat format)
{
  _colorAttachmentFormat = format;
  _renderInfo.colorAttachmentCount = 1;
  _renderInfo.pColorAttachmentFormats = &_colorAttachmentFormat;
}

void PipelineBuilder::set_depth_format(VkFormat format)
{
  _renderInfo.depthAttachmentFormat = format;
}

void PipelineBuilder::disable_depthtest()
{
  _depthStencil.depthTestEnable = VK_FALSE;
  _depthStencil.depthWriteEnable = VK_FALSE;
  _depthStencil.depthCompareOp = VK_COMPARE_OP_NEVER;
  _depthStencil.depthBoundsTestEnable = VK_FALSE;
  _depthStencil.stencilTestEnable = VK_FALSE;
  _depthStencil.front = {};
  _depthStencil.back = {};
  _depthStencil.minDepthBounds = 0.0f;
  _depthStencil.maxDepthBounds = 1.0f;
}