void
Rhi_SetFrameConstant(rhi_cmdlist cmd, byteview data)
{
    ASSERT(data.size <= rhi->limits.FrameConstantSize);

    uint8_t* mem = Rhi_MapBuffer(rhi->constantsUniformBuffer);
    memcpy(mem + ((rhi_cmdlist_data*)cmd)->frameConstantHead, data.data, data.size);
}

void
Rhi_SetPassConstant(rhi_cmdlist cmd, byteview data)
{
    ASSERT(data.size <= rhi->limits.PassConstantSize);

    uint8_t* mem = Rhi_MapBuffer(rhi->constantsUniformBuffer);
    memcpy(mem + ((rhi_cmdlist_data*)cmd)->passConstantHead, data.data, data.size);
}

void
Rhi_SetDrawConstant(rhi_cmdlist cmd, byteview data)
{
    ASSERT(data.size <= rhi->limits.DrawConstantSize);

    uint8_t* mem = Rhi_MapBuffer(rhi->constantsUniformBuffer);
    memcpy(mem + ((rhi_cmdlist_data*)cmd)->drawConstantHead, data.data, data.size);
}

void
Rhi_CmdBindGraphicsPipeline(rhi_cmdlist cmd, rhi_graphics_pipeline pipeline)
{
    VkCommandBuffer cmdbuf = ((rhi_cmdlist_data*)cmd)->cmdbuf;

    rhi_graphics_pipeline_data* Pipeline = (rhi_graphics_pipeline_data*)pipeline;

    vkCmdBindPipeline(cmdbuf, VK_PIPELINE_BIND_POINT_GRAPHICS, Pipeline->pipeline);
    ((rhi_cmdlist_data*)cmd)->boundGraphicsPipeline = pipeline;
    
    VkDescriptorSet descriptorSets[2];
    descriptorSets[0] = rhi->texturesDescriptorTable.set;
    descriptorSets[1] = rhi->SamplersDescriptorTable.set;

    vkCmdBindDescriptorSets(cmdbuf,
                            VK_PIPELINE_BIND_POINT_GRAPHICS,
                            Pipeline->layout,
                            1,
                            COUNT(descriptorSets),
                            descriptorSets,
                            0,
                            nullptr);
}

void
Rhi_CmdPushGraphicsConstants(rhi_cmdlist cmd, uint32_t offset, rhi_shader_stage stage, byteview data)
{
    rhi_cmdlist_data* Cmd = ((rhi_cmdlist_data*)cmd);
    rhi_graphics_pipeline_data* Pipeline = (rhi_graphics_pipeline_data*)Cmd->boundGraphicsPipeline;

    vkCmdPushConstants(Cmd->cmdbuf,
                       Pipeline->layout,
                       ConvertShaderStageIntoVkShaderStageFlags(stage),
                       offset,
                       data.size,
                       data.data);
}

void
Rhi_CmdBindVertexBuffers(alloc* Scratch,
                         rhi_cmdlist cmd,
                         uint32_t firstBinding,
                         rhi_buffer* buffers,
                         uint32_t bufferCount,
                         uint64_t* offsets)
{
    uint8_t* ScratchMark = Scratch->at;
    
    VkBuffer* vkBufs = PUSH_ALLOC(Scratch, VkBuffer, bufferCount);
    VkDeviceSize* vkOffsets = PUSH_ALLOC(Scratch, VkDeviceSize, bufferCount);

    for (uint32_t i = 0; i < bufferCount; ++i) {
        vkBufs[i] = buffers[i]->buffer;
        vkOffsets[i] = offsets[i];
    }

    vkCmdBindVertexBuffers(cmd->cmdbuf,
                           firstBinding,
                           bufferCount,
                           vkBufs,
                           vkOffsets);

    Scratch->at = ScratchMark;
}

void
Rhi_CmdBindIndexBuffer(rhi_cmdlist cmd, rhi_buffer buffer, uint64_t offset)
{
    vkCmdBindIndexBuffer(((rhi_cmdlist_data*)cmd)->cmdbuf,
                         ((rhi_buffer_data*)buffer)->buffer,
                         offset,
                         VK_INDEX_TYPE_UINT16);
}

static void
CmdDrawInternal(rhi_cmdlist_data* Cmd)
{
    uint32_t offsets[3];
    offsets[0] = Cmd->frameConstantHead;
    offsets[1] = Cmd->passConstantHead;
    offsets[2] = Cmd->drawConstantHead;

    vkCmdBindDescriptorSets(Cmd->cmdbuf,
                            VK_PIPELINE_BIND_POINT_GRAPHICS,
                            ((rhi_graphics_pipeline_data*)Cmd->boundGraphicsPipeline)->layout,
                            0,
                            1,
                            &rhi->constantsDescriptorTable.set,
                            COUNT(offsets),
                            offsets);

    Cmd->drawConstantHead += CbvOffset(rhi->limits.DrawConstantSize);
}

void
Rhi_CmdDraw(rhi_cmdlist cmd,
            uint32_t vertexCount,
            uint32_t instanceCount,
            uint32_t firstVertex,
            uint32_t firstInstance)
{
    rhi_cmdlist_data* Cmd = (rhi_cmdlist_data*)cmd;
    CmdDrawInternal(Cmd);
    vkCmdDraw(Cmd->cmdbuf, vertexCount, instanceCount, firstVertex, firstInstance);
}

void
Rhi_CmdDrawIndexed(rhi_cmdlist cmd,
                   uint32_t indexCount,
                   int32_t vertexOffset,
                   uint32_t instanceCount,
                   uint32_t firstIndex,
                   uint32_t firstInstance)
{
    rhi_cmdlist_data* Cmd = (rhi_cmdlist_data*)cmd;
    CmdDrawInternal(Cmd);
    vkCmdDrawIndexed(Cmd->cmdbuf, indexCount, instanceCount, firstIndex, vertexOffset, firstInstance);
}

rhi_graphics_pipeline
Rhi_CreateGraphicsPipeline(alloc* Scratch,
                           const char* Name,
                           rhi_shader vertexShader,
                           rhi_shader pixelShader,
                           rhi_vertex_binding_desc* VertexBindings,
                           uint64_t VertexBindingsCount,
                           rhi_vertex_attribute_desc* VertexAttributes,
                           uint64_t VertexAttributesCount,
                           rhi_texture_format* ColorTargetFormats,
                           uint64_t ColorTargetFormatsCount,
                           rhi_pipeline_flags Flags,
                           rhi_texture_format DepthFormat)
{
    uint8_t* ScratchMark = Scratch->at;

    spv_shader* VertexShader = PUSH_ALLOC(Scratch, spv_shader, 1);
    VertexShader->Stage = Shader_Vertex;
    ProcessSpvShader(VertexShader,
                     ((rhi_shader_data*)vertexShader)->spv,
                     ((rhi_shader_data*)vertexShader)->spvCount,
                     PUSH_ALLOC(Scratch, uint32_t, ((rhi_shader_data*)vertexShader)->spvCount));

    spv_shader* PixelShader = PUSH_ALLOC(Scratch, spv_shader, 1);
    PixelShader->Stage = Shader_Pixel;
    ProcessSpvShader(PixelShader,
                     ((rhi_shader_data*)pixelShader)->spv,
                     ((rhi_shader_data*)pixelShader)->spvCount,
                     PUSH_ALLOC(Scratch, uint32_t, ((rhi_shader_data*)pixelShader)->spvCount));

    rhi_graphics_pipeline_data* Pipeline = nullptr;
    for (uint32_t i = 0; i < rhi->limits.GraphicsPipelinesCount; ++i) {
        if (!(rhi->GraphicsPipelines[i].Flags & VulkanFlag_Used)) {
            Pipeline = &rhi->GraphicsPipelines[i];
            Pipeline->Flags |= VulkanFlag_Used;
            break;
        }
    }
    ASSERT(Pipeline);

    Pipeline->bindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;

    VkVertexInputBindingDescription* vkVertexBindings =
        PUSH_ALLOC(Scratch, VkVertexInputBindingDescription, VertexBindingsCount);
    for (uint32_t i = 0; i < VertexBindingsCount; ++i) {
        rhi_vertex_binding_desc* binding = &VertexBindings[i];
        
        vkVertexBindings[i].binding = binding->binding;
        vkVertexBindings[i].stride = binding->stride;
        vkVertexBindings[i].inputRate = ConvertVulkanInputRate(binding->inputRate);            
    }

    VkVertexInputAttributeDescription* vkVertexAttributes =
        PUSH_ALLOC(Scratch, VkVertexInputAttributeDescription, VertexAttributesCount);
    for (uint32_t i = 0; i < VertexAttributesCount; ++i) {
        rhi_vertex_attribute_desc* attribute = &VertexAttributes[i];

        vkVertexAttributes[i].location = FindLocationBySemantic(VertexShader, attribute->semantic, SpvStorageClassInput);
        vkVertexAttributes[i].binding = attribute->binding;
        vkVertexAttributes[i].format = ConvertVertexFormatIntoVkFormat(attribute->format);
        vkVertexAttributes[i].offset = attribute->offset;
    }

    LinkSpvShaders(VertexShader, PixelShader);
    
    VkShaderModule vsModule;
    {
        VkShaderModuleCreateInfo createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
        createInfo.codeSize = VertexShader->WordCount * sizeof(uint32_t);
        createInfo.pCode = VertexShader->Words;

        VK_CHECK(vkCreateShaderModule(rhi->dev, &createInfo, nullptr, &vsModule));
    }
    
    VkShaderModule psModule;
    {
        VkShaderModuleCreateInfo createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
        createInfo.codeSize = PixelShader->WordCount * sizeof(uint32_t);
        createInfo.pCode = PixelShader->Words;

        VK_CHECK(vkCreateShaderModule(rhi->dev, &createInfo, nullptr, &psModule));
    }

    VkPipelineShaderStageCreateInfo* stages = PUSH_ALLOC(Scratch, VkPipelineShaderStageCreateInfo, 2);

    stages[0].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    stages[0].stage = VK_SHADER_STAGE_VERTEX_BIT;
    stages[0].module = vsModule;
    stages[0].pName = "main";
    
    stages[1].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    stages[1].stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    stages[1].module = psModule;
    stages[1].pName = "main";

    VkPipelineVertexInputStateCreateInfo vertexInputStateCreateInfo{};
    vertexInputStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputStateCreateInfo.vertexBindingDescriptionCount = VertexBindingsCount;
    vertexInputStateCreateInfo.pVertexBindingDescriptions = vkVertexBindings;
    vertexInputStateCreateInfo.vertexAttributeDescriptionCount = VertexAttributesCount;
    vertexInputStateCreateInfo.pVertexAttributeDescriptions = vkVertexAttributes;

    VkPipelineRasterizationStateCreateInfo rasterizerCreateInfo{};
    rasterizerCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterizerCreateInfo.depthClampEnable = VK_FALSE;
    rasterizerCreateInfo.rasterizerDiscardEnable = VK_FALSE;
    rasterizerCreateInfo.polygonMode = VK_POLYGON_MODE_FILL;
    rasterizerCreateInfo.cullMode = ConvertVulkanCullMode(Flags);
    rasterizerCreateInfo.frontFace = ConvertVulkanFrontFace(Flags);
    rasterizerCreateInfo.lineWidth = 1.0f;

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyCreateInfo{};
    inputAssemblyCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssemblyCreateInfo.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

    VkPipelineViewportStateCreateInfo viewportStateCreateInfo{};
    viewportStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportStateCreateInfo.viewportCount = 1;
    viewportStateCreateInfo.scissorCount = 1;

    VkPipelineMultisampleStateCreateInfo multisamplingCreateInfo{};
    multisamplingCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisamplingCreateInfo.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
    multisamplingCreateInfo.sampleShadingEnable = VK_FALSE;
    multisamplingCreateInfo.minSampleShading = 1.0f;
    multisamplingCreateInfo.alphaToCoverageEnable = VK_FALSE;
    multisamplingCreateInfo.alphaToOneEnable = VK_FALSE;

    VkPipelineColorBlendAttachmentState colorBlendAttachment{};
    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;
    colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;

    VkPipelineColorBlendStateCreateInfo colorBlendingCreateInfo{};
    colorBlendingCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlendingCreateInfo.logicOpEnable = VK_FALSE;
    colorBlendingCreateInfo.logicOp = VK_LOGIC_OP_COPY;
    colorBlendingCreateInfo.attachmentCount = 1;
    colorBlendingCreateInfo.pAttachments = &colorBlendAttachment;

    VkDynamicState* dynamicStates = PUSH_ALLOC(Scratch, VkDynamicState, 2);
    dynamicStates[0] = VK_DYNAMIC_STATE_VIEWPORT;
    dynamicStates[1] = VK_DYNAMIC_STATE_SCISSOR;

    VkPipelineDynamicStateCreateInfo dynamicStateCreateInfo{};
    dynamicStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicStateCreateInfo.dynamicStateCount = 2;
    dynamicStateCreateInfo.pDynamicStates = dynamicStates;

    VkFormat* vkColorTargetFormats = PUSH_ALLOC(Scratch, VkFormat, ColorTargetFormatsCount);
    for (uint32_t i = 0; i < ColorTargetFormatsCount; ++i) {
        vkColorTargetFormats[i] = ConvertTextureFormatIntoVkFormat(ColorTargetFormats[i], nullptr);
    }

    VkPipelineRenderingCreateInfo pipelineRenderingCreateInfo{};
    pipelineRenderingCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO;
    pipelineRenderingCreateInfo.colorAttachmentCount = ColorTargetFormatsCount;
    pipelineRenderingCreateInfo.pColorAttachmentFormats = vkColorTargetFormats;
    pipelineRenderingCreateInfo.depthAttachmentFormat = ConvertTextureFormatIntoVkFormat(DepthFormat, nullptr);

    VkDescriptorSetLayout* descriptorSetLayouts = PUSH_ALLOC(Scratch, VkDescriptorSetLayout, 3);
    descriptorSetLayouts[0] = rhi->constantsDescriptorTable.layout;
    descriptorSetLayouts[1] = rhi->texturesDescriptorTable.layout;
    descriptorSetLayouts[2] = rhi->SamplersDescriptorTable.layout;

    VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo{};
    pipelineLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutCreateInfo.setLayoutCount = 3;
    pipelineLayoutCreateInfo.pSetLayouts = descriptorSetLayouts;

    vkCreatePipelineLayout(rhi->dev, &pipelineLayoutCreateInfo, nullptr, &Pipeline->layout);

    VkPipelineDepthStencilStateCreateInfo depthStencilState{};
    depthStencilState.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    depthStencilState.depthTestEnable = Flags & Pipeline_DepthTest;
    depthStencilState.depthWriteEnable = Flags & Pipeline_DepthWrite;
    depthStencilState.depthCompareOp = VK_COMPARE_OP_LESS;
    depthStencilState.depthBoundsTestEnable = false;
    depthStencilState.stencilTestEnable = false;
    depthStencilState.front = {};
    depthStencilState.back = {};
    depthStencilState.minDepthBounds = 0.0f;
    depthStencilState.maxDepthBounds = 1.0f;

    VkGraphicsPipelineCreateInfo pipelineCreateInfo{};
    pipelineCreateInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineCreateInfo.pNext = &pipelineRenderingCreateInfo;
    pipelineCreateInfo.stageCount = 2;
    pipelineCreateInfo.pStages = stages;
    pipelineCreateInfo.pVertexInputState = &vertexInputStateCreateInfo;
    pipelineCreateInfo.pInputAssemblyState = &inputAssemblyCreateInfo;
    pipelineCreateInfo.pViewportState = &viewportStateCreateInfo;
    pipelineCreateInfo.pRasterizationState = &rasterizerCreateInfo;
    pipelineCreateInfo.pMultisampleState = &multisamplingCreateInfo;
    pipelineCreateInfo.pDepthStencilState = &depthStencilState;
    pipelineCreateInfo.pColorBlendState = &colorBlendingCreateInfo;
    pipelineCreateInfo.pDynamicState = &dynamicStateCreateInfo;
    pipelineCreateInfo.layout = Pipeline->layout;
    pipelineCreateInfo.subpass = 0;
    pipelineCreateInfo.basePipelineHandle = VK_NULL_HANDLE;
    pipelineCreateInfo.basePipelineIndex = -1;

    VK_CHECK(vkCreateGraphicsPipelines(rhi->dev, VK_NULL_HANDLE, 1, &pipelineCreateInfo, nullptr, &Pipeline->pipeline));

    VulkanNameHandle(VK_OBJECT_TYPE_PIPELINE, (uint64_t)Pipeline->pipeline, Name);

    vkDestroyShaderModule(rhi->dev, vsModule, nullptr);
    vkDestroyShaderModule(rhi->dev, psModule, nullptr);

    Scratch->at = ScratchMark;

    return Pipeline;
}