//
//  KRPipeline.cpp
//  Kraken Engine
//
//  Copyright 2026 Kearwood Gilbert. All rights reserved.
//  
//  Redistribution and use in source and binary forms, with or without modification, are
//  permitted provided that the following conditions are met:
//  
//  1. Redistributions of source code must retain the above copyright notice, this list of
//  conditions and the following disclaimer.
//  
//  2. Redistributions in binary form must reproduce the above copyright notice, this list
//  of conditions and the following disclaimer in the documentation and/or other materials
//  provided with the distribution.
//  
//  THIS SOFTWARE IS PROVIDED BY KEARWOOD GILBERT ''AS IS'' AND ANY EXPRESS OR IMPLIED
//  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
//  FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL KEARWOOD GILBERT OR
//  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
//  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
//  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
//  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
//  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
//  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//  
//  The views and conclusions contained in the software and documentation are those of the
//  authors and should not be interpreted as representing official policies, either expressed
//  or implied, of Kearwood Gilbert.
//

#include "KRPipeline.h"
#include "assert.h"
#include "nodes/KRLight.h"
#include "nodes/KRDirectionalLight.h"
#include "nodes/KRSpotLight.h"
#include "nodes/KRPointLight.h"
#include "KRContext.h"
#include "KRRenderPass.h"
#include "KRModelView.h"

using namespace hydra;

KRPipeline::KRPipeline(KRContext& context, KrDeviceHandle deviceHandle, const KRRenderPass* renderPass, Vector2i viewport_size, Vector2i scissor_size, const PipelineInfo& info, const char* szKey, const std::vector<KRShader*>& shaders, const VertexBufferLayout* layout)
  : KRContextObject(context)
  , m_deviceHandle(deviceHandle)
{
  std::unique_ptr<KRDevice>& device = getContext().getDeviceManager()->getDevice(m_deviceHandle);

  for (StageInfo& stageInfo : m_stages) {
    PushConstantInfo& pushConstants = stageInfo.pushConstants;
    pushConstants.buffer = nullptr;
    pushConstants.bufferSize = 0;
    memset(pushConstants.size, 0, kPushConstantCount);
    memset(pushConstants.offset, 0, kPushConstantCount * sizeof(int));
    pushConstants.layout = nullptr;
  }

  m_descriptorSetLayout = nullptr;
  m_pipelineLayout = nullptr;
  m_graphicsPipeline = nullptr;
  m_descriptorSets.reserve(KRENGINE_MAX_FRAMES_IN_FLIGHT);

  // TODO - Handle device removal

  strcpy(m_szKey, szKey);

  const int kMaxStages = 4;
  VkPipelineShaderStageCreateInfo stages[kMaxStages];
  memset(static_cast<void*>(stages), 0, sizeof(VkPipelineShaderStageCreateInfo) * kMaxStages);
  size_t stage_count = 0;

  std::vector<VkDescriptorSetLayoutBinding> uboLayoutBindings;

  int attribute_locations[kMaxAttributes] = {};

  uint32_t layout_binding_count = 0;
  for (KRShader* shader : shaders) {
    const SpvReflectShaderModule* reflection = shader->getReflection();
    layout_binding_count += reflection->descriptor_binding_count;
  }
  uboLayoutBindings.reserve(layout_binding_count);

  for (KRShader* shader : shaders) {
    VkShaderModule shaderModule;
    if (!shader->createShaderModule(device->m_logicalDevice, shaderModule)) {
      // failed! TODO - Error handling
    }
    const SpvReflectShaderModule* reflection = shader->getReflection();
    
    for (uint32_t b = 0; b < reflection->descriptor_binding_count; b++) {
      SpvReflectDescriptorBinding& binding_reflect = reflection->descriptor_bindings[b];
      VkDescriptorSetLayoutBinding& binding = uboLayoutBindings.emplace_back();
      memset(&binding, 0, sizeof(VkDescriptorSetLayoutBinding));
      binding.binding = binding_reflect.binding;
      // Note: VkDescriptorType and SpvReflectDescriptorType values match
      binding.descriptorType = static_cast<VkDescriptorType>(binding_reflect.descriptor_type);
      binding.descriptorCount = binding_reflect.count;
      binding.pImmutableSamplers = nullptr;
      binding.stageFlags = shader->getShaderStageFlagBits();
    }

    VkPipelineShaderStageCreateInfo& stageInfo = stages[stage_count++];
    stageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    stageInfo.stage = shader->getShaderStageFlagBits();
    if (stageInfo.stage == VK_SHADER_STAGE_VERTEX_BIT) {

      for (int i = 0; i < reflection->input_variable_count; i++) {
        attribute_locations[i] = -1;
        SpvReflectInterfaceVariable& input_var = *reflection->input_variables[i];
        for (int i = 0; i < kMaxAttributes; i++) {
          const VertexAttributeInfo& attrib = info.layout->attributes[i];
          if (attrib.component == ComponentType::empty) {
            break;
          }
          
          const char* attribBaseName = VertexAttributeNames[(int)attrib.attribute];
          int len = strlen(input_var.name);
          int baseNameLen = strlen(attribBaseName);
          if (len == baseNameLen) {
            if (strcmp(attribBaseName, input_var.name) == 0) {
              attribute_locations[i] = input_var.location;
            }
            break;
          }
          if (len + 1 != baseNameLen) {
            continue;
          }
          char lastChar = input_var.name[len - 1];
          if (lastChar < '0' || lastChar > '9') {
            continue;
          }

          int set = lastChar - '0';
          for (int j = 0; j < i; j++) {
            if (info.layout->attributes[j].attribute == attrib.attribute) {
              --set;
            }
          }
          if (set != 0) {
            continue;
          }
          attribute_locations[i] = input_var.location;
          break;
        }
      }
    }

    initPushConstantStage(shader->getShaderStage(), reflection);
    initDescriptorSetStage(shader->getShaderStage(), reflection);
    stageInfo.module = shaderModule;
    stageInfo.pName = "main";
  }

  VkVertexInputBindingDescription bindingDescription{};
  bindingDescription.binding = 0;
  bindingDescription.stride = (uint32_t)layout->vertexSize;
  bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;

  uint32_t vertexAttributeCount = 0;
  VkVertexInputAttributeDescription vertexAttributeDescriptions[kMaxAttributes]{};

  for (int i = 0; i < kMaxAttributes; i++) {
    VertexAttributeInfo attrib = layout->attributes[i];
    if (attrib.component == ComponentType::empty) {
      break;
    }
    if (attribute_locations[i] == -1) {
      continue;
    }
    VkVertexInputAttributeDescription& desc = vertexAttributeDescriptions[vertexAttributeCount++];
    desc.binding = 0;
    desc.location = attribute_locations[i];
    desc.format = KRMesh::AttributeVulkanFormat(attrib);
    desc.offset = (uint32_t)info.layout->offsets[i];
  }

  VkPipelineVertexInputStateCreateInfo vertexInputInfo{};
  vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
  vertexInputInfo.vertexBindingDescriptionCount = 1;
  vertexInputInfo.pVertexBindingDescriptions = &bindingDescription;
  vertexInputInfo.vertexAttributeDescriptionCount = vertexAttributeCount;
  vertexInputInfo.pVertexAttributeDescriptions = vertexAttributeDescriptions;

  VkPipelineInputAssemblyStateCreateInfo inputAssembly{};
  inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
  switch (info.layout->topology) {
  case Topology::Points:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_POINT_LIST;
    break;
  case Topology::LineStrips:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_LINE_STRIP;
    break;
  case Topology::Lines:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_LINE_LIST;
    break;
  case Topology::Triangles:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    break;
  case Topology::TriangleStrips:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP;
    break;
  case Topology::TriangleFans:
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_FAN;
    break;
  }

  inputAssembly.primitiveRestartEnable = VK_FALSE;

  VkViewport viewport{};
  viewport.x = 0.0f;
  viewport.y = 0.0f;
  viewport.width = static_cast<float>(viewport_size.x);
  viewport.height = static_cast<float>(viewport_size.y);
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;

  VkRect2D scissor{};
  scissor.offset = { 0, 0 };
  scissor.extent.width = scissor_size.x;
  scissor.extent.height = scissor_size.y;

  VkPipelineViewportStateCreateInfo viewportState{};
  viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
  viewportState.viewportCount = 1;
  viewportState.pViewports = &viewport;
  viewportState.scissorCount = 1;
  viewportState.pScissors = &scissor;

  VkPipelineRasterizationStateCreateInfo rasterizer{};
  rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
  rasterizer.depthClampEnable = VK_FALSE;
  rasterizer.rasterizerDiscardEnable = VK_FALSE;
  rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
  rasterizer.lineWidth = 1.0f;
  switch (info.cullMode) {
  case CullMode::kCullBack:
    rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
    break;
  case CullMode::kCullFront:
    rasterizer.cullMode = VK_CULL_MODE_FRONT_BIT;
    break;
  case CullMode::kCullNone:
    rasterizer.cullMode = VK_CULL_MODE_NONE;
    break;
  }
  rasterizer.frontFace = VK_FRONT_FACE_CLOCKWISE;
  rasterizer.depthBiasEnable = VK_FALSE;
  rasterizer.depthBiasConstantFactor = 0.0f;
  rasterizer.depthBiasClamp = 0.0f;
  rasterizer.depthBiasSlopeFactor = 0.0f;

  VkPipelineMultisampleStateCreateInfo multisampling{};
  multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
  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;

  VkPipelineColorBlendAttachmentState colorBlendAttachment{};
  colorBlendAttachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;

  switch (info.rasterMode) {
  case RasterMode::kOpaque:
  case RasterMode::kOpaqueLessTest:
  case RasterMode::kOpaqueNoTest:
  case RasterMode::kOpaqueNoDepthWrite:
    colorBlendAttachment.blendEnable = VK_FALSE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    break;
  case RasterMode::kAlphaBlend:
  case RasterMode::kAlphaBlendNoTest:
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    break;
  case RasterMode::kAdditive:
  case RasterMode::kAdditiveNoTest:
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    break;
  }
  colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
  colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
  colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

  VkPipelineColorBlendStateCreateInfo colorBlending{};
  colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
  colorBlending.logicOpEnable = VK_FALSE;
  colorBlending.logicOp = VK_LOGIC_OP_COPY;
  colorBlending.attachmentCount = 1;
  colorBlending.pAttachments = &colorBlendAttachment;
  colorBlending.blendConstants[0] = 0.0f;
  colorBlending.blendConstants[1] = 0.0f;
  colorBlending.blendConstants[2] = 0.0f;
  colorBlending.blendConstants[3] = 0.0f;

  if (uboLayoutBindings.size()) {
    VkDescriptorSetLayoutCreateInfo layoutInfo{};
    layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount = uboLayoutBindings.size();
    layoutInfo.pBindings = uboLayoutBindings.data();

    if (vkCreateDescriptorSetLayout(device->m_logicalDevice, &layoutInfo, nullptr, &m_descriptorSetLayout) != VK_SUCCESS) {
      // failed! TODO - Error handling
    }
  }

  VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
  pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
  if (uboLayoutBindings.size()) {
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts = &m_descriptorSetLayout;
  } else {
    pipelineLayoutInfo.setLayoutCount = 0;
    pipelineLayoutInfo.pSetLayouts = nullptr;
  }
  pipelineLayoutInfo.pushConstantRangeCount = 0;
  pipelineLayoutInfo.pPushConstantRanges = nullptr;

  if (vkCreatePipelineLayout(device->m_logicalDevice, &pipelineLayoutInfo, nullptr, &m_pipelineLayout) != VK_SUCCESS) {
    // failed! TODO - Error handling
  }

  int iStage = 0;
  for (StageInfo& stageInfo : m_stages) {
    PushConstantInfo& pushConstants = stageInfo.pushConstants;
    if (pushConstants.buffer) {
      VkPipelineLayoutCreateInfo pushConstantsLayoutInfo{};
      pushConstantsLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
      pushConstantsLayoutInfo.setLayoutCount = 0;
      pushConstantsLayoutInfo.pSetLayouts = nullptr;
      pushConstantsLayoutInfo.pushConstantRangeCount = 0;
      pushConstantsLayoutInfo.pPushConstantRanges = nullptr;

      // TODO - We need to support push constants for other shader stages
      VkPushConstantRange push_constant{};
      push_constant.offset = 0;
      push_constant.size = pushConstants.bufferSize;
      push_constant.stageFlags = getShaderStageFlagBitsFromShaderStage(static_cast<ShaderStage>(iStage));
      pushConstantsLayoutInfo.pPushConstantRanges = &push_constant;
      pushConstantsLayoutInfo.pushConstantRangeCount = 1;

      if (vkCreatePipelineLayout(device->m_logicalDevice, &pushConstantsLayoutInfo, nullptr, &pushConstants.layout) != VK_SUCCESS) {
        // failed! TODO - Error handling
      }
    }
    iStage++;
  }

  VkPipelineDepthStencilStateCreateInfo depthStencil{};
  depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
  switch (info.rasterMode) {
  case RasterMode::kOpaque:
  case RasterMode::kOpaqueLessTest:
    depthStencil.depthTestEnable = VK_TRUE;
    depthStencil.depthWriteEnable = VK_TRUE;
    break;
  case RasterMode::kOpaqueNoTest:
    depthStencil.depthTestEnable = VK_FALSE;
    depthStencil.depthWriteEnable = VK_TRUE;
    break;
  case RasterMode::kOpaqueNoDepthWrite:
  case RasterMode::kAlphaBlend:
  case RasterMode::kAdditive:
    depthStencil.depthTestEnable = VK_TRUE;
    depthStencil.depthWriteEnable = VK_FALSE;
    break;
  case RasterMode::kAlphaBlendNoTest:
  case RasterMode::kAdditiveNoTest:
    depthStencil.depthTestEnable = VK_FALSE;
    depthStencil.depthWriteEnable = VK_FALSE;
    break;
  }

  if (info.rasterMode == RasterMode::kOpaqueLessTest) {
    depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
  } else {
    depthStencil.depthCompareOp = VK_COMPARE_OP_LESS_OR_EQUAL;
  }

  depthStencil.depthBoundsTestEnable = VK_FALSE;
  depthStencil.minDepthBounds = 0.0f;
  depthStencil.maxDepthBounds = 1.0f;
  depthStencil.stencilTestEnable = VK_FALSE;
  depthStencil.front = {};
  depthStencil.back = {};

  VkGraphicsPipelineCreateInfo pipelineInfo{};
  pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
  pipelineInfo.stageCount = stage_count;
  pipelineInfo.pStages = stages;
  pipelineInfo.pVertexInputState = &vertexInputInfo;
  pipelineInfo.pInputAssemblyState = &inputAssembly;
  pipelineInfo.pViewportState = &viewportState;
  pipelineInfo.pRasterizationState = &rasterizer;
  pipelineInfo.pMultisampleState = &multisampling;
  pipelineInfo.pDepthStencilState = &depthStencil;
  pipelineInfo.pColorBlendState = &colorBlending;
  pipelineInfo.pDynamicState = nullptr;
  pipelineInfo.layout = m_pipelineLayout;
  pipelineInfo.renderPass = renderPass->m_renderPass;
  pipelineInfo.subpass = 0;
  pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
  pipelineInfo.basePipelineIndex = -1;

  if (vkCreateGraphicsPipelines(device->m_logicalDevice, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &m_graphicsPipeline) != VK_SUCCESS) {
    // Failed! TODO - Error handling
  }
}

KRPipeline::~KRPipeline()
{
  if (m_graphicsPipeline) {
    // TODO: vkDestroyPipeline(device, m_graphicsPipeline, nullptr);
  }
  if (m_pipelineLayout) {
    // TODO: vkDestroyPipelineLayout(device, m_pipelineLayout, nullptr);
  }
  if (m_descriptorSetLayout) {
    // TODO: vkDestroyDescriptorSetLayout(device, m_descriptorSetLayout, nullptr);
  }

  if (getContext().getPipelineManager()->m_active_pipeline == this) {
    getContext().getPipelineManager()->m_active_pipeline = NULL;
  }
  for (StageInfo& stageInfo : m_stages) {
    PushConstantInfo& pushConstants = stageInfo.pushConstants;
    if (pushConstants.layout) {
      // TODO: vkDestroyPipelineLayout(device, pushConstants.layout, nullptr);
    }
    if (pushConstants.buffer) {
      delete pushConstants.buffer;
      pushConstants.buffer = nullptr;
    }
  }
}

void KRPipeline::initPushConstantStage(ShaderStage stage, const SpvReflectShaderModule* reflection)
{
  PushConstantInfo& pushConstants = m_stages[static_cast<int>(stage)].pushConstants;
  for (int i = 0; i < reflection->push_constant_block_count; i++) {
    const SpvReflectBlockVariable& block = reflection->push_constant_blocks[i];
    if (stricmp(block.name, "PushConstants") == 0) {
      if (block.size > 0) {
        pushConstants.buffer = (__uint8_t*)malloc(block.size);
        memset(pushConstants.buffer, 0, block.size);
        pushConstants.bufferSize = block.size;

        // Get push constant offsets
        for (int iUniform = 0; iUniform < kPushConstantCount; iUniform++) {
          for (int iMember = 0; iMember < block.member_count; iMember++) {
            const SpvReflectBlockVariable& member = block.members[iMember];
            if (IsShaderValueName(iUniform, member.name)) {
              pushConstants.offset[iUniform] = member.offset;
              pushConstants.size[iUniform] = member.size;
              if (member.type_description->op == SpvOpTypeVector && member.numeric.scalar.width == 32) {
                switch (member.numeric.vector.component_count) {
                  case 2:
                    pushConstants.type[iUniform] = ShaderValueType::type_vector2;
                    break;
                  case 3:
                    pushConstants.type[iUniform] = ShaderValueType::type_vector3;
                    break;
                  case 4:
                    pushConstants.type[iUniform] = ShaderValueType::type_vector4;
                    break;
                }
              } else if (member.type_description->op == SpvOpTypeBool) {
                pushConstants.type[iUniform] = ShaderValueType::type_bool;
              } else if (member.type_description->op == SpvOpTypeFloat && member.numeric.scalar.width == 32) {
                pushConstants.type[iUniform] = ShaderValueType::type_float32;
              } else if (member.type_description->op == SpvOpTypeFloat && member.numeric.scalar.width == 64) {
                pushConstants.type[iUniform] = ShaderValueType::type_float64;
              } else if (member.type_description->op == SpvOpTypeInt && member.numeric.scalar.width == 32) {
                pushConstants.type[iUniform] = ShaderValueType::type_int32;
              } else if (member.type_description->op == SpvOpTypeInt && member.numeric.scalar.width == 64) {
                pushConstants.type[iUniform] = ShaderValueType::type_int64;
              } else if (member.type_description->op == SpvOpTypeMatrix && member.numeric.scalar.width == 32) {
                if (member.numeric.matrix.column_count == 2 && member.numeric.matrix.row_count == 2) {
                  pushConstants.type[iUniform] = ShaderValueType::type_matrix2;
                } else if (member.numeric.matrix.column_count == 2 && member.numeric.matrix.row_count == 3) {
                  pushConstants.type[iUniform] = ShaderValueType::type_matrix2x3;
                } else if (member.numeric.matrix.column_count == 4 && member.numeric.matrix.row_count == 4) {
                  pushConstants.type[iUniform] = ShaderValueType::type_matrix4;
                }
              }
              // TODO: Support unsigned integer binding?
            }
          }
        }
      }
    }
  }
}

void KRPipeline::initDescriptorSetStage(ShaderStage stage, const SpvReflectShaderModule* reflection)
{
  std::vector<DescriptorSetInfo>& descriptorSets = m_stages[static_cast<int>(stage)].descriptorSets;
  descriptorSets.reserve(reflection->descriptor_set_count);
  for (int i = 0; i < reflection->descriptor_set_count; i++) {
    SpvReflectDescriptorSet descriptorSet = reflection->descriptor_sets[i];
    DescriptorSetInfo& descriptorSetInfo = descriptorSets.emplace_back();
    descriptorSetInfo.bindings.reserve(descriptorSet.binding_count);
    for (int j = 0; j < descriptorSet.binding_count; j++) {
      SpvReflectDescriptorBinding& binding = *descriptorSet.bindings[j];
      DescriptorBinding& descriptorQuery = descriptorSetInfo.bindings.emplace_back();
      
      switch (binding.descriptor_type) {
      case SPV_REFLECT_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
        {
          ImageDescriptorInfo& imageInfo = descriptorQuery.emplace<ImageDescriptorInfo>();
          imageInfo.name = binding.name;
          imageInfo.texture = nullptr;
          imageInfo.sampler = nullptr;
        }
        break;
      case SPV_REFLECT_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
        {
          UniformBufferDescriptorInfo& bufferInfo = descriptorQuery.emplace<UniformBufferDescriptorInfo>();
          bufferInfo.name = binding.name;
          bufferInfo.buffer = nullptr;
        }
        break;
      default:
        // Not supported
        // TODO - Error handling
        break;
      }
    }
  }
}

bool KRPipeline::hasPushConstant(ShaderValue location) const
{
  for (const StageInfo& stageInfo : m_stages) {
    const PushConstantInfo& pushConstants = stageInfo.pushConstants;
    if (pushConstants.size[static_cast<size_t>(location)]) {
      return true;
    }
  }
  return false;
}


bool KRPipeline::setImageBindings(const std::vector<const KRReflectedObject*> objects)
{
  bool success = true;

  for (int stage = 0; stage < static_cast<size_t>(ShaderStage::ShaderStageCount); stage++) {
    StageInfo& stageInfo = m_stages[stage];
    for (DescriptorSetInfo& descriptorSetInfo : stageInfo.descriptorSets) {
      for (DescriptorBinding& binding : descriptorSetInfo.bindings) {
        ImageDescriptorInfo* image = std::get_if<ImageDescriptorInfo>(&binding);
        if (image) {
          bool found = false;
          const KRTextureBinding* binding = nullptr;
          for (const KRReflectedObject* object : objects) {
            KRSampler* sampler = nullptr;
            if (object->getImageBinding(image->name, &binding, &sampler)) {
              if (binding->isBound() && binding->get()->getStreamLevel() > kraken_stream_level::STREAM_LEVEL_OUT) {
                image->texture = binding->get();
                image->sampler = sampler;
                found = true;
              }
              break;
            }
          }

          if (!found) {
            success = false;
            KRContext::Log(KRContext::LOG_LEVEL_ERROR, "Image binding not found: %s", image->name.c_str());
          }
        }
      }
    }
  }

  return success;
}

bool KRPipeline::setPushConstants(const KRCamera* camera, const std::vector<const KRReflectedObject*> objects)
{
  bool success = true;
  for (StageInfo& stageInfo : m_stages) {
    PushConstantInfo& pushConstants = stageInfo.pushConstants;
    for (int i = 0; i < kPushConstantCount; i++) {
      size_t size = pushConstants.size[static_cast<size_t>(i)];
      if (size != 0) {
        void* constant = (pushConstants.buffer + pushConstants.offset[static_cast<size_t>(i)]);
        bool found = false;
        for (const KRReflectedObject* object : objects) {
          if (object->getShaderValue(camera, static_cast<ShaderValue>(i), pushConstants.type[static_cast<size_t>(i)], constant)) {
            found = true;
            break;
          }
        }
        if (!found) {
          success = false;
          KRContext::Log(KRContext::LOG_LEVEL_ERROR, "Push constant not found: %s", getShaderValueName(i));
        }
      }
    }
  }
  return success;
}

void KRPipeline::updateDescriptorBinding()
{
  // TODO - Implement
  // Vulkan Refactoring
}

bool KRPipeline::updatePushConstants(KRNode::RenderInfo& ri, const Matrix4& matModel)
{
  if (!setPushConstants(ri.camera, ri.reflectedObjects)) {
    return false;
  }
  
  for (StageInfo& stageInfo : m_stages) {
    PushConstantInfo& pushConstants = stageInfo.pushConstants;
    if (pushConstants.buffer) {
      vkCmdPushConstants(ri.commandBuffer, pushConstants.layout, VK_SHADER_STAGE_VERTEX_BIT, 0, pushConstants.bufferSize, pushConstants.buffer);
    }
  }
  
  return true;
}

bool KRPipeline::bind(KRNode::RenderInfo& ri, const Matrix4& matModel)
{
  bool success = true;
  KRDirectionalLight* directionalLight = nullptr;
  if (ri.renderPass->getType() != RenderPassType::RENDER_PASS_DEFERRED_LIGHTS && ri.renderPass->getType() != RenderPassType::RENDER_PASS_DEFERRED_GBUFFER && ri.renderPass->getType() != RenderPassType::RENDER_PASS_DEFERRED_OPAQUE && ri.renderPass->getType() != RenderPassType::RENDER_PASS_SHADOWMAP) {
    if (!ri.directional_lights.empty()) {
      directionalLight = ri.directional_lights.front();
    }
  }

  KRModelView modelView(ri.viewport, matModel, directionalLight);

  ri.reflectedObjects.push_back(&modelView);
  ri.reflectedObjects.push_back(ri.viewport);
  ri.reflectedObjects.push_back(&ri.camera->settings);

  if (success) {
    success = updatePushConstants(ri, matModel);
  }

  if (success) {
    success = setImageBindings(ri.reflectedObjects);
  }

  if (success) {
    updateDescriptorBinding();
    updateDescriptorSets();
    bindDescriptorSets(ri.commandBuffer);

    if (ri.pipeline != this) {
      vkCmdBindPipeline(ri.commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_graphicsPipeline);
      ri.pipeline = this;
    }
  }

  ri.reflectedObjects.pop_back();
  ri.reflectedObjects.pop_back();
  ri.reflectedObjects.pop_back();

  return success;
}

void KRPipeline::updateDescriptorSets()
{
  if (m_descriptorSetLayout == VK_NULL_HANDLE) {
    // There are no descriptors
    return;
  }

  std::unique_ptr<KRDevice>& device = getContext().getDeviceManager()->getDevice(m_deviceHandle);

  // If the descriptor sets are not yet allocted, create them
  if (m_descriptorSets.size() == 0) {
    int descriptorSetCount = 0;
    for (int stage = 0; stage < static_cast<size_t>(ShaderStage::ShaderStageCount); stage++) {
      const StageInfo& stageInfo = m_stages[stage];
      descriptorSetCount += stageInfo.descriptorSets.size();
    }
    m_descriptorSets.resize(KRENGINE_MAX_FRAMES_IN_FLIGHT * descriptorSetCount, VK_NULL_HANDLE);
    std::vector<VkDescriptorSetLayout> layouts(KRENGINE_MAX_FRAMES_IN_FLIGHT, m_descriptorSetLayout);
    device->createDescriptorSets(layouts, m_descriptorSets);
  }

  // Update the descriptor sets
  // TODO - We should only do this when the descriptors have changed

  int descriptorSetCount = m_descriptorSets.size() / KRENGINE_MAX_FRAMES_IN_FLIGHT;
  int descriptorSetStart = (getContext().getCurrentFrame() % KRENGINE_MAX_FRAMES_IN_FLIGHT) * descriptorSetCount;
  int descriptorSetIndex = descriptorSetStart;

  std::vector<VkWriteDescriptorSet> descriptorWrites;
  std::vector<VkDescriptorBufferInfo> buffers;
  std::vector<VkDescriptorImageInfo> images;

  for (int stage = 0; stage < static_cast<size_t>(ShaderStage::ShaderStageCount); stage++) {
    StageInfo& stageInfo = m_stages[stage];
    for (DescriptorSetInfo& descriptorSetInfo : stageInfo.descriptorSets) {
      VkDescriptorSet descriptorSet = m_descriptorSets[descriptorSetIndex++];

      int bindingIndex = 0;
      for (DescriptorBinding& binding : descriptorSetInfo.bindings) {
        UniformBufferDescriptorInfo* buffer = std::get_if<UniformBufferDescriptorInfo>(&binding);
        ImageDescriptorInfo* image = std::get_if<ImageDescriptorInfo>(&binding);
        if (buffer) {
          VkDescriptorBufferInfo& bufferInfo = buffers.emplace_back(VkDescriptorBufferInfo{});
          bufferInfo.buffer = buffer->buffer->getBuffer();
          bufferInfo.offset = 0;
          bufferInfo.range = VK_WHOLE_SIZE;

          VkWriteDescriptorSet& descriptorWrite = descriptorWrites.emplace_back(VkWriteDescriptorSet{});
          descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
          descriptorWrite.dstSet = descriptorSet;
          descriptorWrite.dstBinding = bindingIndex;
          descriptorWrite.dstArrayElement = 0;
          descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
          descriptorWrite.descriptorCount = 1;
          descriptorWrite.pBufferInfo = &bufferInfo;
        } else if (image) {
          VkDescriptorImageInfo& imageInfo = images.emplace_back(VkDescriptorImageInfo{});
          imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
          imageInfo.imageView = image->texture->getFullImageView(m_deviceHandle);
          imageInfo.sampler = image->sampler->getSampler(m_deviceHandle);

          VkWriteDescriptorSet& descriptorWrite = descriptorWrites.emplace_back(VkWriteDescriptorSet{});
          descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
          descriptorWrite.dstSet = descriptorSet;
          descriptorWrite.dstBinding = bindingIndex++;
          descriptorWrite.dstArrayElement = 0;
          descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
          descriptorWrite.descriptorCount = 1;
          descriptorWrite.pImageInfo = &imageInfo;
        } else {
          // TODO - Error Handling
          assert(false);
        }
        bindingIndex++;
      }
    }
  }

  if (!descriptorWrites.empty()) {
    vkUpdateDescriptorSets(device->m_logicalDevice, static_cast<uint32_t>(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);
  }
}

void KRPipeline::bindDescriptorSets(VkCommandBuffer& commandBuffer)
{
  if (m_descriptorSets.empty()) {
    return;
  }
  int descriptorSetCount = (int)m_descriptorSets.size() / KRENGINE_MAX_FRAMES_IN_FLIGHT;
  int startDescriptorSet = (getContext().getCurrentFrame() % KRENGINE_MAX_FRAMES_IN_FLIGHT) * descriptorSetCount;
  VkDescriptorSet descriptorSet = m_descriptorSets[startDescriptorSet];
  if (descriptorSet == VK_NULL_HANDLE) {
    return;
  }
  // TODO - Vulkan Refactoring - Support multiple descriptor set binding
  vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipelineLayout, 0, descriptorSetCount, &descriptorSet, 0, nullptr);
}

const char* KRPipeline::getKey() const
{
  return m_szKey;
}

VkPipeline& KRPipeline::getPipeline()
{
  return m_graphicsPipeline;
}

