#ifndef PIPELINES_H
#define PIPELINES_H
// Normal header

#include "Common/Types.h"
#include <cstddef>
#include <cstdint>
#include <fmt/base.h>
#include <vector>
#include <vulkan/vulkan_core.h>

namespace vkutil {
  bool loadShaderModule(
    const char* filePath,
    VkDevice device,
    VkShaderModule* outShaderModule
  );
};

class PipelineBuilder {
  public:
    std::vector<VkPipelineShaderStageCreateInfo>    _shaderStages;

    VkPipelineInputAssemblyStateCreateInfo          _inputAssembly;
    VkPipelineRasterizationStateCreateInfo          _rasterizer;
    VkPipelineColorBlendAttachmentState             _colorBlendAttachment;
    VkPipelineMultisampleStateCreateInfo            _multisampling;
    VkPipelineLayout                                _pipelineLayout;
    VkPipelineDepthStencilStateCreateInfo           _depthStencil;
    VkPipelineRenderingCreateInfo                   _renderInfo;

    PipelineBuilder() { clear(); }

    void clear();

    VkPipeline buildPipeline(
      VkDevice device
    );

    void setShaders(
      VkShaderModule vertexShader,
      VkShaderModule fragmentShader
    );

    void setInputTopology(
      VkPrimitiveTopology topology
    );
};

#endif

#ifdef PIPELINES_IMPL
#ifndef PIPELINES_IMPL_H
#define PIPELINES_IMPL_H
// Implementation
// Imports
#include <fstream>

// Custom imports
#include "initializers.h"

bool vkutil::loadShaderModule(
  const char* filePath,
  VkDevice device,
  VkShaderModule* outShaderModule
) {
  #ifdef DEBUG
  fmt::println("[Engine:Utils] Loading shader located at {}...", filePath);
  #endif
  // Open the file, with the cursor at the end
  std::ifstream file(
    filePath,
    std::ios::ate | std::ios::binary
  );

  if (!file.is_open()) {
    return false;
  }

  // Find the size of the file by looking
  // at the cursor
  size_t fileSize = (size_t)file.tellg();

  // Allocate a big enough vector to suit SPIR-V
  std::vector<uint32_t> buffer(fileSize/ sizeof(uint32_t));

  file.seekg(0); // Put the cursor at the begining

  // Load the file in the buffer
  file.read((char*)buffer.data(), fileSize);

  file.close();

  // Create the shader module
  VkShaderModuleCreateInfo createInfo {
    .sType    = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
    .pNext    = nullptr,
    .codeSize = buffer.size() * sizeof(uint32_t), // Convert to bytes
    .pCode    = buffer.data()
  };

  // Check that everything went well
  VkShaderModule shaderModule;

  if (vkCreateShaderModule(
    device,
    &createInfo,
    nullptr,
    &shaderModule
  ) != VK_SUCCESS){
    return false;
  }
  *outShaderModule = shaderModule;

  #ifdef DEBUG
  fmt::println("[Engine:Utils] Shader loaded from {} and size {}", filePath, fileSize);
  #endif

  return true;
}

void PipelineBuilder::clear() {
  #ifdef DEBUG
  fmt::println("[PipelineBuilder] Clearing pipeline...");
  #endif
  // Clear structs
  _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();

  #ifdef DEBUG
  fmt::println("[PipelineBuilder] Pipeline cleared !");
  #endif
}

VkPipeline PipelineBuilder::buildPipeline (
  VkDevice device
) {
  #ifdef DEBUG
  fmt::println("[PipelineBuilder] Building rasterisation pipeline...");
  #endif
  // Make the viewport state from our stored viewport and scissor
  // This makes it that it won't support multiple viewports or scissors
  VkPipelineViewportStateCreateInfo viewportState {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
    .pNext = nullptr,

    .viewportCount = 1,
    .scissorCount  = 1
  };

  // Setup placeholder color blending.
  // Need to be changed when using transparent objects, which we aren't rn
  // The blending is "no blend", but we write it to the color attachment
  VkPipelineColorBlendStateCreateInfo colorBlending {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
    .pNext = nullptr,

    .logicOpEnable = VK_FALSE,
    .logicOp = VK_LOGIC_OP_COPY,
    .attachmentCount = 1,
    .pAttachments = &_colorBlendAttachment
  };
  // Completely clear VertexInputStateCreateInfo, as we have no need for it
  VkPipelineVertexInputStateCreateInfo _vertexInputInfo {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO
  };

  // Dynamic Viewport
  VkDynamicState state[] = {
    VK_DYNAMIC_STATE_VIEWPORT,
    VK_DYNAMIC_STATE_SCISSOR
  };

  VkPipelineDynamicStateCreateInfo dynamicInfo {
    .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    .dynamicStateCount = 2,
     .pDynamicStates = &state[0]
  };

  // Build the pipeline
  // Uses all the class members info structs
  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
  };

  // Check the creation of the pipeline
  // I don't use the VK_CHECK()
  // 'cause there might be pipeline-specific options
  VkPipeline newPipeline;

  if(
    vkCreateGraphicsPipelines(
      device,
      VK_NULL_HANDLE,
      1,
      &pipelineInfo,
      nullptr,
      &newPipeline
    ) != VK_SUCCESS
  ) {
    fmt::println("[PipelineBuilder] Failed to create pipeline !");
    return VK_NULL_HANDLE; // Graphics pipeline failed to be created
  } else {
    return newPipeline;
  }

}

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

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

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

#endif
#endif